infer_test.py 647 Bytes
Newer Older
zhaoying1's avatar
zhaoying1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("../../chatglm3-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("../../chatglm3-6b", trust_remote_code=True, device='cuda')
model = model.eval()
response, history = model.chat(tokenizer, "你好", history=[])
print(response)
response, history = model.chat(tokenizer, "晚上睡不着应该怎么办", history=history)
print(response)
# print(len(tokenizer))
# vocab_content = tokenizer.get_vocab()
# with open("vocab.txt", "w", encoding="utf-8") as f:
#     for token, index in vocab_content.items():
#         f.write(f"{token} {index}\n")