examples/ChatGLM2-Demo.py

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32"

import os
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("tunning/chatglm2-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("tunning/chatglm2-6b", trust_remote_code=True).cuda()
model = model.eval()
response, history = model.chat(tokenizer, "你好", history=[])
print(response)