Update handler.py

handler.py

@@ -1,4 +1,5 @@
 from transformers import pipeline
+from sklearn.preprocessing import LabelEncoder
 import joblib
 import torch
 import os
@@ -8,7 +9,7 @@ print("Current working directory:", os.getcwd())
 print("Contents of the directory:", os.listdir())
 
 # Load the label encoder
-label_encoder = joblib.load('label_encoder.pkl')  # Ensure the file is in the correct path
+label_encoder = joblib.load('/repository/label_encoder.pkl')  # Use absolute path
 print("Label encoder loaded successfully.")
 
 # Load the model and tokenizer from Hugging Face
@@ -39,16 +40,14 @@ def get_average_sentiment(positive_count, negative_count, neutral_count):
         return "neutral"
 
 class EndpointHandler:
-    def __init__(self):
-        # No need to load the model here since it's loaded globally
+    def __init__(self, model_dir=None):
+        # Model and tokenizer are loaded globally, so no need to reinitialize here
+        # The `model_dir` argument is required by Hugging Face's inference toolkit
         pass
 
     def preprocess(self, data):
         # Extract the input text from the request
-        if isinstance(data, dict):
-            text = data.get("inputs", "")
-        else:
-            text = data  # Fallback if data is not a dictionary
+        text = data.get("inputs", "")
         return text
 
     def inference(self, text):
@@ -116,16 +115,14 @@ class EndpointHandler:
 
     def postprocess(self, output):
         if "error" in output:
-            return {"error": output["error"]}
+            return [{"error": output["error"]}]
         
-        # Return the full output
-        return output
+        # Return only the line-level results as a list
+        return output["line_results"]
+
 
     def __call__(self, data):
         # Main method to handle the request
         text = self.preprocess(data)
         output = self.inference(text)
-        return self.postprocess(output)
-
-# Create an instance of the handler
-handler = EndpointHandler()
+        return self.postprocess(output)