philschmid HF staff commited on
Commit
d7d7e6d
1 Parent(s): f078d4c

Create mii-handler.py

Browse files
Files changed (1) hide show
  1. mii-handler.py +34 -0
mii-handler.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ import mii
3
+ import json
4
+
5
+ class EndpointHandler():
6
+ def __init__(self, path=""):
7
+ self.deploy_name = "bert"
8
+ mii_config = {"dtype": "fp16"}
9
+ mii.deploy(task='text-classification',
10
+ model="philschmid/finbert-tone-endpoint-ds",
11
+ deployment_name=self.deploy_name,
12
+ mii_config=mii_config)
13
+ # create handler for server
14
+ self.pipeline = mii.mii_query_handle(self.deploy_name)
15
+
16
+
17
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
18
+ """
19
+ data args:
20
+ inputs (:obj: `str`)
21
+ date (:obj: `str`)
22
+ Return:
23
+ A :obj:`list` | `dict`: will be serialized and returned
24
+ """
25
+ inputs = data.pop("inputs", data)
26
+ parameters = data.pop("parameters", None)
27
+
28
+ # pass inputs with all kwargs in data
29
+ if parameters is not None:
30
+ prediction = self.pipeline.query({"query": inputs}, **parameters)
31
+ else:
32
+ prediction = self.pipeline.query({"query": inputs})
33
+ # postprocess the prediction
34
+ return prediction.response