fix

gemmademo/_chat.py

@@ -26,7 +26,7 @@ class GradioChat:
         return LlamaCppGemmaModel(name=model_name).load_model()
 
     def _chat(self):
-        def chat_fn(history, message, selected_model):
+        def chat_fn(message, history, selected_model):
             if selected_model != self.current_model_name:
                 self.current_model_name = selected_model
                 self.model = self._load_model(selected_model)  # Reload model when changed
@@ -39,6 +39,7 @@ class GradioChat:
             chat_fn,
             textbox=gr.Textbox(placeholder="What is up?", container=False),
             additional_inputs=[
+                gr.State(self.current_model_name),  # Store selected model state
                 gr.Dropdown(choices=self.model_options, value=self.current_model_name, label="Select Gemma Model"),
                 gr.Dropdown(choices=self.task_options, value="Question Answering", label="Select Task"),
             ],