Update handler.py

handler.py

@@ -3,21 +3,22 @@ from transformers import AutoTokenizer, AutoModel
 
 
 class EndpointHandler:
-    def __init__(self):
-        self.tokenizer = AutoTokenizer.from_pretrained("chatglm2-6b-int4", trust_remote_code=True)
-        self.model = AutoModel.from_pretrained("chatglm2-6b-int4", trust_remote_code=True).half().cuda()
+    def __init__(self, path=""):
+        print('path is' + path)
+        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+        self.model = AutoModel.from_pretrained(path, trust_remote_code=True).half().cuda()
         self.model = self.model.eval()
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-       data args:
-            inputs (:obj: `str`)
-      Return:
-            A :obj:`list` | `dict`: will be serialized and returned
+           data args:
+                inputs (:obj: `str`)
+          Return:
+                A :obj:`list` | `dict`: will be serialized and returned
         """
         # get inputs
         inputs = data.pop("inputs", data)
-    
+
         response, history = self.model.chat(self.tokenizer, inputs, history=[])
-    
+
         return [{'response': response, 'history': history}]