Update app.py

app.py

@@ -16,17 +16,27 @@ current_model_name = None
 # Load selected model
 def load_model(model_name):
     global tokenizer, model, current_model_name
+    
+    # Only load if it's a different model
+    if current_model_name == model_name:
+        return
+        
     full_model_name = f"MaxLSB/{model_name}"
+    print(f"Loading model: {full_model_name}")
     tokenizer = AutoTokenizer.from_pretrained(full_model_name, token=hf_token)
     model = AutoModelForCausalLM.from_pretrained(full_model_name, token=hf_token)
     model.eval()
     current_model_name = model_name
+    print(f"Model loaded: {current_model_name}")
 
 # Initialize default model
 load_model("LeCarnet-8M")
 
 # Streaming generation function
-def respond(message, max_tokens, temperature, top_p):
+def respond(message, max_tokens, temperature, top_p, selected_model):
+    # Ensure the correct model is loaded before generation
+    load_model(selected_model)
+    
     inputs = tokenizer(message, return_tensors="pt")
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=False, skip_special_tokens=True)
 
@@ -57,18 +67,17 @@ def user(message, chat_history):
     chat_history.append([message, None])
     return "", chat_history
 
-# Bot response handler
-def bot(chatbot, max_tokens, temperature, top_p):
+# Bot response handler - UPDATED to pass selected model
+def bot(chatbot, max_tokens, temperature, top_p, selected_model):
     message = chatbot[-1][0]
-    response_generator = respond(message, max_tokens, temperature, top_p)
+    response_generator = respond(message, max_tokens, temperature, top_p, selected_model)
     for response in response_generator:
         chatbot[-1][1] = response
         yield chatbot
 
-# Model selector handler - FIXED
+# Model selector handler
 def update_model(model_name):
     load_model(model_name)
-    # Return the model_name directly instead of using gr.Dropdown.update()
     return model_name
 
 # Clear chat handler
@@ -84,7 +93,6 @@ with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
         </div>
         """)
 
-    # Create the msg_input early, but don't render it yet
     msg_input = gr.Textbox(
         placeholder="Il était une fois un petit garçon",
         label="User Input",
@@ -118,14 +126,13 @@ with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
                 bubble_full_width=False,
                 height=500
             )
-            # Now render the msg_input inside the right column, below the chatbot
             msg_input.render()
 
     # Event Handlers
     model_selector.change(
         fn=update_model,
         inputs=[model_selector],
-        outputs=[model_selector],  # This will update the dropdown value
+        outputs=[model_selector],
     )
 
     msg_input.submit(
@@ -135,7 +142,7 @@ with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
         queue=False
     ).then(
         fn=bot, 
-        inputs=[chatbot, max_tokens, temperature, top_p], 
+        inputs=[chatbot, max_tokens, temperature, top_p, model_selector],  # Pass model_selector
         outputs=[chatbot]
     )