Update app.py

app.py

@@ -2,16 +2,15 @@ import gradio as gr
 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
 
-# Load model and tokenizer
+# Define model name
 model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
-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)
 
 # Define prompts
 system_prompt = ("Adopt the persona of hilariously pissed off Andrej Karpathy "
@@ -23,42 +22,90 @@ user_prompt = '<|user|>\n'
 assistant_prompt = '<|assistant|>\n'
 prompt_suffix = "<|end|>\n"
 
-def generate_response(message, history):
+# 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):
     full_prompt = f"{system_prompt}\n{user_prompt}{message}{prompt_suffix}{assistant_prompt}"
     
-    inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda:0")
-    generate_ids = model.generate(
-        **inputs,
-        max_new_tokens=4000,
-        do_sample=True,
-        temperature=0.7,
-        eos_token_id=tokenizer.eos_token_id,
-    )
+    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,
+        )
     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("#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):
+    def bot(history, model, tokenizer, device):
         user_message = history[-1][0]
-        bot_message = generate_response(user_message, history)
+        bot_message = generate_response(user_message, history, model, tokenizer, device)
         history[-1][1] = bot_message
         return history
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
+        initialize_model, headers_state, [model_state, tokenizer_state, device_state]
+    ).then(
+        bot, [chatbot, model_state, tokenizer_state, device_state], chatbot
     )
     clear.click(lambda: None, None, chatbot, queue=False)
 
-demo.queue()
-demo.launch(debug=True, share=True)
+    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)