scripts/bench.py

"""Benchmark the zero-shot classifier against a small labelled set.

Run from the project root:

    python scripts/bench.py                       # current MODEL_NAME
    python scripts/bench.py MoritzLaurer/deberta-v3-xsmall-zeroshot-v1.1-all-33

Prints accuracy for a sweep of LEXICAL_WEIGHT values and hypothesis templates,
so config changes are made on evidence rather than vibes. The cases are the
English text as it would arrive *after* translation, since that is what the
model actually sees.

The weight sweep runs the model ONCE per case and re-blends the cached scores:
LEXICAL_WEIGHT only affects the blend, never the model, so re-running inference
per weight would just burn time for identical numbers.
"""
from __future__ import annotations

import os
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

from app import classifier                      # noqa: E402
from app.config import DEFAULT_LABELS           # noqa: E402

LABELS = [{"name": n, "description": d} for n, d in DEFAULT_LABELS]

CASES: list[tuple[str, str]] = [
    ("There is no water in our village for three days", "Water Supply"),
    ("The water coming from the tap is muddy and smells bad", "Water Supply"),
    ("A main pipe has burst and water is flooding the street", "Water Supply"),
    ("Our area has been without electricity for three days", "Electricity"),
    ("The transformer near our house is damaged and sparking", "Electricity"),
    ("Voltage keeps fluctuating and our appliances are getting damaged", "Electricity"),
    ("There is a huge pothole on the main road", "Road & Transport"),
    ("The traffic signal at the junction is not working", "Road & Transport"),
    ("The pavement is broken and people cannot walk safely", "Road & Transport"),
    ("Garbage has not been collected in our locality for a week", "Waste & Sanitation"),
    ("The drain is blocked and sewage is overflowing on the street", "Waste & Sanitation"),
    ("The public bin is overflowing and smells terrible", "Waste & Sanitation"),
    ("Mosquitoes are breeding in the stagnant water near the school", "Public Health"),
    ("Stray dogs are attacking children in our neighbourhood", "Public Health"),
    ("Construction noise starts at 5am every day and nobody can sleep", "Noise & Pollution"),
    ("A factory is releasing thick black smoke into the air", "Noise & Pollution"),
    ("Someone is putting up an illegal building without permission", "Building & Property"),
    ("The old building next door is cracked and looks unsafe", "Building & Property"),
]

CANDIDATES = [l["name"] for l in LABELS]


def model_scores(template: str) -> tuple[list[list[tuple[str, float]]], int]:
    """Raw model scores for every case under one hypothesis template."""
    out, latencies = [], []
    for text, _ in CASES:
        started = time.perf_counter()
        r = classifier._pipe(
            text, candidate_labels=CANDIDATES,
            hypothesis_template=template, multi_label=False,
        )
        latencies.append((time.perf_counter() - started) * 1000)
        out.append(list(zip(r["labels"], [float(s) for s in r["scores"]])))
    return out, int(sum(latencies) / len(latencies))


def score(cached, weight, avg_ms, title) -> tuple[int, list[str]]:
    hits, misses = 0, []
    for (text, expected), raw in zip(CASES, cached):
        blended = classifier._blend(list(raw), text, LABELS, weight)
        blended.sort(key=lambda p: p[1], reverse=True)
        got, conf = blended[0]
        if got == expected:
            hits += 1
        else:
            misses.append(f"      {expected:<20} -> {got:<20} ({conf:.2f})  {text[:46]}")
    print(f"  {title:<44} {hits:>2}/{len(CASES)}  ({100 * hits / len(CASES):3.0f}%)  ~{avg_ms}ms")
    return hits, misses


def main() -> int:
    if len(sys.argv) > 1:
        classifier.MODEL_NAME = sys.argv[1]
        classifier._state["model"] = sys.argv[1]

    started = time.time()
    state = classifier.load_model(force=True)
    if state["status"] != "ready":
        print("MODEL FAILED TO LOAD:\n", state["error"][:500])
        return 1
    print(f"\nmodel: {classifier._state['model']}   (loaded in {time.time() - started:.0f}s)\n")

    default_template = "This complaint is about {}."
    cached, avg_ms = model_scores(default_template)

    print("LEXICAL_WEIGHT sweep  (model run once, blend re-applied)")
    best = (-1, 0.0, [])
    for weight in (0.0, 0.15, 0.25, 0.3, 0.4, 0.5, 0.6):
        hits, misses = score(cached, weight, avg_ms, f"lexical_weight = {weight}")
        if hits > best[0]:
            best = (hits, weight, misses)

    print(f"\n  best: lexical_weight={best[1]}  ->  {best[0]}/{len(CASES)}")
    if best[2]:
        print("  remaining misses:")
        print("\n".join(best[2]))

    print("\nhypothesis template sweep")
    score(cached, best[1], avg_ms, f"{default_template!r}")
    for template in (
        "This is a complaint about {}.",
        "The problem is {}.",
        "This text is about {}.",
    ):
        c, ms = model_scores(template)
        score(c, best[1], ms, f"{template!r}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())