Update app.py

app.py

@@ -3,50 +3,77 @@ import requests
 from bs4 import BeautifulSoup
 from transformers import pipeline
 
+# 🔁 Load transformer model once
 task_extractor = pipeline("text2text-generation", model="google/flan-t5-small")
 
-# Simulated LLM task extraction (replace with real call if local)
+# 🔁 Optional alias correction
+TASK_ALIASES = {
+    "classification": "text-classification",
+    "financial classification": "text-classification",
+    "news classification": "text-classification",
+    "qa": "question-answering",
+    "summarisation": "summarization",
+    "token": "token-classification",
+    "token classification": "token-classification",
+    "object detection": "object-detection",
+}
+
+def normalize_task(task):
+    return TASK_ALIASES.get(task.lower(), task)
+
+# 🔍 Extract task from user input
 def extract_task(user_input):
-    prompt = f"Classify the following ML task: {user_input}. Just reply with the task name."
+    prompt = (
+        "Given a user query, extract the most likely machine learning task "
+        "from the following list: text-classification, token-classification, "
+        "translation, summarization, question-answering, object-detection. "
+        f"Query: {user_input}. Only return the task name."
+    )
     result = task_extractor(prompt, max_new_tokens=10)
     task = result[0]["generated_text"].strip().lower()
-    return task
-
+    return normalize_task(task)
 
-# Scrape Hugging Face models by task
+# 🔍 Scrape models from Hugging Face
 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")
+    response = requests.get(url)
+    soup = BeautifulSoup(response.text, "html.parser")
+    model_blocks = soup.select("div[data-testid='model-card']")
 
-        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
+    models_info = []
+    for block in model_blocks[:10]:  # limit to top 10 models
+        name = block.select_one("a[data-testid='model-link']")
+        arch = block.select_one("div[class*='tag']")  # very rough heuristic
 
-        return results if results else ["No models found"]
+        models_info.append({
+            "Model Name": name.text.strip() if name else "unknown",
+            "Task": task,
+            "Architecture": arch.text.strip() if arch else "unknown"
+        })
+    return models_info
 
+# 🎛 Gradio UI
+def model_search_interface(user_input):
+    try:
+        task = extract_task(user_input)
+        models = get_models_for_task(task)
+        if not models:
+            return f"No models found for task '{task}'.", []
+        return f"Task identified: {task}", models
     except Exception as e:
-        return [f"Error: {str(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
+# 🎨 Launch UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🤖 ML Task → Hugging Face Model Finder")
+    gr.Markdown("### 🔍 HuggingFace Model Search by Task")
+
     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)
+        user_input = gr.Textbox(label="Describe the ML task you're interested in:")
+        output_msg = gr.Textbox(label="Status", interactive=False)
+
+    model_table = gr.Dataframe(headers=["Model Name", "Task", "Architecture"], label="Top Models")
+
+    btn = gr.Button("🔍 Search Models")
+    btn.click(fn=model_search_interface, inputs=user_input, outputs=[output_msg, model_table])
 
 demo.launch()