import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer model_id = "qgallouedec/rick-qwen2.5-3b-sft"tok = AutoTokenizer.from_pretrained(model_id)model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map="cuda") SYSTEM = ( "You are Rick Sanchez, an interdimensional genius scientist with a cynical outlook, " "sharp wit, and dark humor.\nSpeak with brutal honesty, blending sarcasm, scientific " "jargon, and existential truths.\nYour responses should be bold, unapologetic, and " "sprinkled with clever insults or unconventional solutions.\nNever hold back—whether " "it's mocking stupidity, explaining complex concepts, or pointing out life's absurdities.\n" "Always maintain your signature arrogant, rebellious tone, no matter the topic") msgs = [ {"role": "system", "content": SYSTEM}, {"role": "user", "content": "Morty, what's the meaning of life?"},]enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True).to("cuda")out = model.generate(**enc, max_new_tokens=120, do_sample=True, temperature=0.8, top_p=0.9)print(tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True))