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
 
 def truncate_list(lst, num):
     if num not in lst:
@@ -91,8 +91,8 @@ def warmup(model):
         prompt += " "
     input_ids = model.tokenizer([prompt]).input_ids
     input_ids = torch.as_tensor(input_ids).to(model.base_model.device)
-    for output_ids in model.ea_generate(input_ids):
-        ol=output_ids.shape[1]
+    outs=model.generate(input_ids)
+    print(outs)
 @spaces.GPU(duration=30)
 def bot(history, temperature, top_p, use_EaInfer, highlight_EaInfer,session_state,):
     if not history:
@@ -269,10 +269,8 @@ parser.add_argument(
 args = parser.parse_args()
 a=torch.tensor(1).cuda()
 print(a)
-model = EaModel.from_pretrained(
-    base_model_path=args.base_model_path,
-    ea_model_path=args.ea_model_path,
-    total_token=args.total_token,
+model = LlamaForCausalLM.from_pretrained(
+    args.base_model_path,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
     load_in_4bit=args.load_in_4bit,