app.py

"""
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()