Update pipeline.py

pipeline.py

@@ -1,12 +1,12 @@
 from typing import  Dict, List, Any
-from optimum.onnxruntime import ORTModelForSequenceClassification
+from optimum.onnxruntime import ORTModelForFeatureExtraction
 from transformers import pipeline, AutoTokenizer
 
 
 class PreTrainedPipeline():
     def __init__(self, path=""):
         # load the optimized model
-        model = ORTModelForSequenceClassification.from_pretrained(path)
+        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)
@@ -30,4 +30,5 @@ class PreTrainedPipeline():
             return [_h[0] for _h in pipeline_output]
         
         embeddings = cls_pooling(self.pipeline(inputs))
-        return embeddings
+        return {"vectors": embeddings}
+