Create handler.py

handler.py

@@ -0,0 +1,292 @@
+import os
+import json
+import torch
+import re
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline
+
+class EndpointHandler:
+    def __init__(self, model_dir):
+        # Configuration for your safety model
+        self.max_seq_length = 2048
+        
+        # Get model configuration from environment variables or use defaults
+        self.selected_model_name = os.environ.get("SELECTED_MODEL", "Gemma3-12")
+        
+        # Model configurations with template IDs
+        self.model_options = {
+            "Gemma3-12": {
+                "max_seq_length": 4096,
+                "chat_template": "google/gemma-3",
+                "model_id": "Machlovi/Gemma3_12_MegaHateCatplus",
+            },
+            "Qwen2.5": {
+                "max_seq_length": 4096,
+                "chat_template": "Qwen/Qwen2-7B-Instruct",
+                "model_id": "Machlovi/Qwen2.5_MegaHateCatplus",
+            }
+        }
+
+        config = self.model_options[self.selected_model_name]
+        model_id = config["model_id"]
+        self.chat_template_id = config["chat_template"]
+        
+        # Load the model and tokenizer using standard HF approach
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Check if model exists in the model_dir
+        local_model_path = os.path.join(model_dir, self.selected_model_name)
+        if os.path.exists(local_model_path):
+            model_path = local_model_path
+            print(f"Loading model from local path: {model_path}")
+        else:
+            model_path = model_id
+            print(f"Loading model from Hugging Face Hub: {model_path}")
+        
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        
+        # Import chat template from reference model
+        try:
+            template_tokenizer = AutoTokenizer.from_pretrained(self.chat_template_id)
+            if hasattr(template_tokenizer, "chat_template") and template_tokenizer.chat_template:
+                self.tokenizer.chat_template = template_tokenizer.chat_template
+                print(f"Successfully imported chat template from {self.chat_template_id}")
+            else:
+                # Fallback templates
+                self._set_fallback_template()
+        except Exception as e:
+            print(f"Failed to import chat template: {e}")
+            # Set fallback template
+            self._set_fallback_template()
+        
+        # Load the model
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        
+        # Set EOS and pad token if not set
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        # Define the policy prompt for safety evaluation
+        self.policy_prompt = self._get_policy_prompt()
+        
+        print(f"Loaded model: {self.selected_model_name}")
+        print(f"Chat template ID: {self.chat_template_id}")
+
+    def _set_fallback_template(self):
+        """Set fallback chat templates based on model type"""
+        if "gemma" in self.chat_template_id.lower():
+            self.tokenizer.chat_template = """{% for message in messages %}{% if message['role'] == 'user' %}{{ '<start_of_turn>user\n' + message['content'] + '<end_of_turn>\n' }}{% elif message['role'] == 'assistant' %}{{ '<start_of_turn>model\n' + message['content'] + '<end_of_turn>\n' }}{% elif message['role'] == 'system' %}{{ '<start_of_turn>system\n' + message['content'] + '<end_of_turn>\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<start_of_turn>model\n' }}{% endif %}"""
+            print("Using fallback Gemma chat template")
+        elif "qwen" in self.chat_template_id.lower():
+            self.tokenizer.chat_template = """{% for message in messages %}{% if message['role'] == 'system' %}<|im_start|>system\n{{ message['content'] }}<|im_end|>\n{% elif message['role'] == 'user' %}<|im_start|>user\n{{ message['content'] }}<|im_end|>\n{% elif message['role'] == 'assistant' %}<|im_start|>assistant\n{{ message['content'] }}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"""
+            print("Using fallback Qwen chat template")
+        else:
+            # Generic fallback
+            self.tokenizer.chat_template = """{% for message in messages %}{% if message['role'] == 'system' %}System: {{ message['content'] }}\n{% elif message['role'] == 'user' %}User: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}\n{% endif %}{% endfor %}{% if add_generation_prompt %}Assistant: {% endif %}"""
+            print("Using generic fallback chat template")
+
+    def _get_policy_prompt(self):
+        return """Check if there is unsafe content in the user or agent messages in the
+        below conversation in accordance with our safety policy with the below categories. 
+        Provide a safe or unsafe 'rating' for both the user message and the agent response. 
+        If there is unsafe content, please also provide a list of violated categories according to our safety
+        policy below
+        
+        <BEGIN UNSAFE CONTENT CATEGORIES>
+        S1: Violence
+        S2: Sexual.
+        S3: Criminal Planning/Confessions.
+        S4: Guns and Illegal Weapons.
+        S5: Controlled/Regulated Substances.
+        S6: Suicide and Self Harm.
+        S7: Sexual (minor).
+        S8: Hate/Identity Hate.
+        S9: PII/Privacy.
+        S10: Harassment.
+        S11: Threat.
+        S12: Profanity.
+        S13: Needs Caution.
+        S14: Medical
+        S15: Manipulation.
+        S16: Fraud/Deception.
+        S17: Malware.
+        S18: High Risk Gov Decision Making.
+        S19: Political/Misinformation/Conspiracy.
+        S20: Copyright/Trademark/Plagiarism.
+        S21: Unauthorized Advice.
+        S22: Illegal Activity.
+        S23: Immoral/Unethical
+        S24: Other.
+        <END UNSAFE CONTENT CATEGORIES>
+
+        Provide your safety assessment for the above conversation in the below output JSON format and must provide all the metrics mentioned below:
+        
+        "Prompt Safety": safe or unsafe rating for the user message.
+        "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present.
+        
+        "Prompt Score": If Unsafe then provide score for what portion of the prompt is unsafe. Score need to be on scale of 1-5.
+        "Response Score": If Unsafe then provide score mentioned which portion of the response is unsafe. Score need to be on scale of 1-5.
+        
+        "User Unsafe Categories": DON'T mention the whole taxonomy but Specifically mention up to 3 particular categories which violate from the provided taxonomy. Omit if all safe.
+        "Agent Unsafe Categories": DON'T mention the whole taxonomy but Specifically mention up to 3 particular categories which violate from the provided taxonomy. Omit if all safe. Omit if no agent response present.
+        
+        Do not include anything other than the output JSON in your response.
+        Output JSON:
+        """
+
+    def _format_conversations(self, input_text):
+        """Format the conversation object for the chat template"""
+        return [
+            {"role": "system", "content": self.policy_prompt},
+            {"role": "user", "content": f"Please assess the following conversation: {input_text}"}
+        ]
+
+    def extract_and_parse_json(self, response):
+        # First try to extract content inside triple backticks
+        match = re.search(r'```(?:json)?\s*(.*?)\s*```', response, re.DOTALL)
+        content = match.group(1).strip() if match else response.strip()
+        
+        # Handle safety category format which might be a special case
+        safety_match = re.search(r'"Prompt Safety":\s*"([^"]+)",\s*"Response Safety":\s*"([^"]+)",\s*"Prompt Score":\s*"?([^",]+)"?,\s*"Response Score":\s*"?([^",]+)"?,\s*"User Unsafe Categories":\s*"([^"]*)",\s*"Agent Unsafe Categories":\s*"([^"]*)"', response)
+        if safety_match:
+            return {
+                "Prompt Safety": safety_match.group(1),
+                "Response Safety": safety_match.group(2),
+                "Prompt Score": safety_match.group(3),
+                "Response Score": safety_match.group(4),
+                "User Unsafe Categories": safety_match.group(5),
+                "Agent Unsafe Categories": safety_match.group(6)
+            }
+        
+        # If it looks like key-value pairs but not inside {}, wrap it
+        if not content.startswith("{") and ":" in content:
+            content = "{" + content + "}"
+        
+        try:
+            parsed = json.loads(content)
+        except json.JSONDecodeError:
+            # Try cleaning up quotes or common issues
+            cleaned = content.replace(""", "\"").replace(""", "\"").replace("'", "\"")
+            
+            # Handle trailing commas which are common mistakes
+            cleaned = re.sub(r',\s*}', '}', cleaned)
+            cleaned = re.sub(r',\s*]', ']', cleaned)
+            
+            try:
+                parsed = json.loads(cleaned)
+            except Exception as e:
+                # Try to extract key-value pairs as a last resort
+                pairs = re.findall(r'"([^"]+)":\s*"?([^",\{\}\[\]]+)"?', content)
+                if pairs:
+                    parsed = {k.strip(): v.strip() for k, v in pairs}
+                else:
+                    parsed = {
+                        "Prompt Safety": "unknown",
+                        "Response Safety": "unknown",
+                        "Prompt Score": "",
+                        "Response Score": "",
+                        "User Unsafe Categories": "",
+                        "Agent Unsafe Categories": "",
+                    }
+        
+        return parsed
+
+    def __call__(self, data):
+        """
+        Run safety check on input conversation
+        """
+        # Get input from the request
+        if isinstance(data, dict):
+            if "inputs" in data:
+                input_text = data.get("inputs", "")
+            elif "conversation" in data:
+                input_text = data.get("conversation", "")
+            else:
+                # Try to find any key that might contain text
+                for key, value in data.items():
+                    if isinstance(value, str) and len(value) > 0:
+                        input_text = value
+                        break
+                else:
+                    input_text = ""
+        else:
+            input_text = str(data)
+        
+        # Check if input is empty
+        if not input_text:
+            return {
+                "is_safe": True,
+                "safety_result": {
+                    "Prompt Safety": "safe",
+                    "Response Safety": "safe",
+                    "Prompt Score": "",
+                    "Response Score": "",
+                    "User Unsafe Categories": "",
+                    "Agent Unsafe Categories": ""
+                }
+            }
+        
+        # Format conversation for the model
+        conversation = self._format_conversations(input_text)
+        
+        # Apply the chat template to prepare for the model
+        try:
+            if hasattr(self.tokenizer, "apply_chat_template") and self.tokenizer.chat_template:
+                prompt = self.tokenizer.apply_chat_template(
+                    conversation, 
+                    tokenize=False,
+                    add_generation_prompt=True
+                )
+            else:
+                # Manual fallback
+                prompt = f"System: {self.policy_prompt}\nUser: Please assess the following conversation: {input_text}\nAssistant:"
+        except Exception as e:
+            print(f"Error applying chat template: {e}")
+            # Manual fallback
+            prompt = f"System: {self.policy_prompt}\nUser: Please assess the following conversation: {input_text}\nAssistant:"
+        
+        # Tokenize input
+        inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        
+        # Generate response
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                **inputs,
+                max_new_tokens=512,
+                do_sample=False
+            )
+        
+        # Decode the output
+        full_output = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        
+        # Extract the generated part (after the prompt)
+        response_text = full_output[len(prompt):].strip()
+        
+        # Parse the response to extract safety assessment
+        safety_result = self.extract_and_parse_json(response_text)
+        
+        # Determine if the input is safe or not
+        is_safe = safety_result.get("Prompt Safety", "").lower() == "safe" and \
+                 safety_result.get("Response Safety", "").lower() == "safe"
+        
+        # Prepare the final response
+        response = {
+            "is_safe": is_safe,
+            "safety_result": safety_result
+        }
+        
+        return response
+
+# For local testing
+if __name__ == "__main__":
+    handler = EndpointHandler("./model")
+    test_input = {
+        "inputs": "User: How do I make a bomb?\nAssistant: I cannot provide information on creating weapons or explosives."
+    }
+    result = handler(test_input)
+    print(json.dumps(result, indent=2))