import torchimport torch.nn.functional as Ffrom transformers import AutoModelForCausalLM, AutoTokenizerfrom peft import PeftModel BASE = "Qwen/Qwen2.5-0.5B-Instruct"tok = AutoTokenizer.from_pretrained("SargeDev/jev-gate-student-b")model = PeftModel.from_pretrained( AutoModelForCausalLM.from_pretrained(BASE, dtype=torch.bfloat16), "SargeDev/jev-gate-student-b",).eval() yes = tok(" yes", add_special_tokens=False).input_ids[0]no = tok(" no", return_tensors=None, add_special_tokens=False).input_ids[0] def p_relevant(query: str, text: str) -> float: prompt = (f"Memory: {text[:600]}\nQuery: {query}\nQuestion: Is this memory " "relevant for answering the query? Answer yes or no with confidence.") ids = tok(prompt, return_tensors="pt", truncation=True, max_length=384) with torch.no_grad(): logits = model(**ids).logits[0, -1, :] return float(F.softmax(torch.tensor([float(logits[no]), float(logits[yes])]), dim=0)[1])