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)