Update handler.py

handler.py

@@ -1,22 +1,29 @@
 import json
 from datetime import datetime
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
 
 class EndpointHandler:
     def __init__(self, model_dir):
         self.model_dir = model_dir
 
     def load(self):
-        print("Model loading skipped for Hello World API.")
+        """
+        Load a simple DistilBERT model for text classification.
+        """
+        model_name = "distilbert-base-uncased-finetuned-sst-2-english"  # Pretrained model for sentiment analysis
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        print(f"Loaded model: {model_name}")
 
     def __call__(self, inputs):
         """
-        Hugging Face expects the `EndpointHandler` to be callable.
-        So, we define `__call__` instead of `predict`.
+        Process user input and classify the text using DistilBERT.
         """
         try:
-            # Hugging Face sends input as a LIST or DICT, so handle both cases
+            # Handle different input formats
             if isinstance(inputs, list) and len(inputs) > 0:
-                user_text = inputs[0]  # Extract text from the list
+                user_text = inputs[0]
             elif isinstance(inputs, dict) and "inputs" in inputs:
                 user_text = inputs["inputs"]
             else:
@@ -25,8 +32,22 @@ class EndpointHandler:
             # Generate timestamp
             current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
-            # Return formatted message
-            return {"message": f"Received at {current_time}: {user_text}"}
+            # Tokenize input text
+            inputs = self.tokenizer(user_text, return_tensors="pt")
+
+            # Perform inference
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+
+            # Get predicted label (0 = negative, 1 = positive)
+            predicted_label = torch.argmax(outputs.logits).item()
+            label_map = {0: "negative", 1: "positive"}
+
+            return {
+                "timestamp": current_time,
+                "input_text": user_text,
+                "predicted_label": label_map[predicted_label]
+            }
 
         except Exception as e:
             return {"error": f"Unexpected error: {str(e)}"}