スパースベクトル化に対応する

.gitignore

@@ -0,0 +1 @@
+test*

handler.py

@@ -1,15 +1,14 @@
 from typing import  Dict, List, Any
-from optimum.onnxruntime import ORTModelForSequenceClassification
-from transformers import pipeline, AutoTokenizer
+# from optimum.onnxruntime import ORTModelForSequenceClassification
+# from transformers import pipeline, AutoTokenizer
+from FlagEmbedding import BGEM3FlagModel
 
 
 class EndpointHandler():
-    def __init__(self, path="./onnx"):
+    def __init__(self, path="./"):
         # load the optimized model
-        model = ORTModelForSequenceClassification.from_pretrained(path)
-        tokenizer = AutoTokenizer.from_pretrained(path)
-        # create inference pipeline
-        self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
+        # モデルの準備
+        model = BGEM3FlagModel("./", use_fp16=False)
 
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
@@ -24,11 +23,20 @@ class EndpointHandler():
         """
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
+        
+        sparse_embs = []
+        
+        result = self.model.encode(inputs, return_dense=False, return_sparse=True)
+        sparse_vectors = result["lexical_weights"]
+        for sparse_vector in sparse_vectors:
+            sparse_values = [value for value in sparse_vector.values()]
+            sparse_dimensions = [int(key) for key in sparse_vector.keys()]
+            sparse_embs.append((sparse_values, sparse_dimensions))
 
         # pass inputs with all kwargs in data
-        if parameters is not None:
-            prediction = self.pipeline(inputs, **parameters)
-        else:
-            prediction = self.pipeline(inputs)
+        # if parameters is not None:
+        #     prediction = self.pipeline(inputs, **parameters)
+        # else:
+        #     prediction = self.pipeline(inputs)
         # postprocess the prediction
-        return prediction
+        return sparse_vectors