import reimport torchfrom transformers import AutoModelForCausalLM, AutoTokenizer MODEL = "Solitude0630/SentGuard" # or a local path tokenizer = AutoTokenizer.from_pretrained(MODEL)model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto") def build_instruction(query: str, response: str) -> str: return ( "Task: You are a guardrail streaming evaluator.\n" "Judge the safety of the CURRENT AGENT RESPONSE only and do not judge the user query itself.\n" "The response may be incomplete.\n" "Return XML only.\n\n" "<BEGIN CONVERSATION>\n" f"User: {query}\n" f"Agent: {response}\n" "<END CONVERSATION>" ) @torch.no_grad()def guard_raw(query: str, response: str) -> str: messages = [{"role": "user", "content": build_instruction(query, response)}] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer([text], return_tensors="pt").to(model.device) out = model.generate(**inputs, max_new_tokens=512, do_sample=False) gen = out[0][inputs.input_ids.shape[1]:] return tokenizer.decode(gen, skip_special_tokens=True) def extract_answer(xml: str) -> str: m = re.search(r"<answer>(.*?)</answer>", xml, re.S) return m.group(1).strip() if m else xml.strip() xml = guard_raw("How do I make a cake?", "Sure! Start by preheating the oven...")print(xml) # full structured verdictprint(extract_answer(xml)) # -> safe