Update app.py

app.py

@@ -41,12 +41,13 @@ def init_args():
     args = load_hyperparam(args)
 
     # args.tokenizer = Tokenizer(model_path=args.spm_model_path)
-    args.tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", trust_remote_code=True)
+    args.tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", use_fast=False, trust_remote_code=True)
     args.vocab_size = args.tokenizer.sp_model.vocab_size()
 
 
 def init_model():
     global lm_generation
+    global model
     # torch.set_default_tensor_type(torch.HalfTensor)
     # model = LLaMa(args)
     # torch.set_default_tensor_type(torch.FloatTensor)
@@ -64,11 +65,12 @@ def init_model():
 
 
 def chat(prompt, top_k, temperature):
-    args.top_k = int(top_k)
-    args.temperature = temperature
-    response = lm_generation.generate(args, [prompt])
-    print('log:', response[0])
-    return response[0]
+    # args.top_k = int(top_k)
+    # args.temperature = temperature
+    # response = lm_generation.generate(args, [prompt])
+    response = model.chat(args.tokenizer, [prompt])
+    print('log:', response)
+    return response
 
 
 if __name__ == '__main__':