Update app.py

app.py

@@ -1,42 +1,24 @@
 import gradio as gr
-from functools import lru_cache
-
-# Cache model loading to optimize performance
-@lru_cache(maxsize=3)
-def load_hf_model(model_name):
-    return gr.load(f"models/{model_name}", src="huggingface")
-
-# Load all models at startup
-MODELS = {
-    "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": load_hf_model("deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"),
-    "deepseek-ai/DeepSeek-R1": load_hf_model("deepseek-ai/DeepSeek-R1"),
-    "deepseek-ai/DeepSeek-R1-Zero": load_hf_model("deepseek-ai/DeepSeek-R1-Zero")
-}
+from models import demo  # Import the demo object from models.py
 
 # --- Chatbot function ---
 def chatbot(input_text, history, model_choice, system_message, max_new_tokens, temperature, top_p):
     history = history or []
     
-    # Get the selected model component
-    model_component = MODELS[model_choice]
-    
     # Create payload for the model
     payload = {
-        "inputs": input_text,  # Directly pass the input text
-        "parameters": {
-            "max_new_tokens": max_new_tokens,
-            "temperature": temperature,
-            "top_p": top_p,
-            "return_full_text": False  # Only return the generated text
-        }
+        "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:
-        response = model_component(**payload)  # Pass payload as keyword arguments
-        if isinstance(response, list) and len(response) > 0:
-            # Extract the generated text from the response
-            assistant_response = response[0].get("generated_text", "No response generated.")
+        response = demo(payload)  # Use the demo object directly
+        if isinstance(response, dict) and "choices" in response:
+            assistant_response = response["choices"][0]["message"]["content"]
         else:
             assistant_response = "Unexpected model response format."
     except Exception as e:
@@ -48,7 +30,7 @@ def chatbot(input_text, history, model_choice, system_message, max_new_tokens, t
     return history, history, ""
 
 # --- Gradio Interface ---
-with gr.Blocks(theme=gr.themes.Soft(), title="DeepSeek Chatbot") as demo:
+with gr.Blocks(theme=gr.themes.Soft(), title="DeepSeek Chatbot") as app:
     gr.Markdown(
         """
         # DeepSeek Chatbot
@@ -69,9 +51,9 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DeepSeek Chatbot") as demo:
     with gr.Row():
         with gr.Accordion("Options", open=True):
             model_choice = gr.Radio(
-                choices=list(MODELS.keys()),
+                choices=["DeepSeek-R1-Distill-Qwen-32B", "DeepSeek-R1", "DeepSeek-R1-Zero"],
                 label="Choose a Model",
-                value="deepseek-ai/DeepSeek-R1"
+                value="DeepSeek-R1"
             )
             with gr.Accordion("Optional Parameters", open=False):
                 system_message = gr.Textbox(
@@ -104,4 +86,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DeepSeek Chatbot") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    app.launch()