import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
base = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-7B-Instruct", torch_dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(base, "rotalabs/sakshi-judge-indic")
tok = AutoTokenizer.from_pretrained("rotalabs/sakshi-judge-indic")
SYSTEM = ("You are a groundedness judge for banking documents. Given a CONTEXT "
"and an ANSWER, decide whether every claim in the answer is supported "
"by the context. Reply with exactly one word: grounded or hallucinated.")
prompt = tok.apply_chat_template(
[{"role": "system", "content": SYSTEM},
{"role": "user", "content": "CONTEXT:\n...\n\nANSWER:\n...\n\nVerdict:"}],
add_generation_prompt=True, return_tensors="pt").to(model.device)
logits = model(prompt).logits[0, -1]
g = tok("grounded", add_special_tokens=False)["input_ids"][0]
h = tok("hallucinated", add_special_tokens=False)["input_ids"][0]
probs = torch.softmax(torch.stack([logits[g], logits[h]]).float(), dim=0)