Update app.py

app.py

@@ -3,61 +3,47 @@ import spaces
 import transformers_gradio
 
 # Load models
-llama_demo = gr.load(name="allenai/Llama-3.1-Tulu-3-8B", src=transformers_gradio.registry)
-llama_demo.fn = spaces.GPU()(llama_demo.fn)
-
-# Modified OLMo loading - pass additional parameters to handle chat interface
-olmo_demo = gr.load(
-    name="akhaliq/olmo-anychat", 
-    src="spaces",
-    api_name=False,  # Disable API endpoint
-    title="OLMo Chat",  # These parameters will be passed to ChatInterface
-    description="A chat interface for the OLMo model"
-)
-
-# Create the interface
-with gr.Blocks() as demo:
-    model_dropdown = gr.Dropdown(
-        choices=["allenai/Llama-3.1-Tulu-3-8B", "akhaliq/olmo-anychat"],
-        value="allenai/Llama-3.1-Tulu-3-8B",
-        label="Select Model"
-    )
-    
-    # Create columns for each model
-    with gr.Column(visible=True) as llama_column:
-        gr.ChatInterface(
-            fn=lambda message, history: next(llama_demo.fn(message)),
-            title="Llama 3.1 Tulu",
-            description="A chat interface for the Llama 3.1 Tulu model"
-        )
+models = {
+    "Llama": {
+        "name": "allenai/Llama-3.1-Tulu-3-8B",
+        "src": transformers_gradio.registry
+    },
+    "OLMo": {
+        "name": "akhaliq/olmo-anychat",
+        "src": "spaces"
+    }
+}
+
+def load_model(model_choice):
+    model_info = models[model_choice]
+    demo = gr.load(name=model_info["name"], src=model_info["src"])
     
-    with gr.Column(visible=False) as olmo_column:
-        # Wrap the olmo_demo with error handling if needed
-        if hasattr(olmo_demo, 'fn'):
-            olmo_demo.fn = lambda message, history: next(olmo_demo.fn(message)) if olmo_demo.fn(message) else "I apologize, but I couldn't generate a response. Please try again."
-        olmo_demo  # Just render the loaded interface directly
+    if model_choice == "Llama":
+        demo.fn = spaces.GPU()(demo.fn)
     
-    # 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")
-        ]
+    # Disable API names
+    for fn in demo.fns.values():
+        fn.api_name = False
     
-    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()
+    return demo
+
+if __name__ == "__main__":
+    with gr.Blocks() as interface:
+        model_dropdown = gr.Dropdown(
+            choices=list(models.keys()),
+            value="Llama",
+            label="Select Model"
+        )
+        
+        demo_container = gr.Group()
+        
+        def update_demo(model_choice):
+            return load_model(model_choice)
+        
+        model_dropdown.change(
+            fn=update_demo,
+            inputs=[model_dropdown],
+            outputs=[demo_container]
+        )
+        
+    interface.launch()