1 files changed, 93 insertions, 0 deletions
diff --git a/worldalign/evaluate_prefix.py b/worldalign/evaluate_prefix.py
new file mode 100644
index 0000000..2d53a51
--- /dev/null
+++ b/worldalign/evaluate_prefix.py
@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .common import dtype_for_device, read_json, write_json
+from .evaluate import unigram_f1
+from .io import load_feature_pair, select_rows
+from .models import load_prefix
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--manifest", default="artifacts/manifest.json")
+    p.add_argument("--vision", default="artifacts/vision.pt")
+    p.add_argument("--text", default="artifacts/text.pt")
+    p.add_argument("--prefix", default="artifacts/prefix.pt")
+    p.add_argument("--split", choices=["val", "test"], default="val")
+    p.add_argument("--device", default="cuda:1")
+    p.add_argument("--samples", type=int, default=100)
+    p.add_argument("--max-new-tokens", type=int, default=32)
+    p.add_argument("--output", default="artifacts/prefix_evaluation.json")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    manifest = read_json(args.manifest)
+    _, text, _, tlookup = load_feature_pair(args.vision, args.text)
+    rows = manifest[args.split][: args.samples]
+    semantic = select_rows(text["features"], tlookup, rows)
+    caption_lookup = {
+        int(row): caps for row, caps in zip(text["rows"], text["all_captions"])
+    }
+    prefix, prefix_state = load_prefix(args.prefix, args.device)
+    tokenizer = AutoTokenizer.from_pretrained(text["model"])
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    dtype = dtype_for_device(args.device)
+    lm = AutoModelForCausalLM.from_pretrained(
+        text["model"], torch_dtype=dtype
+    ).to(args.device)
+    lm.eval()
+
+    generated = []
+    with torch.inference_mode():
+        for chunk in semantic.split(16):
+            embeds = prefix(chunk.to(args.device)).to(dtype)
+            attention = torch.ones(
+                embeds.shape[:2], dtype=torch.long, device=args.device
+            )
+            output = lm.generate(
+                inputs_embeds=embeds,
+                attention_mask=attention,
+                max_new_tokens=args.max_new_tokens,
+                do_sample=False,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+            generated.extend(tokenizer.batch_decode(output, skip_special_tokens=True))
+
+    records = []
+    for row, caption in zip(rows, generated):
+        references = caption_lookup[int(row)]
+        records.append(
+            {
+                "row": int(row),
+                "generated": caption,
+                "references": references,
+                "unigram_f1": unigram_f1(caption, references),
+            }
+        )
+    result = {
+        "split": args.split,
+        "samples": len(records),
+        "mean_best_reference_unigram_f1": sum(
+            r["unigram_f1"] for r in records
+        )
+        / max(len(records), 1),
+        "examples": records[:25],
+        "training": prefix_state["training"],
+    }
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    write_json(args.output, result)
+    print(json.dumps(result, indent=2, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()