File size: 2,838 Bytes
7245f41
9f59898
 
 
7245f41
0c6d118
9f59898
7245f41
0c6d118
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9f59898
0c6d118
 
 
 
 
 
9f59898
 
0c6d118
7245f41
0c6d118
9f59898
 
0c6d118
 
 
7245f41
0c6d118
 
 
 
7245f41
0c6d118
 
 
 
 
 
7245f41
0c6d118
 
 
 
 
 
 
 
9f59898
0c6d118
9f59898
0c6d118
9f59898
0c6d118
 
9f59898
0c6d118
 
 
 
 
 
 
9f59898
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
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()