Update app.py

app.py

@@ -5,14 +5,12 @@ import torch
 
 # Global dictionary to store loaded models, keyed by model name.
 loaded_models = {}
+# Global variable to track the currently loaded model's name.
+current_model_name = ""
 
 @spaces.GPU
 def load_model(model_name: str):
-    """
-    Loads the model and tokenizer and stores them in a global dictionary.
-    Returns a status message and the name of the loaded model.
-    """
-    global loaded_models
+    global loaded_models, current_model_name
     try:
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
@@ -21,23 +19,20 @@ def load_model(model_name: str):
         )
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         loaded_models[model_name] = (model, tokenizer)
-        return f"Model '{model_name}' loaded successfully.", model_name
+        current_model_name = model_name  # update global state
+        return f"Model '{model_name}' loaded successfully."
     except Exception as e:
-        return f"Failed to load model '{model_name}': {str(e)}", ""
+        return f"Failed to load model '{model_name}': {str(e)}"
 
 @spaces.GPU
-def generate(prompt, history, current_model_name):
-    """
-    Generates a response for the given prompt using the loaded model.
-    If the model (based on the current model name) isn’t loaded, it informs the user.
-    """
-    global loaded_models
+def generate(prompt, history):
+    global loaded_models, current_model_name
     if current_model_name == "" or current_model_name not in loaded_models:
         return "Please load a model first by entering a model name and clicking the Load Model button."
     
     model, tokenizer = loaded_models[current_model_name]
     
-    # Prepare the messages (with a system prompt)
+    # Prepare the messages (with a system prompt and the user's prompt)
     messages = [
         {"role": "system", "content": "Je bent een vriendelijke, behulpzame assistent."},
         {"role": "user", "content": prompt}
@@ -47,7 +42,6 @@ def generate(prompt, history, current_model_name):
         tokenize=False,
         add_generation_prompt=True
     )
-    
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
     
     generated_ids = model.generate(
@@ -70,19 +64,16 @@ with gr.Blocks() as demo:
         model_name_input = gr.Textbox(
             label="Model Name", 
             value="agentica-org/DeepScaleR-1.5B-Preview",
-            placeholder="Enter model name"
+            placeholder="Enter model name (e.g., agentica-org/DeepScaleR-1.5B-Preview)"
         )
         load_button = gr.Button("Load Model")
     load_status = gr.Textbox(label="Status", interactive=False)
-    # Hidden state to store the currently loaded model's name.
-    model_state = gr.State("")
     
-    # When the button is clicked, load_model() returns both a status message and the model name.
-    load_button.click(fn=load_model, inputs=model_name_input, outputs=[load_status, model_state])
+    # When the Load Model button is clicked, load_model is called.
+    load_button.click(fn=load_model, inputs=model_name_input, outputs=load_status)
     
     gr.Markdown("## Chat Interface")
-    # The chat interface now passes the hidden model_state into the generate function.
-    chat_interface = gr.ChatInterface(fn=generate, extra_inputs=[model_state])
+    # Create the chat interface without extra_inputs.
+    chat_interface = gr.ChatInterface(fn=generate)
 
-# Launch the Gradio app (share=True to get a public link if desired).
 demo.launch(share=True)