Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
 import spaces
-from transformers import OlmoeForCausalLM, AutoTokenizer
 import torch
 import subprocess
 import sys
-import os
 
-# Force upgrade transformers to the latest version
-subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "transformers"])
+# 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"])
+
+from transformers import OlmoeForCausalLM, AutoTokenizer
 
 model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
 
@@ -37,18 +37,20 @@ 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."
     
-    messages = [{"role": "user", "content": message}]
+    messages = [{"role": "system", "content": system_prompt},
+                {"role": "user", "content": message}]
+    
     inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
     
     with torch.no_grad():
         generate_ids = model.generate(
-            **inputs,
+            inputs,
             max_new_tokens=max_new_tokens,
             do_sample=True,
             temperature=temperature,
             eos_token_id=tokenizer.eos_token_id,
         )
-    response = tokenizer.decode(generate_ids[0], skip_special_tokens=True)
+    response = tokenizer.decode(generate_ids[0, inputs.shape[1]:], skip_special_tokens=True)
     return response.strip()
 
 css = """
@@ -84,4 +86,4 @@ with gr.Blocks(css=css) as demo:
 
 if __name__ == "__main__":
     demo.queue(api_open=False)
-    demo.launch(debug=True, show_api=False, share=True )
+    demo.launch(debug=True, show_api=False)