Spaces:
Running on Zero
Running on Zero
| """Grounding demo — is a claim supported by the document it cites? | |
| Runs BOTH grounding models side by side: | |
| * grounding-en (open weights, English) | |
| * grounding-multilingual (commercial; loaded from a PRIVATE repo via the HF_TOKEN Space secret, | |
| so its outputs are shown but the weights are never downloadable here) | |
| Each score is a calibrated support probability: p = softmax(logits / T)[entailment_index]. Temperature T | |
| is per-model (monotonic — it fixes the confidence value, not the ranking). | |
| """ | |
| import os | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| EI = 0 # entailment class index for the *-zeroshot-v2.0 heads (id2label {0: entailment, 1: not_entailment}) | |
| MODELS = [ | |
| {"key": "en", "name": "grounding-en", "tag": "open · English", | |
| "repo": os.environ.get("EN_REPO", "nutrientdocs/grounding-en"), "T": 1.2866, "private": False}, | |
| {"key": "multi", "name": "grounding-multilingual", "tag": "commercial · 15+ languages", | |
| "repo": os.environ.get("ML_REPO", "nutrientdocs/grounding-multilingual-private"), | |
| "subfolder": "weights", "T": 0.9432, "private": True}, | |
| ] | |
| TOKEN = os.environ.get("HF_TOKEN") # Space secret; required for the private commercial model | |
| _loaded = {} | |
| # Model selector -> which model keys to score. A model card deep-links here with ?model=en|multi so | |
| # clicking "Try it" from a model page opens the demo focused on that model. | |
| SELECTION = { | |
| "grounding-en (English)": ["en"], | |
| "grounding-multilingual": ["multi"], | |
| "Both (compare)": ["en", "multi"], | |
| } | |
| def _from_query(request: gr.Request): | |
| m = (request.query_params.get("model") or "").lower() | |
| if m == "en": | |
| return "grounding-en (English)" | |
| if m in ("multi", "multilingual", "bge"): | |
| return "grounding-multilingual" | |
| return "Both (compare)" | |
| def _ensure(spec): | |
| """Load tokenizer+model onto CPU and cache. Runs OUTSIDE the GPU window (no GPU needed to download), | |
| so a ZeroGPU allocation is spent only on the forward pass. Stores the Exception if a private model | |
| can't be reached (e.g. missing HF_TOKEN secret).""" | |
| if spec["key"] in _loaded: | |
| return _loaded[spec["key"]] | |
| kw = {"token": TOKEN} if TOKEN else {} # both models are private; token is org-scoped | |
| if spec.get("subfolder"): | |
| kw["subfolder"] = spec["subfolder"] | |
| try: | |
| tok = AutoTokenizer.from_pretrained(spec["repo"], **kw) | |
| model = AutoModelForSequenceClassification.from_pretrained(spec["repo"], **kw).eval() | |
| _loaded[spec["key"]] = (tok, model) | |
| except Exception as e: | |
| _loaded[spec["key"]] = e | |
| return _loaded[spec["key"]] | |
| def _forward(premise, hypothesis, keys): | |
| """ZeroGPU-allocated: move each selected model to CUDA and score. Returns key -> calibrated support | |
| probability (None if the model failed to load).""" | |
| scores = {} | |
| for spec in MODELS: | |
| if spec["key"] not in keys: | |
| continue | |
| got = _loaded.get(spec["key"]) | |
| if not isinstance(got, tuple): | |
| scores[spec["key"]] = None | |
| continue | |
| tok, model = got | |
| model = model.to("cuda") | |
| enc = tok(premise, hypothesis, truncation=True, max_length=1024, return_tensors="pt").to("cuda") | |
| with torch.no_grad(): | |
| logits = model(**enc).logits | |
| scores[spec["key"]] = torch.softmax(logits / spec["T"], dim=-1)[0, EI].item() # calibrated | |
| return scores | |
| def grade(premise, hypothesis, selection): | |
| # Generator: emit an immediate status line so the click has visible feedback while the (potentially | |
| # slow) cold start runs — first request downloads the weights to CPU AND cold-starts a shared ZeroGPU | |
| # allocation, which together can take ~30-60s. Warm requests only wait on the GPU handoff. | |
| if not premise.strip() or not hypothesis.strip(): | |
| yield "Enter a document premise and a claim to check." | |
| return | |
| keys = SELECTION.get(selection, ["en", "multi"]) | |
| cold = [s for s in MODELS if s["key"] in keys and s["key"] not in _loaded] | |
| if cold: | |
| yield ("⏳ **Waking up…** the first check downloads the model and cold-starts a shared GPU — " | |
| "this can take **~30–60s**. Later checks are near-instant.") | |
| else: | |
| yield "⏳ Scoring on the GPU…" | |
| for spec in MODELS: # download/load on CPU first, so the GPU window is inference-only | |
| if spec["key"] in keys: | |
| _ensure(spec) | |
| scores = _forward(premise, hypothesis, keys) | |
| lines = ["| Model | | Grounded support | Verdict |", "|---|---|---:|---|"] | |
| for spec in MODELS: | |
| if spec["key"] not in keys: | |
| continue | |
| p = scores.get(spec["key"]) | |
| if p is None: | |
| lines.append(f"| **{spec['name']}** | {spec['tag']} | _unavailable_ | " | |
| "_(commercial — set HF_TOKEN)_ |") | |
| continue | |
| verdict = "✅ grounded" if p >= 0.5 else "❌ not grounded" | |
| bar = "█" * round(p * 10) + "░" * (10 - round(p * 10)) | |
| lines.append(f"| **{spec['name']}** | {spec['tag']} | `{bar}` **{p:.3f}** | {verdict} |") | |
| yield "\n".join(lines) | |
| from examples import EXAMPLES # noqa: E402 (real, labeled rows; single source of truth, verified) | |
| with gr.Blocks(title="Grounding demo", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown( | |
| "# Does the document actually support this claim?\n" | |
| "Paste a **document premise** (a table or passage) and a **claim**. The grounding model scores " | |
| "whether the document *supports* the claim. Scores are temperature-calibrated support " | |
| "probabilities; the verdict is the model's decision at 0.5.\n\n" | |
| "Pick an example below (its **expected answer** is shown) or type your own. The non-English " | |
| "examples are cases the English model gets wrong — switch to the **multilingual** model.\n\n" | |
| "→ [grounding-en (open)](https://huggingface.co/nutrientdocs/grounding-en) · " | |
| "[grounding-multilingual (commercial)](https://huggingface.co/nutrientdocs/grounding-multilingual) · " | |
| "[leaderboard](https://huggingface.co/spaces/nutrientdocs/grounding-leaderboard) · " | |
| "[benchmark](https://huggingface.co/datasets/nutrientdocs/grounding-benchmark)") | |
| model_sel = gr.Radio(list(SELECTION), value="Both (compare)", label="Model") | |
| with gr.Row(): | |
| premise = gr.Textbox(label="Document premise", lines=8) | |
| hypothesis = gr.Textbox(label="Claim to check", lines=8) | |
| expected = gr.Textbox(label="Expected answer (ground truth for the selected example)", | |
| interactive=False) | |
| btn = gr.Button("Check grounding", variant="primary") | |
| gr.Markdown("<sub>⏱️ The first check cold-starts a shared GPU (~30–60s); later checks are fast.</sub>") | |
| out = gr.Markdown() | |
| btn.click(grade, [premise, hypothesis, model_sel], out, show_progress="full") | |
| gr.Examples(EXAMPLES, [premise, hypothesis, expected]) # fills the boxes incl. expected answer | |
| gr.Markdown( | |
| '## About the author\n' | |
| '<a href="https://nutrient.io/">' | |
| '<img src="https://avatars2.githubusercontent.com/u/1527679?v=3&s=200" height="80" /></a>\n\n' | |
| "This demo is maintained and funded by [Nutrient](https://nutrient.io/) — " | |
| "The deterministic document infrastructure enterprises run their highest-stakes workflows on: " | |
| "replayable output, clear exceptions, and full audit trails on the messy, regulated documents where AI alone breaks.") | |
| demo.load(_from_query, None, model_sel) # ?model=en|multi presets the selector | |
| if __name__ == "__main__": | |
| demo.launch() | |