Update app.py

app.py

@@ -4,8 +4,9 @@ import torch
 import subprocess
 import sys
 
-# Force install the specific transformers version from the GitHub PR
+# Install required packages
 subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "--force-reinstall", "--no-deps", "accelerate", "git+https://github.com/Muennighoff/transformers.git@olmoe"])
+subprocess.run('pip install flash-attn --no-build-isolation --no-deps', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 from transformers import OlmoeForCausalLM, AutoTokenizer
 
@@ -19,7 +20,8 @@ try:
         trust_remote_code=True, 
         torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
         low_cpu_mem_usage=True,
-        device_map="auto"
+        device_map="auto",
+        _attn_implementation="flash_attention_2"  # Enable Flash Attention 2
     )
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 except Exception as e:
@@ -32,29 +34,24 @@ system_prompt = ("Adopt the persona of hilariously pissed off Andrej Karpathy "
                  "while always answering questions in full first principles analysis type of thinking "
                  "without using any analogies and always showing full working code or output in his answers.")
 
-# Define a chat template as a string
 chat_template = "<|system|>{system_message}<|end|><|user|>{user_message}<|end|><|assistant|>"
 
 @spaces.GPU
 def generate_response(message, history, temperature, max_new_tokens):
     if model is None or tokenizer is None:
-        return "Model or tokenizer not loaded properly. Please check the logs."
-    
-    # Construct the full prompt
+        yield "Model or tokenizer not loaded properly. Please check the logs."
+        return
+
     full_prompt = chat_template.format(system_message=system_prompt, user_message=message)
-    
     inputs = tokenizer(full_prompt, return_tensors="pt").to(DEVICE)
     
-    with torch.no_grad():
-        generate_ids = model.generate(
-            inputs.input_ids,
-            max_new_tokens=max_new_tokens,
-            do_sample=True,
-            temperature=temperature,
-            eos_token_id=tokenizer.eos_token_id,
-        )
-    response = tokenizer.decode(generate_ids[0, inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    return response.strip()
+    streamer = tokenizer.stream(inputs.input_ids, model, temperature=temperature, max_new_tokens=max_new_tokens)
+    
+    collected_tokens = []
+    for token in streamer:
+        collected_tokens.append(token)
+        partial_text = tokenizer.decode(collected_tokens, skip_special_tokens=True)
+        yield partial_text.strip()
 
 css = """
   #output {
@@ -65,7 +62,7 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE")
+    gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE (Now with Flash Attention 2!)")
     chatbot = gr.Chatbot(elem_id="output")
     msg = gr.Textbox(label="Meow")
     with gr.Row():
@@ -78,15 +75,17 @@ with gr.Blocks(css=css) as demo:
 
     def bot(history, temp, max_tokens):
         user_message = history[-1][0]
-        bot_message = generate_response(user_message, history, temp, max_tokens)
-        history[-1][1] = bot_message
-        return history
+        bot_message = ""
+        for token in generate_response(user_message, history, temp, max_tokens):
+            bot_message = token
+            history[-1][1] = bot_message
+            yield history
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         bot, [chatbot, temperature, max_new_tokens], chatbot
     )
-    clear.click(lambda: None, None, chatbot, queue=False)
+    clear.click(lambda: None, None, chatbot, queue=True)
 
 if __name__ == "__main__":
     demo.queue(api_open=True)
-    demo.launch(debug=True, show_api=True, share=True)
+    demo.launch(debug=True, show_api=True)