Update app.py

app.py

@@ -1,19 +1,19 @@
 import gradio as gr
 import spaces
-from transformers import OlmoeForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import subprocess
 import sys
 
-# Force install the specific transformers version from the GitHub PR
-subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "git+https://github.com/Muennighoff/transformers.git@olmoe"])
+# Force install the latest transformers version and flash attention
+subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "transformers", "flash-attn"])
 
 model_name = "allenai/OLMoE-1B-7B-0924"
 
 # Wrap model loading in a try-except block to handle potential errors
 try:
     DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    model = OlmoeForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16).to(DEVICE)
+    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype="auto", _attn_implementation="flash_attention_2").to(DEVICE)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 except Exception as e:
     print(f"Error loading model: {e}")
@@ -58,7 +58,7 @@ css = """
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# Nisten's Karpathy Chatbot with OSS olMoE")
     chatbot = gr.Chatbot(elem_id="output")
-    msg = gr.Textbox(label="Your message")
+    msg = gr.Textbox(label="Your prompt")
     with gr.Row():
         temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
         max_new_tokens = gr.Slider(minimum=50, maximum=4000, value=1000, step=50, label="Max New Tokens")