""" O*NET Task -> AI Capability Classifier A fine-tuned DistilBERT model that maps work tasks to a 9-category AI capability taxonomy. Three modes: 1. Classify - type any task, get the model's predicted capability + confidence across all 9 2. Authored vs Model - for tasks already in the corpus, compare the human-authored label to the model 3. Browse - search/filter the full 18,796-task corpus """ import gradio as gr import pandas as pd import torch from transformers import AutoTokenizer, AutoModelForSequenceClassification import torch.nn.functional as F # ---------------------------------------------------------------------------- # Load model (from the Hub) and corpus (shipped with the Space) # ---------------------------------------------------------------------------- MODEL_ID = "abandekar-dev/onet-capability-classifier" tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID) model.eval() id2label = model.config.id2label LABELS = [id2label[i] for i in range(len(id2label))] # Short gloss for each capability (shown under predictions for context) GLOSS = { "INPUT": "Enter/update data into systems", "EXTRACT": "Pull structured data from unstructured sources", "CLASSIFY": "Categorize inputs into predefined groups", "MATCH": "Find correspondences across datasets", "DETECT": "Identify anomalies/exceptions from expected patterns", "GENERATE": "Create new content from context", "ORCHESTRATE": "Chain multi-step workflows with conditional logic", "PREDICT": "Forecast outcomes from historical patterns", "CONVERSE": "Natural language interaction for resolution", } corpus = pd.read_csv("corpus.csv") corpus["text_lower"] = corpus["text"].str.lower() FUNCTIONS = ["All"] + sorted(corpus["function"].dropna().unique().tolist()) # ---------------------------------------------------------------------------- # Inference # ---------------------------------------------------------------------------- def classify(text): """Return dict of label -> probability for a single task description.""" if not text or not text.strip(): return {} enc = tokenizer(text, truncation=True, max_length=128, return_tensors="pt") enc.pop("token_type_ids", None) # DistilBERT does not accept this argument with torch.no_grad(): logits = model(**enc).logits probs = F.softmax(logits, dim=-1)[0].tolist() return {LABELS[i]: probs[i] for i in range(len(LABELS))} def predict_mode(text): scores = classify(text) if not scores: return "Enter a task description above.", {} top = max(scores, key=scores.get) summary = f"### {top}\n**{GLOSS[top]}**\n\nConfidence: {scores[top]*100:.1f}%" return summary, scores def lookup_mode(text): """Find an exact/near corpus match; compare authored label to model prediction.""" if not text or not text.strip(): return "Enter or select a task.", {}, "" q = text.strip().lower() hit = corpus[corpus["text_lower"] == q] if hit.empty: hit = corpus[corpus["text_lower"].str.contains(q[:60], regex=False, na=False)] scores = classify(text) top = max(scores, key=scores.get) if hit.empty: note = ( "**Not found in corpus** — this is a novel task, so only the model can answer.\n\n" f"Model predicts: **{top}** ({scores[top]*100:.1f}%)" ) return note, scores, "" row = hit.iloc[0] authored = row["label"] agree = "match" if authored == top else "differ" icon = "✓" if authored == top else "✗" note = ( f"**Authored label:** {authored} \n" f"**Model prediction:** {top} ({scores[top]*100:.1f}%) \n\n" f"{icon} They **{agree}**." ) meta = f"Occupation: {row['occupation']} · Function: {row['function']} · SOC: {row['soc']}" return note, scores, meta def browse(query, function, limit): df = corpus if function and function != "All": df = df[df["function"] == function] if query and query.strip(): q = query.strip().lower() df = df[df["text_lower"].str.contains(q, regex=False, na=False)] df = df.head(int(limit)) return df[["text", "label", "occupation", "function"]].rename( columns={"text": "Task", "label": "Capability", "occupation": "Occupation", "function": "Function"} ) # ---------------------------------------------------------------------------- # Theme — dark, Root accent (#C4782A), zero radius, mono labels # ---------------------------------------------------------------------------- CSS = """ :root { --root:#C4782A; --root-light:#E8974A; --ink:#0E0F0C; --soil:#1C1D18; --bark:#2A2B24; --ash:#A8AB9C; --parchment:#F0EDE4; } .gradio-container { background:var(--ink) !important; color:var(--parchment) !important; font-family:'Instrument Sans',system-ui,sans-serif !important; } * { border-radius:0 !important; } h1,h2,h3 { color:var(--parchment) !important; font-family:Georgia,'DM Serif Display',serif !important; } .tab-nav button { font-family:'IBM Plex Mono',monospace !important; text-transform:uppercase; letter-spacing:0.05em; font-size:12px !important; color:var(--ash) !important; } .tab-nav button.selected { color:var(--root) !important; border-bottom:2px solid var(--root) !important; } button.primary { background:var(--root) !important; color:var(--ink) !important; font-family:'IBM Plex Mono',monospace !important; text-transform:uppercase; letter-spacing:0.05em; border:none !important; } button.primary:hover { background:var(--root-light) !important; } label span, .label-wrap span { font-family:'IBM Plex Mono',monospace !important; text-transform:uppercase; letter-spacing:0.04em; font-size:11px !important; color:var(--ash) !important; } input,textarea,.dropdown { background:var(--soil) !important; color:var(--parchment) !important; border:1px solid var(--bark) !important; } table { font-size:13px !important; } thead { background:var(--soil) !important; } """ INTRO = """ # Task → AI Capability Classifier A fine-tuned **DistilBERT** model mapping work tasks to a 9-category AI capability taxonomy, trained on 18,796 labeled O*NET tasks. Type a task to classify it, compare the model against the authored labels, or browse the corpus. """ with gr.Blocks(css=CSS, title="Task → Capability Classifier") as demo: gr.Markdown(INTRO) with gr.Tab("Classify"): gr.Markdown("Enter any task description. The model returns its predicted capability and confidence across all nine.") inp = gr.Textbox(label="Task description", lines=3, placeholder="e.g. Reconcile vendor invoices against purchase orders and flag discrepancies") btn = gr.Button("Classify", variant="primary") out_md = gr.Markdown() out_lbl = gr.Label(num_top_classes=9, label="All capabilities") btn.click(predict_mode, inp, [out_md, out_lbl]) with gr.Tab("Authored vs Model"): gr.Markdown("Paste a task that exists in the corpus to see the human-authored label beside the model's prediction. Novel tasks fall back to the model alone.") inp2 = gr.Textbox(label="Task description", lines=3) btn2 = gr.Button("Compare", variant="primary") out_md2 = gr.Markdown() out_meta = gr.Markdown() out_lbl2 = gr.Label(num_top_classes=9, label="Model scores") btn2.click(lookup_mode, inp2, [out_md2, out_lbl2, out_meta]) with gr.Tab("Browse corpus"): gr.Markdown("Search and filter all 18,796 authored task→capability mappings.") with gr.Row(): q = gr.Textbox(label="Search task text", scale=3) fn = gr.Dropdown(FUNCTIONS, value="All", label="Function", scale=1) lim = gr.Slider(10, 200, value=50, step=10, label="Max rows", scale=1) tbl = gr.Dataframe(headers=["Task", "Capability", "Occupation", "Function"], wrap=True) for c in (q, fn, lim): c.change(browse, [q, fn, lim], tbl) demo.load(browse, [q, fn, lim], tbl) if __name__ == "__main__": demo.launch()