Update app.py

app.py

@@ -11,7 +11,7 @@ except:
 import torch
 from fastchat.model import get_conversation_template
 import re
-from transformers import LlamaForCausalLM
+from transformers import LlamaForCausalLM,AutoTokenizer
 
 def truncate_list(lst, num):
     if num not in lst:
@@ -89,7 +89,7 @@ def warmup(model):
     prompt = conv.get_prompt()
     if args.model_type == "llama-2-chat":
         prompt += " "
-    input_ids = model.tokenizer([prompt]).input_ids
+    input_ids = tokenizer([prompt]).input_ids
     input_ids = torch.as_tensor(input_ids).to(model.base_model.device)
     outs=model.generate(input_ids)
     print(outs)
@@ -278,6 +278,7 @@ model = LlamaForCausalLM.from_pretrained(
     device_map="auto",
 )
 model.eval()
+tokenizer=AutoTokenizer.from_pretrained(args.base_model_path)
 warmup(model)
 
 custom_css = """