Update app.py

app.py

@@ -1,22 +1,34 @@
 import gradio as gr
+import subprocess
 
-# Placeholder for model loading (adjust as needed for your specific models)
-def load_model(model_name):
-    print(f"Loading {model_name}...")
-    # Simulate different model behaviors (replace with actual model logic)
-    if model_name == "DeepSeek-R1-Distill-Qwen-32B":
-        return lambda input_text, history: f"Distilled Model Response to: {input_text}"
-    elif model_name == "DeepSeek-R1":
-        return lambda input_text, history: f"Base Model Response to: {input_text}"
-    elif model_name == "DeepSeek-R1-Zero":
-        return lambda input_text, history: f"Zero Model Response to: {input_text}"
-    else:
-        return lambda input_text, history: f"Default Response to: {input_text}"
+# Function to load a model using Hugging Face Spaces and enable GPU
+def load_model_with_gpu(model_name):
+    print(f"Attempting to load {model_name} with GPU enabled...")
+    try:
+        # Use subprocess to run hf.space_info and get GPU setting
+        result = subprocess.run(
+            ["python", "-c", f"from huggingface_hub import space_info; print(space_info('{model_name}').hardware)"],
+            capture_output=True,
+            text=True,
+            check=True
+        )
+        hardware = result.stdout.strip()
+        print(f"Hardware for {model_name}: {hardware}")
+
+        demo = gr.load(name=model_name, src="spaces")
+
+        # Return the loaded model demo
+        print(f"Successfully loaded {model_name}")
+        return demo
+
+    except Exception as e:
+        print(f"Error loading model {model_name}: {e}")
+        return None
 
-# Load the models (placeholder functions here)
-deepseek_r1_distill = load_model("DeepSeek-R1-Distill-Qwen-32B")
-deepseek_r1 = load_model("DeepSeek-R1")
-deepseek_r1_zero = load_model("DeepSeek-R1-Zero")
+# Load the models with GPU enabled (if available)
+deepseek_r1_distill = load_model_with_gpu("deepseek-ai/DeepSeek-R1-Distill-Qwen-32B")
+deepseek_r1 = load_model_with_gpu("deepseek-ai/DeepSeek-R1")
+deepseek_r1_zero = load_model_with_gpu("deepseek-ai/DeepSeek-R1-Zero")
 
 # --- Chatbot function ---
 def chatbot(input_text, history, model_choice, system_message, max_new_tokens, temperature, top_p):
@@ -24,22 +36,29 @@ def chatbot(input_text, history, model_choice, system_message, max_new_tokens, t
     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":
-        model_function = deepseek_r1_distill
-    elif model_choice == "DeepSeek-R1":
-        model_function = deepseek_r1
-    elif model_choice == "DeepSeek-R1-Zero":
-        model_function = deepseek_r1_zero
+    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:
-        model_function = lambda x, h: "Please select a model."
+        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
+        
+    # Adjust the call to the model, remove default_value if not applicable
+    model_output = model_demo(input_text, history, max_new_tokens, temperature, top_p, system_message)
 
-    # Simulate model response with parameters
-    response = model_function(input_text, history)
-    # Format the response for display (without parameter details in the main chat)
-    display_response = f"{response}"
-
-    history.append((input_text, display_response))
-    return history, history, "", model_choice, system_message, max_new_tokens, temperature, top_p  # Clear input, keep other parameters
+    # 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
 
 # --- Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -63,7 +82,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     # Options moved below the chat interface
     with gr.Row():
-        with gr.Accordion("Options", open=True): # Changed label to "Options"
+        with gr.Accordion("Options", open=True):
             model_choice = gr.Radio(
                 choices=["DeepSeek-R1-Distill-Qwen-32B", "DeepSeek-R1", "DeepSeek-R1-Zero"],
                 label="Choose a Model",
@@ -84,7 +103,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 top_p = gr.Slider(
                     minimum=0.10, maximum=1.00, value=0.90, label="Top-p (nucleus sampling)"
                 )
-    
 
     # Maintain chat history
     chat_history = gr.State([])