app.py

import os
import json
import re
from pathlib import Path
from typing import Dict, List, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from transformers import AutoModel, AutoTokenizer
from huggingface_hub import hf_hub_download

# ---------------------------------------------------------------------------
# Model definition (mirrors train_unified_multihead.py)
# ---------------------------------------------------------------------------

MODEL_REPO = os.environ.get(
    "MODEL_REPO",
    "asansanwal/wet-iab-mdl-modernbert-unified-multihead-20260718-192232-v1"
)
HF_TOKEN = os.environ.get("HF_TOKEN", "")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_LENGTH = 256


class UnifiedMultiHeadModel(nn.Module):
    def __init__(self, encoder: nn.Module, heads: Dict[str, nn.Linear], tier_order: List[str]):
        super().__init__()
        self.encoder = encoder
        self.heads = nn.ModuleDict(heads)
        self.tier_order = tier_order

    def forward(self, input_ids, attention_mask, **_):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        # ModernBERT: use last_hidden_state[:, 0, :] (CLS)
        if hasattr(out, "last_hidden_state"):
            hidden = out.last_hidden_state[:, 0, :]
        else:
            hidden = out[0][:, 0, :]
        return {tier: self.heads[tier](hidden) for tier in self.tier_order}


# ---------------------------------------------------------------------------
# Load model (cached after first call)
# ---------------------------------------------------------------------------

_model = None
_tokenizer = None
_meta = None


def _load():
    global _model, _tokenizer, _meta

    token = HF_TOKEN or None

    # Download meta.json
    meta_path = hf_hub_download(MODEL_REPO, "meta.json", token=token)
    with open(meta_path) as f:
        _meta = json.load(f)

    tiers = list(_meta["tiers"].keys())
    model_name = _meta.get("model_name", "answerdotai/ModernBERT-base")

    # Download heads.pt
    heads_path = hf_hub_download(MODEL_REPO, "heads.pt", token=token)

    # Load encoder from the encoder/ subfolder inside the repo
    encoder = AutoModel.from_pretrained(
        MODEL_REPO,
        subfolder="encoder",
        token=token,
    )
    _tokenizer = AutoTokenizer.from_pretrained(
        MODEL_REPO,
        subfolder="encoder",
        token=token,
    )

    heads_state = torch.load(heads_path, map_location="cpu", weights_only=True)
    heads: Dict[str, nn.Linear] = {}
    for tier in tiers:
        num_labels = _meta["tiers"][tier]["num_labels"]
        head = nn.Linear(encoder.config.hidden_size, num_labels)
        head.weight = nn.Parameter(heads_state[f"{tier}.weight"])
        head.bias = nn.Parameter(heads_state[f"{tier}.bias"])
        heads[tier] = head

    _model = UnifiedMultiHeadModel(encoder, heads, tiers).to(DEVICE).eval()


def get_model():
    if _model is None:
        _load()
    return _model, _tokenizer, _meta


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def predict(text: str, top_k_t2: int = 5, top_k_t3: int = 5) -> Tuple[str, str, str]:
    if not text or not text.strip():
        return "Please enter some text.", "", ""

    model, tokenizer, meta = get_model()
    enc = tokenizer(
        text.strip(),
        max_length=MAX_LENGTH,
        truncation=True,
        padding=True,
        return_tensors="pt",
    ).to(DEVICE)

    with torch.no_grad():
        logits = model(**enc)

    results = {}
    for tier, lgt in logits.items():
        probs = F.softmax(lgt[0], dim=-1).cpu().tolist()
        id_to_label = {v: k for k, v in meta["tiers"][tier]["label_to_id"].items()}
        ranked = sorted(enumerate(probs), key=lambda x: -x[1])
        results[tier] = [(id_to_label[i], p) for i, p in ranked]

    def fmt_tier(tier: str, top_k: int, emoji: str) -> str:
        rows = results[tier][:top_k]
        lines = [f"### {emoji} {tier.upper()} — IAB Taxonomy Classification\n"]
        for rank, (label, score) in enumerate(rows, 1):
            bar = "█" * int(score * 20) + "░" * (20 - int(score * 20))
            lines.append(f"**{rank}. {label}**  \n`{bar}` {score*100:.1f}%\n")
        return "\n".join(lines)

    t1_md = fmt_tier("tier1", 5, "🏷️")
    t2_md = fmt_tier("tier2", top_k_t2, "📂")
    t3_md = fmt_tier("tier3", top_k_t3, "🔖")
    return t1_md, t2_md, t3_md


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

EXAMPLES = [
    ["CNN is your source for breaking news, latest news and video from politics, business, world news, health, entertainment, technology and sports."],
    ["AutoTrader is the UK's largest digital automotive marketplace for buying and selling new and used cars."],
    ["NerdWallet: Expert advice on personal finance, including banking, credit cards, mortgages, investments and loans."],
    ["Nike official online store. Free delivery and returns on eligible orders. Shop the latest range of shoes, clothing and accessories."],
    ["Coursera offers online courses, specializations, and degrees from top universities and companies."],
    ["Real Madrid CF official website. Match results, squad, fixtures and latest news about Real Madrid."],
    ["The Guardian – latest news, sport and comment from the Guardian, the world's leading liberal voice."],
    ["OpenAI is an AI research and deployment company. Our mission is to ensure that artificial general intelligence benefits all of humanity."],
]

