Update app.py

app.py

@@ -3,17 +3,16 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Global variables to store the loaded model and tokenizer.
-model = None
-tokenizer = None
+# Global dictionary to store loaded models, keyed by model name.
+loaded_models = {}
 
 @spaces.GPU
 def load_model(model_name: str):
     """
-    Loads the model and tokenizer given the model name.
-    Returns a status message.
+    Loads the model and tokenizer and stores them in a global dictionary.
+    Returns a status message and the name of the loaded model.
     """
-    global model, tokenizer
+    global loaded_models
     try:
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
@@ -21,20 +20,24 @@ def load_model(model_name: str):
             device_map="auto"
         )
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        return f"Model '{model_name}' loaded successfully."
+        loaded_models[model_name] = (model, tokenizer)
+        return f"Model '{model_name}' loaded successfully.", model_name
     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):
+def generate(prompt, history, current_model_name):
     """
     Generates a response for the given prompt using the loaded model.
-    If the model is not loaded, informs the user to load it first.
+    If the model (based on the current model name) isn’t loaded, it informs the user.
     """
-    if model is None or tokenizer is None:
+    global loaded_models
+    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."
-
-    # Prepare the chat history (here, a simple system prompt is added)
+    
+    model, tokenizer = loaded_models[current_model_name]
+    
+    # Prepare the messages (with a system prompt)
     messages = [
         {"role": "system", "content": "Je bent een vriendelijke, behulpzame assistent."},
         {"role": "user", "content": prompt}
@@ -44,6 +47,7 @@ def generate(prompt, history):
         tokenize=False,
         add_generation_prompt=True
     )
+    
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
     
     generated_ids = model.generate(
@@ -70,13 +74,15 @@ with gr.Blocks() as demo:
         )
         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() is called.
-    load_button.click(fn=load_model, inputs=model_name_input, outputs=load_status)
+    # 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])
     
     gr.Markdown("## Chat Interface")
-    # The ChatInterface calls generate() which uses the loaded model.
-    chat_interface = gr.ChatInterface(fn=generate)
+    # The chat interface now passes the hidden model_state into the generate function.
+    chat_interface = gr.ChatInterface(fn=generate, extra_inputs=[model_state])
 
-# Launch the Gradio app (using share=True if you wish to share it publicly).
+# Launch the Gradio app (share=True to get a public link if desired).
 demo.launch(share=True)