from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODEL_DIR = "."

tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_DIR,
    trust_remote_code=True,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None
)

question = "Calculate: 8603 + 2303"
print(model.synapse_generate(tokenizer, question, return_trace=True))