Linsad commited on
Commit
0e501f1
1 Parent(s): 50aa043

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +3 -3
handler.py CHANGED
@@ -2,10 +2,10 @@ 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
 
 
2
  from transformers import AutoTokenizer, AutoModel
3
 
4
 
5
+ class EndpointHandler:
6
  def __init__(self):
7
+ self.tokenizer = AutoTokenizer.from_pretrained("chatglm2-6b-int4", trust_remote_code=True)
8
+ self.model = AutoModel.from_pretrained("chatglm2-6b-int4", trust_remote_code=True).half().cuda()
9
  self.model = self.model.eval()
10
 
11