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,AutoTokenizer
+
 
 def truncate_list(lst, num):
     if num not in lst:
@@ -73,7 +73,7 @@ def highlight_text(text, text_list,color="black"):
 
     return result
 
-@spaces.GPU(duration=30)
+@spaces.GPU(duration=60)
 def warmup(model):
     model.cuda()
     conv = get_conversation_template(args.model_type)
@@ -90,12 +90,13 @@ def warmup(model):
     prompt = conv.get_prompt()
     if args.model_type == "llama-2-chat":
         prompt += " "
-    input_ids = tokenizer([prompt]).input_ids
+    input_ids = model.tokenizer([prompt]).input_ids
     input_ids = torch.as_tensor(input_ids).to(model.base_model.device)
-    outs=model.generate(input_ids)
-    print(outs)
-@spaces.GPU(duration=30)
+    for output_ids in model.ea_generate(input_ids):
+        ol=output_ids.shape[1]
+@spaces.GPU(duration=60)
 def bot(history, temperature, top_p, use_EaInfer, highlight_EaInfer,session_state,):
+    model.cuda()
     if not history:
         return history, "0.00 tokens/s", "0.00", session_state
     pure_history = session_state.get("pure_history", [])
@@ -270,17 +271,17 @@ parser.add_argument(
 args = parser.parse_args()
 a=torch.tensor(1).cuda()
 print(a)
-model = LlamaForCausalLM.from_pretrained(
-    args.base_model_path,
+model = EaModel.from_pretrained(
+    base_model_path=args.base_model_path,
+    ea_model_path=args.ea_model_path,
+    total_token=args.total_token,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
     load_in_4bit=args.load_in_4bit,
     load_in_8bit=args.load_in_8bit,
     device_map="auto",
 )
-
 model.eval()
-tokenizer=AutoTokenizer.from_pretrained(args.base_model_path)
 warmup(model)
 
 custom_css = """
@@ -327,4 +328,4 @@ with gr.Blocks(css=custom_css) as demo:
     )
     stop_button.click(fn=None, inputs=None, outputs=None, cancels=[send_event,regenerate_event,enter_event])
 demo.queue()
-demo.launch(share=True)
+demo.launch()