Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import OlmoeForCausalLM, AutoTokenizer
 import torch
 import subprocess
 import sys
@@ -8,12 +8,13 @@ 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"])
 
-model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
+model_name = "allenai/OLMoE-1B-7B-0924"
 
 # Wrap model loading in a try-except block to handle potential errors
 try:
-    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16).cuda().eval()
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    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)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
 except Exception as e:
     print(f"Error loading model: {e}")
     model = None
@@ -24,48 +25,56 @@ 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.")
 
-user_prompt = '<|user|>\n'
-assistant_prompt = '<|assistant|>\n'
-prompt_suffix = "<|end|>\n"
-
 @spaces.GPU
-def generate_response(message, history):
+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{user_prompt}{message}{prompt_suffix}{assistant_prompt}"
+    full_prompt = f"{system_prompt}\n\nHuman: {message}\n\nAssistant:"
+    
+    inputs = tokenizer(full_prompt, return_tensors="pt")
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
     
-    inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda:0")
     with torch.no_grad():
         generate_ids = model.generate(
             **inputs,
-            max_new_tokens=1000,
+            max_length=inputs['input_ids'].shape[1] + max_new_tokens,
             do_sample=True,
-            temperature=0.7,
-            eos_token_id=tokenizer.eos_token_id,
+            temperature=temperature,
         )
-    response = tokenizer.batch_decode(generate_ids[:, inputs['input_ids'].shape[1]:], 
-                                      skip_special_tokens=True, 
-                                      clean_up_tokenization_spaces=False)[0]
-    return response.strip()
+    response = tokenizer.decode(generate_ids[0], skip_special_tokens=True)
+    # Extract only the assistant's response
+    assistant_response = response.split("Assistant:")[-1].strip()
+    return assistant_response
+
+css = """
+  #output {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+  }
+"""
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Pissed Off Karpathy Chatbot")
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
+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")
+    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")
     clear = gr.Button("Clear")
 
     def user(user_message, history):
         return "", history + [[user_message, None]]
 
-    def bot(history):
+    def bot(history, temp, max_tokens):
         user_message = history[-1][0]
-        bot_message = generate_response(user_message, history)
+        bot_message = generate_response(user_message, history, temp, max_tokens)
         history[-1][1] = bot_message
         return history
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
+        bot, [chatbot, temperature, max_new_tokens], chatbot
     )
     clear.click(lambda: None, None, chatbot, queue=False)