File size: 848 Bytes
df93a8c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | import os, torch
from transformers import AutoModelForCausalLM, AutoTokenizer
M="Qwen/Qwen2-0.5B"
tok=AutoTokenizer.from_pretrained(M, token=os.environ.get("HF_TOKEN"))
model=AutoModelForCausalLM.from_pretrained(M, torch_dtype=torch.bfloat16, device_map="cuda", token=os.environ.get("HF_TOKEN")).eval()
P=["The Company's total revenue for the fiscal year","Net income increased primarily due to",
"The capital of France is","Water is made of hydrogen and"]
for p in P:
ids=tok(p,return_tensors="pt").to("cuda")
out=model.generate(**ids,max_new_tokens=50,do_sample=True,temperature=0.8,top_k=40,top_p=0.95,
repetition_penalty=1.8,no_repeat_ngram_size=3)
txt=tok.decode(out[0][ids["input_ids"].shape[1]:],skip_special_tokens=True)
print(f"[stock] {p!r} -> {txt!r}",flush=True)
print("STOCK_DONE",flush=True)
|