Update chat_ner.py

taishan1994 · May 24, 2023 · 2053893 · 2053893
1 parent bb5d759
commit 2053893
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/chat_ner.py b/chat_ner.py
@@ -1,3 +1,4 @@
+import os
 import torch
 import json
 from pprint import pprint
@@ -18,12 +19,12 @@
 
 model = AutoModel.from_pretrained(args["model_dir"],  trust_remote_code=True).half().cuda()
 model = model.eval()
-model = PeftModel.from_pretrained(model, args["save_dir"], torch_dtype=torch.float32, trust_remote_code=True)
+model = PeftModel.from_pretrained(model, os.path.join(args["save_dir"], "adapter_model"), torch_dtype=torch.float32, trust_remote_code=True)
 model.half().cuda()
 model.eval()
 
 while True:
     inp = input("用户 >>> ")
     response, history = model.chat(tokenizer, inp, history=[])
     print("ChatNER >>> ", response))
-    print("="*100)
+    print("="*100)