Update handler.py

handler.py

@@ -15,9 +15,9 @@ class EndpointHandler:
         """
         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)
+        model_name = "./json_extraction_point_activity"  # Pretrained model for sentiment analysis
+        self.tokenizer = T5Tokenizer.from_pretrained(model_name)
+        self.model = T5ForConditionalGeneration.from_pretrained(model_name)
         self.model.eval()  # Set model to evaluation mode (no training)
         print(f"Loaded model: {model_name}")
 
@@ -41,21 +41,18 @@ class EndpointHandler:
             current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
             # Tokenize input text
-            encoded_input = self.tokenizer(user_text, return_tensors="pt")
+            input_ids = self.tokenizer(user_text, return_tensors="pt").input_ids
 
             # Perform inference
             with torch.no_grad():
-                outputs = self.model(**encoded_input)
-
-            # 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)}"}
+                output_ids = model.generate(input_ids, max_length=100, temperature=0.3)
+
+            json_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+            
+            try:
+                return json.loads(json_output)
+            except:
+                return json_output
+    
+            except Exception as e:
+                return {"error": f"Unexpected error: {str(e)}"}