Update app.py

app.py

@@ -1,55 +1,52 @@
 import gradio as gr
-from model_tools import extract_task, scrape_huggingface_models
-
-# Final agent with Gradio UI
-def run_agent(user_query: str):
-    """
-    Given a user query, extracts the ML task, finds relevant models, and formats results in markdown.
-    This function is used for Gradio UI interaction.
-    """
-    try:
-        # 1. Extract the standard ML task (e.g., "text-classification")
-        task = extract_task(user_query)
-
-        # 2. Get relevant models for the task
-        models = scrape_huggingface_models(task)
-
-        if not models:
-            return f"❌ No models found for task `{task}`. Try refining your query."
+import requests
+from bs4 import BeautifulSoup
+from transformers import pipeline
 
-        # 3. Format response as a markdown table
-        response = f"### 🔍 Models for task: `{task}`\n\n"
-        response += "| Model Name | Task | Architecture |\n"
-        response += "|------------|------|---------------|\n"
+task_extractor = pipeline("text2text-generation", model="google/flan-t5-small")
 
-        for model in models:
-            name = model.get("model_name", "unknown")
-            task_name = model.get("task", "unknown")
-            arch = model.get("architecture", "unknown")
-            response += f"| [{name}](https://huggingface.co/{name}) | {task_name} | {arch} |\n"
+# Simulated LLM task extraction (replace with real call if local)
+def extract_task(user_input):
+    prompt = f"Classify the following ML task: {user_input}. Just reply with the task name."
+    result = task_extractor(prompt, max_new_tokens=10)
+    task = result[0]["generated_text"].strip().lower()
+    return task
 
-        return response
-
-    except Exception as e:
-        return f"❌ Error: {str(e)}"
 
-# Gradio interface for deployment
-def gradio_ui():
-    with gr.Blocks() as demo:
-        gr.Markdown("# Hugging Face Model Finder Agent")
-        gr.Markdown("Enter a task description, and I'll find suitable ML models for you!")
-
-        # User input for task description
-        user_input = gr.Textbox(label="Describe the ML Task", placeholder="e.g., 'I need a text summarization model'", lines=2)
-
-        # Output for model search results
-        output = gr.Markdown()
+# Scrape Hugging Face models by task
+def get_models_for_task(task):
+    url = f"https://huggingface.co/models?pipeline_tag={task}"
+    headers = {"User-Agent": "Mozilla/5.0"}
+    try:
+        res = requests.get(url, headers=headers)
+        soup = BeautifulSoup(res.text, "html.parser")
 
-        # Connect the input/output to the agent
-        user_input.submit(run_agent, inputs=user_input, outputs=output)
+        results = []
+        for a in soup.find_all("a", class_="flex items-center gap-2"):
+            model_name = a.get("href", "").strip("/").split("/")[-1]
+            if model_name:
+                results.append(model_name)
+            if len(results) >= 5:
+                break
 
-    return demo
+        return results if results else ["No models found"]
 
-# Run the Gradio interface (will run locally, and can be deployed to Spaces)
-if __name__ == "__main__":
-    gradio_ui().launch()
+    except Exception as e:
+        return [f"Error: {str(e)}"]
+
+# Combined predict function
+def predict(user_input):
+    task = extract_task(user_input)
+    models = get_models_for_task(task)
+    return f"🧠 Task: {task}\n📦 Models:\n" + "\n".join(models)
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 🤖 ML Task → Hugging Face Model Finder")
+    with gr.Row():
+        input_box = gr.Textbox(label="Describe your ML task")
+        submit_btn = gr.Button("🔍 Find Models")
+    output_box = gr.Textbox(label="Suggested Models", lines=10)
+    submit_btn.click(predict, inputs=input_box, outputs=output_box)
+
+demo.launch()