Spaces:
Running
Running
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()
|