| |
| """Annotate prompts with DeepSeek (two passes + arbitration) into LLM-labelled gold sets. |
| |
| Scope of this entry point: text-only tasks (spatial, action, objects). Scene items |
| need frames and real captions and are produced by the caption/vision pipeline. |
| |
| Every provider call is charged to ``output/annotation/BUDGET-annotate-<name>.json`` |
| before it is sent; ``--budget`` is a hard cap for the run. |
| |
| Usage:: |
| |
| # MovieGen long-prompt pilot (decision B1: existing long prompts first) |
| uv run --no-sync python scripts/annotate_gold.py --name longprompt-pilot \\ |
| --source moviegen --buckets 50-99,100-200 --budget 300 |
| |
| # full run over a bucket mix |
| uv run --no-sync python scripts/annotate_gold.py --name gold-v1 \\ |
| --source moviegen+seed --buckets 1-14,15-49,50-99,100-200 --limit-per-task 120 --budget 4000 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from collections import Counter, defaultdict |
| import json |
| from pathlib import Path |
| import sys |
| from typing import Any, Iterable, Sequence |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / "src")) |
|
|
| from vbench_prompts_compile import annotate as A |
| from vbench_prompts_compile import records as R |
| from vbench_prompts_compile import sources as S |
| from vbench_prompts_compile.llm import ChatClient, RequestBudget, load_tokens |
|
|
|
|
| def build_moviegen_items(tasks: Sequence[str], buckets: Sequence[str], limit_per_task: int | None, seed: int) -> list[A.Item]: |
| import random |
|
|
| refs = [ref for ref in S.load_moviegen_prompts() if R.length_bucket(ref.word_count) in set(buckets)] |
| rng = random.Random(seed) |
| rng.shuffle(refs) |
| items: list[A.Item] = [] |
| for task in tasks: |
| for index, ref in enumerate(refs): |
| if limit_per_task is not None and index >= limit_per_task: |
| break |
| items.append( |
| A.Item( |
| item_id=f"{task}-moviegen-{ref.sha256[:12]}", |
| task=task, |
| prompt=ref.text, |
| group_id=f"moviegen:{ref.sha256[:12]}", |
| bucket=R.length_bucket(ref.word_count), |
| source="moviegen", |
| source_id=f"{ref.source_file}:{ref.line}", |
| meta={"word_count": ref.word_count, "licence": "CC-BY-NC-4.0"}, |
| ) |
| ) |
| return items |
|
|
|
|
| def build_seed_items(tasks: Sequence[str], buckets: Sequence[str], limit_per_task: int | None, seed: int) -> list[A.Item]: |
| """Seeds from the existing weak corpora (VG templates, K400 templates, fixtures).""" |
| import random |
|
|
| rng = random.Random(seed + 7) |
| pools: dict[str, list[A.Item]] = defaultdict(list) |
| vg = [row for row in R.read_jsonl(ROOT / "data" / "processed" / "formal-0001" / "candidates" / "spatial.jsonl")] |
| rng.shuffle(vg) |
| for row in vg[: (limit_per_task or 200)]: |
| for task in ("spatial", "objects"): |
| if task not in tasks or task not in ("spatial", "objects"): |
| continue |
| prompt = row["input"]["prompt"] |
| if R.length_bucket(R.word_count(prompt)) not in set(buckets): |
| continue |
| pools[task].append( |
| A.Item( |
| item_id=f"{task}-vg-{row['source_id']}", |
| task=task, |
| prompt=prompt, |
| group_id=row["group_id"], |
| bucket=R.length_bucket(R.word_count(prompt)), |
| source="visual_genome_prompt", |
| source_id=row["source_id"], |
| meta={"licence": "CC-BY-4.0"}, |
| ) |
| ) |
| k400 = R.read_jsonl(ROOT / "data" / "processed" / "formal-0001" / "candidates" / "action.jsonl") |
| rng.shuffle(k400) |
| for row in k400[: (limit_per_task or 200)]: |
| if "action" not in tasks: |
| continue |
| prompt = row["input"]["prompt"] |
| if R.length_bucket(R.word_count(prompt)) not in set(buckets): |
| continue |
| pools["action"].append( |
| A.Item( |
| item_id=f"action-k400-{row['source_id']}", |
| task="action", |
| prompt=prompt, |
| group_id=row["group_id"], |
| bucket=R.length_bucket(R.word_count(prompt)), |
| source="k400_prompt", |
| source_id=row["source_id"], |
| ) |
| ) |
| items: list[A.Item] = [] |
| for task in tasks: |
| items.extend(pools.get(task, [])[: (limit_per_task or len(pools.get(task, [])))]) |
| return items |
|
|
|
|
| def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--name", required=True) |
| parser.add_argument("--tasks", default="spatial,action,objects") |
| parser.add_argument("--source", default="moviegen", help="moviegen | seed | moviegen+seed") |
| parser.add_argument("--buckets", default="1-14,15-49,50-99,100-200") |
| parser.add_argument("--limit-per-task", type=int, default=None) |
| parser.add_argument("--budget", type=int, default=200, help="hard cap of provider requests for this run") |
| parser.add_argument("--provider", default="deepseek") |
| parser.add_argument("--model", default=None) |
| parser.add_argument("--token-file", default=str(ROOT / "token.txt")) |
| parser.add_argument("--seed", type=int, default=20260919) |
| parser.add_argument("--max-tokens", type=int, default=384) |
| parser.add_argument("--dry-run", action="store_true") |
| return parser.parse_args(argv) |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| args = parse_args(argv) |
| tasks = [t.strip() for t in args.tasks.split(",") if t.strip()] |
| buckets = [b.strip() for b in args.buckets.split(",") if b.strip()] |
| items: list[A.Item] = [] |
| if "moviegen" in args.source: |
| items += build_moviegen_items(tasks, buckets, args.limit_per_task, args.seed) |
| if "seed" in args.source: |
| items += build_seed_items(tasks, buckets, args.limit_per_task, args.seed) |
| if not items: |
| print("no items selected", file=sys.stderr) |
| return 2 |
|
|
| out_dir = ROOT / "data" / "gold" / args.name |
| out_dir.mkdir(parents=True, exist_ok=True) |
| R.write_jsonl(out_dir / "items.jsonl", [item.as_dict() for item in items]) |
| print(json.dumps({"items": len(items), "by_task": dict(Counter(i.task for i in items)), "by_bucket": dict(Counter(i.bucket for i in items))}, ensure_ascii=False)) |
| if args.dry_run: |
| return 0 |
|
|
| tokens = load_tokens(args.token_file) |
| client = ChatClient(provider=args.provider, api_key=tokens.get(args.provider), model=args.model) |
| budget = RequestBudget(ROOT / "output" / "annotation", f"annotate-{args.name}", authorized_total=args.budget) |
| vocab = S.load_k400() |
| resolver = vocab.resolve |
|
|
| gold: list[dict[str, Any]] = [] |
| quarantine: list[dict[str, Any]] = [] |
| reports: list[dict[str, Any]] = [] |
| for item in items: |
| if budget.remaining() <= 0: |
| print("budget exhausted; stopping", file=sys.stderr) |
| break |
| result = A.annotate_item( |
| client, |
| item, |
| action_resolver=resolver, |
| max_tokens=args.max_tokens, |
| charge=lambda entry: budget.charge(entry, provider=args.provider), |
| ) |
| entry = { |
| "item_id": item.item_id, |
| "task": item.task, |
| "bucket": item.bucket, |
| "status": result.status, |
| "agreement": result.agreement, |
| "errors": result.errors, |
| "target": result.target, |
| "usage": result.usage, |
| } |
| reports.append(entry) |
| if result.target is not None: |
| gold.append(R.add_length_meta(A.result_to_record(result, item, action_resolver=resolver))) |
| else: |
| quarantine.append({**entry, "prompt": item.prompt, "passes": result.passes}) |
| budget.append({"kind": "result", "item_id": item.item_id, "status": result.status, "agreement": result.agreement, "errors": result.errors}) |
|
|
| R.write_jsonl(out_dir / "gold.jsonl", gold) |
| R.write_jsonl(out_dir / "quarantine.jsonl", quarantine) |
| status_counts = Counter(entry["status"] for entry in reports) |
| report = { |
| "name": args.name, |
| "source": args.source, |
| "buckets": buckets, |
| "tasks": tasks, |
| "items": len(items), |
| "processed": len(reports), |
| "status_counts": dict(status_counts), |
| "agreement_rate": (sum(1 for entry in reports if entry["agreement"]) / len(reports)) if reports else 0.0, |
| "by_task": {task: dict(Counter(entry["status"] for entry in reports if entry["task"] == task)) for task in tasks}, |
| "by_bucket": {bucket: dict(Counter(entry["status"] for entry in reports if entry["bucket"] == bucket)) for bucket in buckets}, |
| "error_counts": dict(Counter(error for entry in reports for error in entry["errors"])), |
| "target_kinds": { |
| "action_other": sum(1 for record in gold if record["task"] == "action" and R.ACTION_OTHER in record["target"]["actions"]), |
| "empty_results": sum(1 for record in gold if isinstance(record["target"], dict) and any(isinstance(v, list) and not v for v in record["target"].values())), |
| }, |
| "budget": budget.summary(), |
| "quality": A.ANNOTATION_QUALITY, |
| "note": "LLM labels (two passes + arbitration), not human gold", |
| } |
| (out_dir / "report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") |
| print(json.dumps(report, ensure_ascii=False, indent=2)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
