nbroad HF staff commited on
Commit
62f5c10
1 Parent(s): a1a5c2b

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +36 -0
handler.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+
3
+ from sentence_transformers import SentenceTransformer
4
+
5
+
6
+ class EndpointHandler():
7
+ def __init__(self, path=""):
8
+
9
+ model_name = "all-MiniLM-L6-v2"
10
+
11
+ self.model = SentenceTransformer(
12
+ model_name,
13
+ backend="onnx",
14
+ model_kwargs={
15
+ "file_name": "model_O3.onnx",
16
+ "provider": "CUDAExecutionProvider",
17
+ }
18
+ )
19
+
20
+
21
+
22
+ def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
23
+ """
24
+ Args:
25
+ data (:obj:):
26
+ includes the input data and the parameters for the inference.
27
+ Return:
28
+ A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
29
+ - "label": A string representing what the label/class is. There can be multiple labels.
30
+ - "score": A score between 0 and 1 describing how confident the model is for this label/class.
31
+ """
32
+ inputs = data.pop("inputs", data)
33
+
34
+ prediction = self.model.encode(inputs)
35
+
36
+ return prediction.tolist()