Update app.py

app.py

@@ -1,20 +1,19 @@
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+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"])
 
-model_name = "allenai/OLMoE-1B-7B-0924"
+model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
 
 # Wrap model loading in a try-except block to handle potential errors
 try:
     DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    model = AutoModelForCausalLM.from_pretrained(
+    model = OlmoeForCausalLM.from_pretrained(
         model_name, 
         trust_remote_code=True, 
         torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
@@ -37,24 +36,25 @@ 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."
     
-    full_prompt = f"{system_prompt}\n\nHuman: {message}\n\nAssistant:"
+    messages = [{"role": "system", "content": system_prompt},
+                {"role": "user", "content": message}]
     
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(DEVICE)
+    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, inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    response = tokenizer.decode(generate_ids[0, inputs.shape[1]:], skip_special_tokens=True)
     return response.strip()
 
 css = """
   #output {
-    height: 500px; 
+    height: 900px; 
     overflow: auto; 
     border: 1px solid #ccc; 
   }
@@ -85,4 +85,4 @@ with gr.Blocks(css=css) as demo:
 
 if __name__ == "__main__":
     demo.queue(api_open=False)
-    demo.launch(debug=True, show_api=False)
+    demo.launch(debug=True, show_api=True, share=True)