Update sentpipeline.py

sentpipeline.py

@@ -2,21 +2,36 @@ from transformers import Pipeline
 from sentence_transformers import SentenceTransformer
 import torch
 
+
 class SentimentModelPipe(Pipeline):
     def __init__(self, **kwargs):
         Pipeline.__init__(self, **kwargs)
-        self.smodel = SentenceTransformer(kwargs.get("embedding_model", "sentence-transformers/all-MiniLM-L6-v2"))
+        self.smodel = SentenceTransformer(
+            kwargs.get("embedding_model", "sentence-transformers/all-MiniLM-L6-v2")
+        )
+        self.class_map = kwargs.get(
+            "class_map",
+            {0: "sad", 1: "joy", 2: "love", 3: "anger", 4: "fear", 5: "surprise"},
+        )
 
     def _sanitize_parameters(self, **kw):
         return {}, {}, {}
 
     def preprocess(self, inputs):
         return self.smodel.encode(inputs, convert_to_tensor=True)
-    
+
     def postprocess(self, outputs):
-        return outputs
-    
+        if isinstance(outputs, torch.Tensor):
+            outputs = [outputs]
+        results = []
+        for out in outputs:
+            r = []
+            for i, l in enumerate(out):
+                r.append({"label": self.class_map[i], "score": l.item()})
+            results.append(r)
+        return results
+
     def _forward(self, tensor):
         with torch.no_grad():
             out = self.model(tensor)
-        return out
+        return out