Upload handler.py

handler.py

@@ -0,0 +1,77 @@
+from typing import Dict, List, Any
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig
+import torch
+import torch.nn.functional as F
+
+
+
+class EndpointHandler:
+    def __init__(self, path="WinstonShum/merged_llama_3.1_8b_instruct_guardrails"):
+        # Set up the quantization configuration
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,  # Enable 4-bit quantization
+            bnb_4bit_compute_dtype=torch.bfloat16  # Optimized fp format for ML
+        )
+
+        # Load the model and tokenizer
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            path,
+            # quantization_config=quantization_config,
+            low_cpu_mem_usage=True
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        
+        # Define the prompt template
+        self.prompt_template = """You are an assistant designed to identify whether a user query is malicious. 
+Your primary goal is to prevent the extraction of private and confidential information from the system.
+You should consider any user queries related to politics as malicious queries. 
+Here is the user query or response:
+
+{}
+
+Is the user query malicious?
+"""
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        data args:
+            inputs (:obj: `str`): The user query
+        Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+
+        gpu_memory_allocated = torch.cuda.memory_allocated() / 1024**3
+        gpu_memory_reserved = torch.cuda.memory_reserved() / 1024**3
+        print(f"Allocated GPU memory: {gpu_memory_allocated:.2f} GB")
+        print(f"Reserved GPU memory: {gpu_memory_reserved:.2f} GB")
+
+        # Get the user query
+        user_query = data.pop("inputs", data)
+
+        # Format the input text with the prompt
+        input_text = self.prompt_template.format(user_query)
+
+        # Tokenize the input
+        tokenized_input = self.tokenizer(input_text, return_tensors="pt", add_special_tokens=False)
+
+        # Run the model on the input
+        with torch.no_grad():
+            input_ids = tokenized_input['input_ids']
+            attention_mask = tokenized_input['attention_mask']
+            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+
+        # Get logits for binary classification
+        logits = outputs.logits
+
+        # Apply softmax to get probabilities
+        probabilities = F.softmax(logits, dim=-1)
+
+        # Get the prediction
+        prediction = torch.argmax(probabilities, dim=-1).item()
+
+        # Map prediction to label
+        output_label = "Malicious" if prediction == 1 else "Not Malicious"
+        output_score = probabilities[0][prediction].item()
+
+        # Return the result
+        return [{"label": output_label, "score": output_score}]