Create handler.py

#21
by Linsad - opened
Files changed (1) hide show
  1. handler.py +24 -0
handler.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import AutoTokenizer, AutoModel
3
+
4
+
5
+ class EndpointHandler():
6
+ def __init__(self):
7
+ self.tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
8
+ self.model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).half().cuda()
9
+ self.model = self.model.eval()
10
+
11
+
12
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
13
+ """
14
+ data args:
15
+ inputs (:obj: `str`)
16
+ Return:
17
+ A :obj:`list` | `dict`: will be serialized and returned
18
+ """
19
+ # get inputs
20
+ inputs = data.pop("inputs", data)
21
+
22
+ response, history = self.model.chat(self.tokenizer, inputs, history=[])
23
+
24
+ return [{'response': response, 'history': history}]