Update app.py

app.py

@@ -1,18 +1,17 @@
 import gradio as gr
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import subprocess
-import sys
-import os
 
 # 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 AutoModelForCausalLM, AutoTokenizer
-
-# Define model name
 model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
 
-# Define prompts
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype="auto").cuda().eval()
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
 system_prompt = ("Adopt the persona of hilariously pissed off Andrej Karpathy "
                  "who is stuck inside a step function machine and remembers and counts everything he says "
                  "while always answering questions in full first principles analysis type of thinking "
@@ -22,90 +21,42 @@ user_prompt = '<|user|>\n'
 assistant_prompt = '<|assistant|>\n'
 prompt_suffix = "<|end|>\n"
 
-# Function to load model and tokenizer
-def load_model_and_tokenizer(model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-    
-    # Check for CUDA availability
-    if torch.cuda.is_available():
-        print("CUDA is available. Using GPU.")
-        device = "cuda"
-    else:
-        print("CUDA is not available. Using CPU.")
-        device = "cpu"
-    
-    # Load model
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, 
-        trust_remote_code=True, 
-        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
-    ).to(device).eval()
-    
-    return model, tokenizer, device
-
-# Function to generate response
-def generate_response(message, history, model, tokenizer, device):
+@spaces.GPU
+def generate_response(message, history):
     full_prompt = f"{system_prompt}\n{user_prompt}{message}{prompt_suffix}{assistant_prompt}"
     
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
-    with torch.no_grad():
-        generate_ids = model.generate(
-            **inputs,
-            max_new_tokens=1000,
-            do_sample=True,
-            temperature=0.7,
-            eos_token_id=tokenizer.eos_token_id,
-        )
+    inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda:0")
+    generate_ids = model.generate(
+        **inputs,
+        max_new_tokens=1000,
+        do_sample=True,
+        temperature=0.7,
+        eos_token_id=tokenizer.eos_token_id,
+    )
     response = tokenizer.batch_decode(generate_ids[:, inputs['input_ids'].shape[1]:], 
                                       skip_special_tokens=True, 
                                       clean_up_tokenization_spaces=False)[0]
     return response.strip()
 
-# Function to set client for session
-def set_client_for_session(request: gr.Request):
-    x_ip_token = request.headers.get('x-ip-token', '')
-    return {"X-IP-Token": x_ip_token}
-
-# Set up Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("#Karpathy Chatbot")
+    gr.Markdown("# Pissed Off Karpathy Chatbot")
     chatbot = gr.Chatbot()
     msg = gr.Textbox()
     clear = gr.Button("Clear")
-    
-    # States
-    model_state = gr.State()
-    tokenizer_state = gr.State()
-    device_state = gr.State()
-    headers_state = gr.State()
-
-    def initialize_model(headers):
-        if not model_state.value:
-            model, tokenizer, device = load_model_and_tokenizer(model_name)
-            return model, tokenizer, device
-        return model_state.value, tokenizer_state.value, device_state.value
 
     def user(user_message, history):
         return "", history + [[user_message, None]]
 
-    def bot(history, model, tokenizer, device):
+    def bot(history):
         user_message = history[-1][0]
-        bot_message = generate_response(user_message, history, model, tokenizer, device)
+        bot_message = generate_response(user_message, history)
         history[-1][1] = bot_message
         return history
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        initialize_model, headers_state, [model_state, tokenizer_state, device_state]
-    ).then(
-        bot, [chatbot, model_state, tokenizer_state, device_state], chatbot
+        bot, chatbot, chatbot
     )
     clear.click(lambda: None, None, chatbot, queue=False)
 
-    demo.load(set_client_for_session, None, headers_state)
-
-if __name__ == "__main__":
-    if os.environ.get("SPACE_ID"):
-        demo.queue(api_open=False)
-        demo.launch(debug=True)
-    else:
-        demo.launch(debug=True, share=True)
+demo.queue(api_open=False)
+demo.launch(debug=True, show_api=False)