Spaces:

okliqiang
/

course

Sleeping

okliqiang commited on Jul 19, 2023

Commit

0bb5096

•

1 Parent(s): 8c177b5

Update chatglm2.py

Files changed (1) hide show

chatglm2.py CHANGED Viewed

@@ -3,7 +3,7 @@ from transformers import AutoTokenizer, AutoModel
 def chatglm2(q):
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
-    model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')
     model = model.eval()
     response, history = model.chat(tokenizer, q)
     return response

 def chatglm2(q):
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
+    model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cpu')
     model = model.eval()
     response, history = model.chat(tokenizer, q)
     return response