from pathlib import Path
import argparse,json,torch
from safetensors.torch import load_file
from transformers import AutoTokenizer
from canopy_r3 import CanopyConfig,CanopyForCausalLM
p=argparse.ArgumentParser();p.add_argument("prompt");args=p.parse_args()
root=Path(__file__).resolve().parent
torch.set_num_threads(2)
device="cuda" if torch.cuda.is_available() else "cpu"
model=CanopyForCausalLM(CanopyConfig(**json.loads((root/"native_config.json").read_text())))
model.load_state_dict(load_file(str(root/"model.safetensors")),strict=True)
model=model.to(device=device,dtype=torch.bfloat16).eval()
tokenizer=AutoTokenizer.from_pretrained(root,local_files_only=True)
ids=tokenizer.encode("\nUser: "+args.prompt+"\nAssistant: ",add_special_tokens=False,return_tensors="pt").to(device)
with torch.inference_mode():
 out=model.generate(ids,max_new_tokens=32,temperature=0.,repetition_penalty=1.,ptrm_stochastic_scale=0.,eos_token_id=0,enable_prefix_sliding=False)
tokens=out[0,ids.shape[1]:].tolist()
end=next((i for i,t in enumerate(tokens) if t in (0,2)),len(tokens))
print(tokenizer.decode(tokens[:end],skip_special_tokens=True).strip())