File size: 508 Bytes
b399ea6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
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))
|