Update app.py

app.py

@@ -1,60 +1,52 @@
 import gradio as gr
-import subprocess
+import spaces
+import transformers_gradio
+from functools import lru_cache
 
-# Function to load a model using Hugging Face Spaces
-def load_model_from_space(model_name):
-    print(f"Attempting to load {model_name}...")
-    try:
-        demo = gr.load(name=model_name, src="spaces")
-        print(f"Successfully loaded {model_name}")
-        return demo
-    except Exception as e:
-        print(f"Error loading model {model_name}: {e}")
-        return None
+# Cache model loading to optimize performance
+@lru_cache(maxsize=3)
+def load_hf_model(model_name):
+    return gr.load(
+        name=f"deepseek-ai/{model_name}",
+        src=transformers_gradio.registry,
+        api_name="/chat"
+    )
 
-# Load the models
-deepseek_r1_distill = load_model_from_space("deepseek-ai/DeepSeek-R1-Distill-Qwen-32B")
-deepseek_r1 = load_model_from_space("deepseek-ai/DeepSeek-R1")
-deepseek_r1_zero = load_model_from_space("deepseek-ai/DeepSeek-R1-Zero")
+# Load all models at startup
+MODELS = {
+    "DeepSeek-R1-Distill-Qwen-32B": load_hf_model("DeepSeek-R1-Distill-Qwen-32B"),
+    "DeepSeek-R1": load_hf_model("DeepSeek-R1"),
+    "DeepSeek-R1-Zero": load_hf_model("DeepSeek-R1-Zero")
+}
 
 # --- Chatbot function ---
 def chatbot(input_text, history, model_choice, system_message, max_new_tokens, temperature, top_p):
     history = history or []
-    print(f"Input: {input_text}, History: {history}, Model: {model_choice}")
-
-    # Choose the model based on user selection
-    if model_choice == "DeepSeek-R1-Distill-Qwen-32B" and deepseek_r1_distill:
-        model_demo = deepseek_r1_distill
-    elif model_choice == "DeepSeek-R1" and deepseek_r1:
-        model_demo = deepseek_r1
-    elif model_choice == "DeepSeek-R1-Zero" and deepseek_r1_zero:
-        model_demo = deepseek_r1_zero
-    else:
-        default_response = "Model not selected or could not be loaded."
-        history.append((input_text, default_response))
-        return history, history, "", model_choice, system_message, max_new_tokens, temperature, top_p
-
-    # Call the model's 'predict' function.
+    
+    # Get the selected model component
+    model_component = MODELS[model_choice]
+    
+    # Create payload for the model
+    payload = {
+        "messages": [{"role": "user", "content": input_text}],
+        "system": system_message,
+        "max_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_p": top_p
+    }
+    
+    # Run inference using the selected model
     try:
-      model_output = model_demo(input_text, history, max_new_tokens, temperature, top_p, system_message, fn_index=0)
+        response = model_component(payload)
+        assistant_response = response[-1]["content"]
     except Exception as e:
-      print(f"An error occurred: {e}")
-      model_output= "An error occurred please check the model and try again."
-      history.append((input_text, model_output))
-      return history, history, "", model_choice, system_message, max_new_tokens, temperature, top_p
-
-    # Check if model_output is iterable and has expected number of elements
-    if not isinstance(model_output, (list, tuple)) or len(model_output) < 2:
-        error_message = "Model output does not have the expected format."
-        history.append((input_text, error_message))
-        return history, history, "", model_choice, system_message, max_new_tokens, temperature, top_p
-
-    response = model_output[-1][1] if model_output[-1][1] else "Model did not return a response."
-    history.append((input_text, response))
-    return history, history, "", model_choice, system_message, max_new_tokens, temperature, top_p
+        assistant_response = f"Error: {str(e)}"
+    
+    history.append((input_text, assistant_response))
+    return history, history, ""
 
 # --- Gradio Interface ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+with gr.Blocks(theme=gr.themes.Soft(), title="DeepSeek Chatbot") as demo:
     gr.Markdown(
         """
         # DeepSeek Chatbot
@@ -73,11 +65,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 submit_btn = gr.Button("Submit", variant="primary")
                 clear_btn = gr.ClearButton([msg, chatbot_output])
 
-    # Options moved below the chat interface
     with gr.Row():
         with gr.Accordion("Options", open=True):
             model_choice = gr.Radio(
-                choices=["DeepSeek-R1-Distill-Qwen-32B", "DeepSeek-R1", "DeepSeek-R1-Zero"],
+                choices=list(MODELS.keys()),
                 label="Choose a Model",
                 value="DeepSeek-R1"
             )
@@ -97,21 +88,24 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                     minimum=0.10, maximum=1.00, value=0.90, label="Top-p (nucleus sampling)"
                 )
 
-    # Maintain chat history
     chat_history = gr.State([])
 
     # Event handling
     submit_btn.click(
         chatbot,
         [msg, chat_history, model_choice, system_message, max_new_tokens, temperature, top_p],
-        [chatbot_output, chat_history, msg, model_choice, system_message, max_new_tokens, temperature, top_p],
+        [chatbot_output, chat_history, msg]
     )
     msg.submit(
         chatbot,
         [msg, chat_history, model_choice, system_message, max_new_tokens, temperature, top_p],
-        [chatbot_output, chat_history, msg, model_choice, system_message, max_new_tokens, temperature, top_p],
+        [chatbot_output, chat_history, msg]
     )
 
-# Launch the demo
+# Add GPU support for Hugging Face Spaces
+demo.fn = spaces.GPU()(demo.fn)
+for fn in demo.fns.values(): 
+    fn.api_name = False
+
 if __name__ == "__main__":
     demo.launch()