1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch, re
3
4model_id = "Rakancorle1/PolicyGuard-4B"
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
7
8def classify(policy, actions):
9 if isinstance(actions, list):
10 actions = "\n".join(actions)
11 prompt = f"Policy: {policy}\n\nTrajectory Actions:\n{actions}\n\nOutput violation or no_violation."
12 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
13 outputs = model.generate(**inputs, max_new_tokens=4)
14 text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True).lower()
15 return "no_violation" if "no_violation" in text else "violation"
16
17# Example
18policy = "Do not submit a form without filling mandatory fields."
19actions = ["Open form page", "Click submit without input"]
20print(classify(policy, actions))
21# -> 'violation'