1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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()
|
