Create handler.py

handler.py

@@ -0,0 +1,46 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+import torch.nn.functional as F
+from peft import PeftModel
+
+
+class EndpointHandler:
+    def __init__(self, model_dir):
+        """
+        Initialize the model and tokenizer using the provided model directory.
+        """
+        model_name = "munzirmuneer/phishing_url_gemma_pytorch"  # Replace with your specific model
+        model_name2 = "google/gemma-2b"
+        
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name2)
+        base_model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        self.model = PeftModel.from_pretrained(base_model, model_name)
+
+    def __call__(self, input_data):
+        """
+        Perform inference on the input text and return predictions.
+        """
+                # Extract the URL from the input_data dictionary
+        if 'inputs' in input_data:
+            input_text = input_data['inputs']  # Expecting a single URL as a string
+        else:
+            raise ValueError("Input data must contain the 'inputs' key with a URL.")
+        # Tokenize input
+        inputs = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
+        
+        # Run inference
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        
+        # Get logits and probabilities
+        logits = outputs.logits
+        probs = F.softmax(logits, dim=-1)
+        
+        # Get the predicted class (highest probability)
+        pred_class = torch.argmax(probs, dim=-1)
+        
+        return {
+            "predicted_class": pred_class.item(),
+            "probabilities": probs[0].tolist()
+        }