ABOUT_MD = """
## IAB Content Taxonomy Classifier

**Model:** ModernBERT-base fine-tuned on 33.5 million multilingual web-domain examples  
**Architecture:** Unified multi-head encoder — one shared backbone, three independent classification heads  
**Coverage:** IAB Tech Lab Content Taxonomy 3.0

| Tier | Categories | Test Accuracy |
|------|-----------|---------------|
| Tier 1 (broad topic) | 27 | **98.1%** |
| Tier 2 (sub-category) | 434 | 79.8% |
| Tier 3 (specific topic) | 217 | 56.1% |

---

### Data Pipeline

The model was trained through an 8-stage pipeline:

1. **LLM Label Correction** — 98K Kaggle domains corrected via AWS Bedrock (49.7% original labels were wrong)
2. **IAB Seed XL** — LLM-generated synthetic examples for all 678 IAB taxonomy nodes
3. **Common Crawl WET/WAT** — Real domain text from CC-MAIN-2026-25 (~100K shards)
4. **Unified Hierarchical Dataset** — Combined corrected + synthetic + crawl data
5. **Argos GPU Translation** — 11-language expansion (en, zh, hi, es, fr, ar, bn, pt, ru, id, ur)
6. **Argos CPU Translation** — 15 additional languages (de, ja, sw, mr, te, tr, ta, vi, ko, it, th, fa, pl, uk, nl)
7. **All-26 Merge** — 33.5M train rows across 26 languages
8. **Sharded Fine-tuning** — 4 × 8.4M stratified shards, ModernBERT-base backbone

Full pipeline documentation: [PIPELINE.md](https://huggingface.co/datasets/asansanwal/wet-iab-ds-multilingual-unified-training-all26-v1/blob/main/PIPELINE.md)

---

### Use Cases

- **Programmatic Advertising** — Brand-safe contextual targeting aligned to IAB taxonomy
- **Content Moderation** — Automatic category flagging at ingestion
- **Search & Discovery** — Topic classification for crawled/indexed content
- **Compliance** — GARM framework alignment via IAB category mapping
- **Publisher Monetisation** — Automated inventory categorisation for DSP/SSP integrations

---

*Demo model trained on English-source data. Multilingual 26-language model in training.*  
*For API access, integration, or licensing enquiries — contact via HuggingFace.*
"""

with gr.Blocks(
    title="IAB Content Taxonomy Classifier | QuickPod AI",
    theme=gr.themes.Soft(),
    css="""
    footer { display: none !important; }
    """,
) as demo:
    gr.HTML("""
    <div style="text-align:center; padding: 24px 0 8px 0;">
      <h1 style="font-size:2rem; font-weight:700; margin:0;">
        🏷️ IAB Content Taxonomy Classifier
      </h1>
      <p style="color:#64748b; margin-top:6px; font-size:1rem;">
        Hierarchical web content classification · 27 tier-1 · 434 tier-2 · 217 tier-3 categories
      </p>
    </div>
    """)

    with gr.Tabs():
        with gr.Tab("🔍 Classify"):
            with gr.Row():
                with gr.Column(scale=1):
                    text_input = gr.Textbox(
                        label="Enter page title, description, keywords, or URL text",
                        placeholder="e.g.  'BBC Sport – live football scores, rugby, cricket, F1 and tennis news'",
                        lines=4,
                        max_lines=8,
                    )
                    with gr.Row():
                        classify_btn = gr.Button("Classify", variant="primary", scale=2)
                        clear_btn = gr.Button("Clear", scale=1)

                    top_k_t2 = gr.Slider(1, 10, value=5, step=1, label="Tier-2 results to show")
                    top_k_t3 = gr.Slider(1, 10, value=5, step=1, label="Tier-3 results to show")

                    gr.Examples(
                        examples=EXAMPLES,
                        inputs=text_input,
                        label="Example inputs",
                        examples_per_page=4,
                    )

                with gr.Column(scale=1):
                    out_t1 = gr.Markdown(label="Tier 1", elem_classes=["output-tier"])
                    out_t2 = gr.Markdown(label="Tier 2", elem_classes=["output-tier"])
                    out_t3 = gr.Markdown(label="Tier 3", elem_classes=["output-tier"])

            classify_btn.click(
                fn=predict,
                inputs=[text_input, top_k_t2, top_k_t3],
                outputs=[out_t1, out_t2, out_t3],
            )
            clear_btn.click(
                fn=lambda: ("", "", "", ""),
                outputs=[text_input, out_t1, out_t2, out_t3],
            )

        with gr.Tab("📖 About & Pipeline"):
            gr.Markdown(ABOUT_MD)

demo.launch(show_error=True)