Update README.md

README.md

@@ -8,13 +8,24 @@ pipeline_tag: text-classification
 
 加载代码：
 ```python
+import torch
+import os
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+
+# 载入Tokenizer与模型
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
-config = AutoConfig.from_pretrained("glm-emotion", trust_remote_code=True, pre_seq_len=128)
-model = AutoModel.from_pretrained("glm-emotion", config=config, trust_remote_code=True)
+config = AutoConfig.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, pre_seq_len=128)
+model = AutoModel.from_pretrained("THUDM/chatglm2-6b", config=config, trust_remote_code=True)
+# 载入预训练参数
 prefix_state_dict = torch.load(os.path.join("glm-emotion", "pytorch_model.bin"))
 new_prefix_state_dict = {}
 for k, v in prefix_state_dict.items():
     if k.startswith("transformer.prefix_encoder."):
         new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
 model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
+
+model.half().cuda().eval()
+
+response, history = model.chat(tokenizer, "说出下列句子的情绪类型：到北京五环25分钟车程，938离我住处只有300米，饮食一条街更是北京饮食街的翻版", history=[])
+print(response)
 ```