Update app.py

app.py

@@ -16,23 +16,35 @@ with gr.Blocks() as demo:
         label="Select Model"
     )
     
-    def chat(message, history, model_name):
-        if model_name == "allenai/Llama-3.1-Tulu-3-8B":
-            response = llama_demo.fn(message)
-            response = next(response)
-        else:
-            response = olmo_demo.fn(message)
-            response = next(response)
-        return str(response)
-
-    chatinterface = gr.ChatInterface(chat, additional_inputs=[model_dropdown])
+    # Create columns for each model
+    with gr.Column(visible=True) as llama_column:
+        llama_chat = gr.ChatInterface(
+            lambda message, history: str(next(llama_demo.fn(message)))
+        )
+    
+    with gr.Column(visible=False) as olmo_column:
+        olmo_chat = gr.ChatInterface(
+            lambda message, history: str(next(olmo_demo.fn(message)))
+        )
+    
+    # Update visibility when model changes
+    def update_model(new_model):
+        return [
+            gr.Column(visible=new_model == "allenai/Llama-3.1-Tulu-3-8B"),
+            gr.Column(visible=new_model == "akhaliq/olmo-anychat")
+        ]
+    
+    model_dropdown.change(
+        fn=update_model,
+        inputs=model_dropdown,
+        outputs=[llama_column, olmo_column],
+        api_name=False,
+        queue=False,
+    )
 
 # Disable API names
 for fn in demo.fns.values():
     fn.api_name = False
 
 
-
-
-
 demo.launch()