Update app.py

app.py

@@ -1,29 +1,20 @@
-import subprocess
-subprocess.run(
-    'pip install flash-attn --no-build-isolation',
-    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
-    shell=True
-)
 import os
 import time
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import gradio as gr
+from threading import Thread
 
-MODEL_LIST = ["internlm/internlm2_5-7b-chat", "internlm/internlm2_5-7b-chat-1m"]
+MODEL_LIST = ["meta-llama/Meta-Llama-3.1-8B-Instruct"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-MODEL_ID = os.environ.get("MODEL_ID", None)
-MODEL_NAME = MODEL_ID.split("/")[-1]
+MODEL = os.environ.get("MODEL_ID")
 
-TITLE = "<h1><center>internlm2.5-7b-chat</center></h1>"
+TITLE = "<h1><center>Mistral-Nemo</center></h1>"
 
-DESCRIPTION = f"""
-<h3>MODEL NOW: <a href="https://hf.co/{MODEL_ID}">{MODEL_NAME}</a></h3>
-"""
 PLACEHOLDER = """
 <center>
-<p>InternLM2.5 has open-sourced a 7 billion parameter base model<br> and a chat model tailored for practical scenarios.</p>
+<p>Hi! How can I help you today?</p>
 </center>
 """
 
@@ -40,14 +31,19 @@ h3 {
 }
 """
 
+device = "cuda" # for GPU usage or "cpu" for CPU usage
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID, 
-    torch_dtype=torch.float16,
-    attn_implementation="flash_attention_2",
-    trust_remote_code=True).cuda()
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    MODEL,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    ignore_mismatched_sizes=True)
+terminators = [
+    tokenizer.eos_token_id,
+    tokenizer.convert_tokens_to_ids("<|eot_id|>")
+]
 
-model = model.eval()
 
 @spaces.GPU()
 def stream_chat(
@@ -57,28 +53,49 @@ def stream_chat(
     max_new_tokens: int = 1024, 
     top_p: float = 1.0, 
     top_k: int = 20, 
-    penalty: float = 1.2
+    penalty: float = 1.2,
 ):
     print(f'message: {message}')
     print(f'history: {history}')
-    for resp, history in model.stream_chat(
-        tokenizer,
-        query = message,
-        history = history,
+
+    conversation = []
+    for prompt, answer in history:
+        conversation.extend([
+            {"role": "user", "content": prompt}, 
+            {"role": "assistant", "content": answer},
+        ])
+
+    conversation.append({"role": "user", "content": message})
+
+    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
+    
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    
+    generate_kwargs = dict(
+        input_ids=input_ids, 
         max_new_tokens = max_new_tokens,
         do_sample = False if temperature == 0 else True,
         top_p = top_p,
         top_k = top_k,
         temperature = temperature,
-    ):
-        yield resp
+        eos_token_id=terminators,
+        streamer=streamer,
+    )
 
+    with torch.no_grad():
+        thread = Thread(target=model.generate, kwargs=generate_kwargs)
+        thread.start()
+        
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield buffer
 
+            
 chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)
 
 with gr.Blocks(css=CSS, theme="soft") as demo:
     gr.HTML(TITLE)
-    gr.HTML(DESCRIPTION)
     gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
     gr.ChatInterface(
         fn=stream_chat,
@@ -99,7 +116,7 @@ with gr.Blocks(css=CSS, theme="soft") as demo:
                 maximum=8192,
                 step=1,
                 value=1024,
-                label="Max New Tokens",
+                label="Max new tokens",
                 render=False,
             ),
             gr.Slider(
@@ -138,4 +155,4 @@ with gr.Blocks(css=CSS, theme="soft") as demo:
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()