from transformers import AutoTokenizer, AutoModelForCausalLMimport torch, re model_id = "Rakancorle1/PolicyGuard-4B"tokenizer = AutoTokenizer.from_pretrained(model_id)model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto") def classify(policy, actions): if isinstance(actions, list): actions = "\n".join(actions) prompt = f"Policy: {policy}\n\nTrajectory Actions:\n{actions}\n\nOutput violation or no_violation." inputs = tokenizer(prompt, return_tensors="pt").to(model.device) outputs = model.generate(**inputs, max_new_tokens=4) text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True).lower() return "no_violation" if "no_violation" in text else "violation" # Examplepolicy = "Do not submit a form without filling mandatory fields."actions = ["Open form page", "Click submit without input"]print(classify(policy, actions))# -> 'violation'