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

๐Ÿท๏ธ IAB Content Taxonomy Classifier

Hierarchical web content classification ยท 27 tier-1 ยท 434 tier-2 ยท 217 tier-3 categories

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