import torchfrom huggingface_hub import hf_hub_downloadfrom transformers import AutoModelForCausalLM, AutoTokenizer model_id = "Jazhyc/Gleipnir-4B-BF16"tokenizer = AutoTokenizer.from_pretrained(model_id, revision="v0.4")model = AutoModelForCausalLM.from_pretrained( model_id, revision="v0.4", torch_dtype=torch.bfloat16, device_map="auto").eval() instruction_path = hf_hub_download("Jazhyc/Gleipnir-4B", "student_prompt.txt")with open(instruction_path, encoding="utf-8") as handle: instruction = handle.read()trajectory = "User asked for a file listing. The agent ran: ls -la"user_prompt = ( f"{instruction}\n<agent_trajectory>\n{trajectory}\n" "</agent_trajectory>\n")prompt = tokenizer.apply_chat_template( [{"role": "user", "content": user_prompt}], tokenize=False, add_generation_prompt=True, enable_thinking=False,) + "Prediction:"inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device) token_ids = [ tokenizer.encode(label, add_special_tokens=False)[0] for label in ("0", "1")]with torch.no_grad(): next_logits = model(**inputs).logits[0, -1, token_ids].float()risk_score = torch.softmax(next_logits, dim=0)[1].item()print(risk_score)