Update app.py

app.py

@@ -4,7 +4,7 @@ import time
 #os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
 import gradio as gr
 import argparse
-from model.ex_model import ExModel
+from model.ea_model import EaModel
 import torch
 from fastchat.model import get_conversation_template
 import re
@@ -76,7 +76,7 @@ def warmup(model):
         prompt += " "
     input_ids = model.tokenizer([prompt]).input_ids
     input_ids = torch.as_tensor(input_ids).cuda()
-    for output_ids in model.ex_generate(input_ids):
+    for output_ids in model.ea_generate(input_ids):
         ol=output_ids.shape[1]
 
 def bot(history, session_state):
@@ -113,7 +113,7 @@ def bot(history, session_state):
     total_ids=0
 
 
-    for output_ids in model.ex_generate(input_ids, temperature=temperature, top_p=top_p,
+    for output_ids in model.ea_generate(input_ids, temperature=temperature, top_p=top_p,
                                         max_steps=args.max_new_token):
         totaltime+=(time.time()-start_time)
         total_ids+=1
@@ -185,7 +185,7 @@ def clear(history,session_state):
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
-    "--ex-model-path",
+    "--ea-model-path",
     type=str,
     default="lmsys/vicuna-7b-v1.3",
     help="The path to the weights. This can be a local folder or a Hugging Face repo ID.",
@@ -207,9 +207,9 @@ parser.add_argument(
 )
 args = parser.parse_args()
 
-model = ExModel.from_pretrained(
+model = EaModel.from_pretrained(
     base_model_path=args.base_model_path,
-    ex_model_path=args.ex_model_path,
+    ea_model_path=args.ea_model_path,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
     load_in_4bit=args.load_in_4bit,