95 lines
3.7 KiB
Python
95 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
import yaml
|
|
from torch.utils.data import DataLoader
|
|
|
|
from data.dataset import EvidenceDataset, collate_evidence
|
|
from models.evidence_model import EvidenceExtractorModel
|
|
|
|
|
|
def slot_counts(predictions: dict[str, torch.Tensor], targets: list[torch.Tensor]) -> dict[str, int]:
|
|
subjects = predictions["subject_logits"].argmax(-1).cpu()
|
|
predicates = predictions["predicate_logits"].argmax(-1).cpu()
|
|
objects = predictions["object_logits"].argmax(-1).cpu()
|
|
valid = predictions["valid_logits"].argmax(-1).cpu().bool()
|
|
counts = {"tp": 0, "fp": 0, "fn": 0, "component_correct": 0, "component_total": 0}
|
|
for batch_index, facts in enumerate(targets):
|
|
predicted = {
|
|
(int(subjects[batch_index, slot]), int(predicates[batch_index, slot]), int(objects[batch_index, slot]))
|
|
for slot in range(subjects.shape[1])
|
|
if valid[batch_index, slot]
|
|
}
|
|
expected = {tuple(map(int, row.tolist())) for row in facts}
|
|
counts["tp"] += len(predicted & expected)
|
|
counts["fp"] += len(predicted - expected)
|
|
counts["fn"] += len(expected - predicted)
|
|
for triple in predicted:
|
|
best = max((sum(a == b for a, b in zip(triple, target)) for target in expected), default=0)
|
|
counts["component_correct"] += best
|
|
counts["component_total"] += 3
|
|
return counts
|
|
|
|
|
|
def summarize(counts: dict[str, int]) -> dict[str, float]:
|
|
precision = counts["tp"] / max(1, counts["tp"] + counts["fp"])
|
|
recall = counts["tp"] / max(1, counts["tp"] + counts["fn"])
|
|
return {
|
|
"precision": precision,
|
|
"recall": recall,
|
|
"f1": 2 * precision * recall / max(1e-12, precision + recall),
|
|
"component_accuracy": counts["component_correct"] / max(1, counts["component_total"]),
|
|
}
|
|
|
|
|
|
@torch.no_grad()
|
|
def evaluate(
|
|
model: EvidenceExtractorModel, loader: DataLoader, device: torch.device
|
|
) -> dict[str, Any]:
|
|
model.eval()
|
|
totals = {
|
|
"visual": {"tp": 0, "fp": 0, "fn": 0, "component_correct": 0, "component_total": 0},
|
|
"text": {"tp": 0, "fp": 0, "fn": 0, "component_correct": 0, "component_total": 0},
|
|
}
|
|
for batch in loader:
|
|
outputs = model(batch["images"].to(device), batch["rules_text"])
|
|
for modality, target_name in (("visual", "visual_facts"), ("text", "text_facts")):
|
|
counts = slot_counts(outputs[modality], batch[target_name])
|
|
for key, value in counts.items():
|
|
totals[modality][key] += value
|
|
return {modality: summarize(value) for modality, value in totals.items()}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
|
|
parser.add_argument("--checkpoint", type=Path, required=True)
|
|
parser.add_argument("--backbone", choices=("tiny", "clip"), default=None)
|
|
parser.add_argument("--split", default="test_id")
|
|
args = parser.parse_args()
|
|
config = yaml.safe_load(args.config.read_text(encoding="utf-8"))
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = EvidenceExtractorModel(config, backbone_name=args.backbone).to(device)
|
|
state = torch.load(args.checkpoint, map_location=device, weights_only=True)
|
|
model.load_state_dict(state["model"])
|
|
dataset = EvidenceDataset(
|
|
Path(config["data"]["root"]) / f"{args.split}.jsonl",
|
|
config["data"]["image_size"],
|
|
)
|
|
loader = DataLoader(
|
|
dataset,
|
|
batch_size=config["training"]["batch_size"],
|
|
collate_fn=collate_evidence,
|
|
)
|
|
print(json.dumps(evaluate(model, loader, device), indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|