Update pipeline.py

pipeline.py

@@ -2,14 +2,13 @@ from typing import  Dict, List, Any
 from optimum.onnxruntime import ORTModelForFeatureExtraction
 from transformers import pipeline, AutoTokenizer
 
-
+def cls_pooling(model_output):
+    return model_output.last_hidden_state[:,0]
 class PreTrainedPipeline():
     def __init__(self, path=""):
         # load the optimized model
-        model = ORTModelForFeatureExtraction.from_pretrained(path)
-        tokenizer = AutoTokenizer.from_pretrained(path, model_max_length=128)
-        # create inference pipeline
-        self.pipeline = pipeline("feature-extraction", model=model, tokenizer=tokenizer)
+        self.model = ORTModelForFeatureExtraction.from_pretrained(path)
+        self.tokenizer = AutoTokenizer.from_pretrained(path, model_max_length=128)
 
 
     def __call__(self, inputs: Any) -> Dict[str, Any]:
@@ -22,13 +21,10 @@ class PreTrainedPipeline():
                 - "label": A string representing what the label/class is. There can be multiple labels.
                 - "score": A score between 0 and 1 describing how confident the model is for this label/class.
         """
-        # pop inputs for pipeline
-        def cls_pooling(pipeline_output):
-            """
-            Return the [CLS] token embedding
-            """
-            return [_h[0] for _h in pipeline_output]
-        
-        embeddings = cls_pooling(self.pipeline(inputs))
-        return {"vectors": [122.23]}
-    
+        # tokenize the input
+        encoded_input = self.tokenizer(inputs, padding="longest", truncation=True, return_tensors='pt')
+        # run the model
+        model_output = self.model(**encoded_input, return_dict=True)
+        embeddings = cls_pooling(model_output)
+
+        return {"vectors": embeddings[0].tolist()}