File size: 4,697 Bytes
212e9d7 | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | from __future__ import annotations
import argparse
import csv
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import DataLoader
from train_pet_text_alignment import PETTextAlignmentModel, PETTextDataset, collate_pet_text, load_pet_model
def retrieval_metrics(logits: np.ndarray) -> dict[str, float]:
ranks = []
for i in range(logits.shape[0]):
order = np.argsort(-logits[i])
ranks.append(int(np.where(order == i)[0][0]) + 1)
ranks = np.asarray(ranks)
return {
"recall@1": float(np.mean(ranks <= 1)),
"recall@5": float(np.mean(ranks <= 5)),
"recall@10": float(np.mean(ranks <= 10)),
"mrr": float(np.mean(1.0 / ranks)),
"median_rank": float(np.median(ranks)),
}
def split_regions(value: str) -> set[str]:
return {item for item in str(value).split("|") if item}
def factuality(logits: np.ndarray, lows: list[str], highs: list[str], k: int = 5) -> dict[str, float]:
top_text = np.argmax(logits, axis=1)
low_scores = []
high_scores = []
for query_idx, text_idx in enumerate(top_text.tolist()):
query_low = split_regions(lows[query_idx])
query_high = split_regions(highs[query_idx])
text_low = split_regions(lows[text_idx])
text_high = split_regions(highs[text_idx])
low_scores.append(len(query_low & text_low) / max(min(k, len(query_low)), 1))
high_scores.append(len(query_high & text_high) / max(min(k, len(query_high)), 1))
return {
"retrieved_text_low_overlap": float(np.mean(low_scores)),
"retrieved_text_high_overlap": float(np.mean(high_scores)),
}
def main() -> None:
parser = argparse.ArgumentParser(description="Evaluate controlled PET-to-region-text alignment.")
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--test-csv", type=Path, required=True)
parser.add_argument("--csv-out", type=Path, default=None)
parser.add_argument("--batch-size", type=int, default=8)
parser.add_argument("--num-workers", type=int, default=2)
parser.add_argument("--max-length", type=int, default=None)
args = parser.parse_args()
from transformers import AutoTokenizer
ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
saved = argparse.Namespace(**ckpt["args"])
saved.train_csv = saved.train_csv
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
pet_model = load_pet_model(saved.pet_checkpoint, saved, device)
tokenizer = AutoTokenizer.from_pretrained(saved.text_model)
model = PETTextAlignmentModel(pet_model, saved.text_model, saved.embed_dim or 256).to(device)
model.load_state_dict(ckpt["model"], strict=True)
model.eval()
output_size = tuple(saved.output_size)
dataset = PETTextDataset(args.test_csv, output_size=output_size)
loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=collate_pet_text)
pet_chunks = []
text_chunks = []
lows: list[str] = []
highs: list[str] = []
max_length = args.max_length or saved.max_length
with torch.no_grad():
for batch in loader:
image = batch["image"].to(device, non_blocking=True)
tokens = tokenizer(batch["text"], padding=True, truncation=True, max_length=max_length, return_tensors="pt")
tokens = {k: v.to(device) for k, v in tokens.items()}
pet_chunks.append(model.encode_pet(image).cpu())
text_chunks.append(model.encode_text(tokens).cpu())
lows.extend(batch["low_regions"])
highs.extend(batch["high_regions"])
pet_z = torch.cat(pet_chunks, dim=0)
text_z = torch.cat(text_chunks, dim=0)
logits = (pet_z @ text_z.T).numpy()
metrics = {"samples": float(logits.shape[0])}
metrics.update({f"pet_to_text_{k}": v for k, v in retrieval_metrics(logits).items()})
metrics.update({f"text_to_pet_{k}": v for k, v in retrieval_metrics(logits.T).items()})
metrics.update(factuality(logits, lows, highs))
for key, value in metrics.items():
print(f"{key}={value:.6f}")
if args.csv_out:
args.csv_out.parent.mkdir(parents=True, exist_ok=True)
write_header = not args.csv_out.exists()
with args.csv_out.open("a", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=["checkpoint", "test_csv", *metrics.keys()])
if write_header:
writer.writeheader()
writer.writerow({"checkpoint": str(args.checkpoint), "test_csv": str(args.test_csv), **metrics})
if __name__ == "__main__":
main()
|