Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient, list_models
 import os
 import json
 
@@ -20,7 +20,7 @@ def save_prompts():
     with open(PROMPTS_FILE, "w") as file:
         json.dump(system_prompts, file, indent=4)
 
-def chat_with_model(user_input, system_prompt):
+def chat_with_model(user_input, system_prompt, selected_model):
     """Send user input to the model and return its response."""
     messages = [
         {"role": "system", "content": system_prompt},
@@ -29,7 +29,7 @@ def chat_with_model(user_input, system_prompt):
 
     try:
         result = client.chat.completions.create(
-            model="HuggingFaceH4/zephyr-7b-beta",
+            model=selected_model,
             messages=messages,
             temperature=0.5,
             max_tokens=2048,
@@ -50,32 +50,50 @@ def get_prompt(name):
     """Retrieve a system prompt by name."""
     return system_prompts.get(name, "")
 
+def fetch_models(task):
+    """Fetch models for a specific task from Hugging Face Hub."""
+    try:
+        models = list_models(filter=f"pipeline_tags:{task}")
+        return [model.modelId for model in models]
+    except Exception as e:
+        return [f"Error fetching models: {str(e)}"]
+
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("## Hugging Face Chatbot with Gradio")
+    gr.Markdown("## Hugging Face Chatbot with Dynamic Model Selection")
 
     with gr.Row():
         with gr.Column():
+            # Task selection
+            task_selector = gr.Dropdown(
+                choices=["text-generation", "image-classification", "text-classification", "translation"],
+                label="Select Task",
+                value="text-generation"
+            )
+
+            # Model selector
+            model_selector = gr.Dropdown(choices=[], label="Select Model")
+
+            # System prompt and input
             system_prompt_name = gr.Dropdown(choices=list(system_prompts.keys()), label="Select System Prompt")
             system_prompt_content = gr.TextArea(label="System Prompt", value=get_prompt("default"), lines=4)
             save_prompt_button = gr.Button("Save System Prompt")
 
             user_input = gr.TextArea(label="Enter your prompt", placeholder="Describe the character or request a detailed description...", lines=4)
             submit_button = gr.Button("Generate")
-        
+
         with gr.Column():
             output = gr.TextArea(label="Model Response", interactive=False, lines=10)
 
-    def load_prompt(name):
-        return get_prompt(name)
+    # Update model list when task changes
+    def update_model_list(task):
+        models = fetch_models(task)
+        return gr.Dropdown.update(choices=models, value=models[0] if models else None)
 
-    system_prompt_name.change(
-        lambda name: (name, get_prompt(name)), 
-        inputs=[system_prompt_name], 
-        outputs=[system_prompt_name, system_prompt_content]
-    )
+    # Event bindings
+    task_selector.change(update_model_list, inputs=[task_selector], outputs=[model_selector])
     save_prompt_button.click(update_prompt, inputs=[system_prompt_name, system_prompt_content], outputs=[])
-    submit_button.click(chat_with_model, inputs=[user_input, system_prompt_content], outputs=[output])
+    submit_button.click(chat_with_model, inputs=[user_input, system_prompt_content, model_selector], outputs=[output])
 
 # Run the app
-demo.launch()
+demo.launch()