Spaces:
Running
Running
| """Zero-shot sentiment baselines on the seeded IMDB sample. | |
| Runs each zero-shot model on the same 2,000-review IMDB sample and saves accuracy | |
| plus a Wilson 95% CI, as comparison points for the fine-tuned models. | |
| Saves artifacts/zero_shot_baseline.json. | |
| """ | |
| import os | |
| import sys | |
| import json | |
| import math | |
| THIS_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| PROJECT_ROOT = os.path.dirname(os.path.dirname(THIS_DIR)) | |
| for _p in (PROJECT_ROOT, THIS_DIR): | |
| if _p not in sys.path: | |
| sys.path.insert(0, _p) | |
| import evaluate as ev | |
| # Same architecture as our fine-tuned model first, then a larger zero-shot model. | |
| MODELS = [ | |
| ("typeform/distilbert-base-uncased-mnli", "distilbert-mnli (same architecture)"), | |
| ("facebook/bart-large-mnli", "bart-large-mnli (407M)"), | |
| ] | |
| def wilson_ci(acc, n, z=1.96): | |
| if n == 0: | |
| return [0.0, 0.0] | |
| denom = 1 + z * z / n | |
| center = (acc + z * z / (2 * n)) / denom | |
| half = z * math.sqrt(acc * (1 - acc) / n + z * z / (4 * n * n)) / denom | |
| return [round(center - half, 4), round(center + half, 4)] | |
| def measure(model_id, texts, labels, device): | |
| from transformers import pipeline | |
| clf = pipeline("zero-shot-classification", model=model_id, device=device) | |
| candidate = ["positive", "negative"] | |
| capped = [t[:4000] for t in texts] | |
| preds = [] | |
| chunk = 100 | |
| for i in range(0, len(texts), chunk): | |
| out = clf(capped[i:i + chunk], candidate_labels=candidate, batch_size=16) | |
| if isinstance(out, dict): | |
| out = [out] | |
| for o in out: | |
| preds.append(1 if o["labels"][0] == "positive" else 0) | |
| print(f" {min(i + chunk, len(texts))}/{len(texts)}", flush=True) | |
| acc = sum(int(p == l) for p, l in zip(preds, labels)) / len(labels) | |
| return {"accuracy": round(acc, 4), "ci": wilson_ci(acc, len(labels)), "n": len(labels)} | |
| def main(): | |
| import torch | |
| texts, labels = ev.load_imdb_test(max_per_class=1000, seed=42) | |
| labels = [int(x) for x in labels] | |
| device = "mps" if torch.backends.mps.is_available() else "cpu" | |
| out_path = os.path.join(PROJECT_ROOT, "artifacts", "zero_shot_baseline.json") | |
| os.makedirs(os.path.dirname(out_path), exist_ok=True) | |
| results = {"sample": "imdb seeded sample, max_per_class=1000, seed=42", | |
| "device": device, "models": {}} | |
| for model_id, label in MODELS: | |
| print(f"=== {label} ({model_id}) on {len(texts)} reviews, device {device} ===", flush=True) | |
| r = measure(model_id, texts, labels, device) | |
| r["label"] = label | |
| results["models"][model_id] = r | |
| with open(out_path, "w", encoding="utf-8") as f: | |
| json.dump(results, f, indent=2) | |
| print(f"RESULT {label}: acc {r['accuracy']} CI {r['ci']} (n={r['n']})", flush=True) | |
| print(f"saved {out_path}") | |
| if __name__ == "__main__": | |
| main() | |