File size: 872 Bytes
216ce1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | """Minimal local generation example for QuadOrbit-40M."""
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_dir = Path(__file__).resolve().parent
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_dir,
trust_remote_code=True,
torch_dtype=dtype,
).to(device)
inputs = tokenizer("The future of language models", return_tensors="pt").to(device)
with torch.no_grad():
generated = model.generate(
**inputs,
max_new_tokens=30,
do_sample=True,
temperature=0.8,
top_k=50,
use_cache=False,
)
print(tokenizer.decode(generated[0], skip_special_tokens=True))
|