import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer name = "beyoru/linne-fern-1-0"tok = AutoTokenizer.from_pretrained(name)model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16, device_map="auto") state = "My Stripe integration keeps failing. Please help ASAP."question = "Which team should handle this request?"options = {"billing": "Charges and refunds", "technical_support": "Integration errors", "sales": "Questions about buying a product"} letters = [chr(65 + i) for i in range(len(options))]lines = [f"{l}: {k} - {v}" for l, (k, v) in zip(letters, options.items())]user = "\n".join([state, "", f"Question: {question}", *lines, "Answer with the letter of one option only."])prompt = tok.apply_chat_template([{"role": "user", "content": user}], tokenize=False, add_generation_prompt=True, enable_thinking=False)ids = tok(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)with torch.no_grad(): logits = model(**ids).logits[0, -1]label_ids = [tok(prompt + l, add_special_tokens=False)["input_ids"][-1] for l in letters]probs = torch.softmax(logits[label_ids].float(), dim=0)print({k: round(float(p), 4) for k, p in zip(options, probs)})