app.py

import gradio as gr
import requests
from bs4 import BeautifulSoup
from transformers import pipeline

# 🔁 Load transformer model once
task_extractor = pipeline("text2text-generation", model="google/flan-t5-small")

# 🔁 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 = (
        "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 normalize_task(task)

# 🔍 Scrape models from Hugging Face
def get_models_for_task(task):
    url = f"https://huggingface.co/models?pipeline_tag={task}"
    response = requests.get(url)
    soup = BeautifulSoup(response.text, "html.parser")
    model_blocks = soup.select("div[data-testid='model-card']")

    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

        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)}", []

# 🎨 Launch UI
with gr.Blocks() as demo:
    gr.Markdown("### 🔍 HuggingFace Model Search by Task")

    with gr.Row():
        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()