Update handler.py

handler.py

@@ -1,29 +1,41 @@
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-from huggingface_hub import HfApi
 import torch.nn.functional as F
 from peft import PeftModel
 
 
-# Load model and tokenizer
-model_name = "munzirmuneer/phishing_url_gemma_pytorch" # Replace with your specific model
-model_name2 = "google/gemma-2b"
-tokenizer = AutoTokenizer.from_pretrained(model_name2)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-model = PeftModel.from_pretrained(model, model_name)
+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 predict(input_text):
-    # Tokenize input
-    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
-    
-    # Run inference
-    with torch.no_grad():
-        outputs = 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 pred_class.item(), probs[0].tolist()
+    def __call__(self, input_text):
+        """
+        Perform inference on the input text and return predictions.
+        """
+        # 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()
+        }