181 lines
6.8 KiB
Python
181 lines
6.8 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 data.schema import FACT_TYPE, NULL_TYPE, RULE_TYPE
|
|
from losses.matching_loss import hungarian_matching_loss
|
|
from models.evidence_model import EvidenceExtractorModel
|
|
|
|
COUNT_KEYS = ("tp", "fp", "fn", "component_correct", "component_total")
|
|
|
|
|
|
def empty_counts() -> dict[str, int]:
|
|
return {key: 0 for key in COUNT_KEYS}
|
|
|
|
|
|
def decode_rows(predictions: dict[str, torch.Tensor], batch_index: int) -> list[tuple[int, ...]]:
|
|
types = predictions["type_logits"].argmax(-1)[batch_index].cpu()
|
|
subjects = predictions["subject_logits"].argmax(-1)[batch_index].cpu()
|
|
predicates = predictions["predicate_logits"].argmax(-1)[batch_index].cpu()
|
|
objects = predictions["object_logits"].argmax(-1)[batch_index].cpu()
|
|
consequent_predicates = predictions["consequent_predicate_logits"].argmax(-1)[
|
|
batch_index
|
|
].cpu()
|
|
consequent_objects = predictions["consequent_object_logits"].argmax(-1)[batch_index].cpu()
|
|
rows = []
|
|
for slot, state_type in enumerate(types.tolist()):
|
|
if state_type == NULL_TYPE:
|
|
continue
|
|
consequence = (
|
|
(int(consequent_predicates[slot]), int(consequent_objects[slot]))
|
|
if state_type == RULE_TYPE
|
|
else (-1, -1)
|
|
)
|
|
rows.append(
|
|
(
|
|
state_type,
|
|
int(subjects[slot]),
|
|
int(predicates[slot]),
|
|
int(objects[slot]),
|
|
*consequence,
|
|
)
|
|
)
|
|
return rows
|
|
|
|
|
|
def slot_counts(
|
|
predictions: dict[str, torch.Tensor],
|
|
targets: list[torch.Tensor],
|
|
state_type: int | None = None,
|
|
) -> dict[str, int]:
|
|
counts = empty_counts()
|
|
for batch_index, target_tensor in enumerate(targets):
|
|
predicted = set(decode_rows(predictions, batch_index))
|
|
expected = {tuple(map(int, row.tolist())) for row in target_tensor}
|
|
if state_type is not None:
|
|
predicted = {row for row in predicted if row[0] == state_type}
|
|
expected = {row for row in expected if row[0] == state_type}
|
|
counts["tp"] += len(predicted & expected)
|
|
counts["fp"] += len(predicted - expected)
|
|
counts["fn"] += len(expected - predicted)
|
|
for row in predicted:
|
|
compatible = [target for target in expected if target[0] == row[0]]
|
|
width = 6 if row[0] == RULE_TYPE else 4
|
|
best = max(
|
|
(sum(a == b for a, b in zip(row[:width], target[:width])) for target in compatible),
|
|
default=0,
|
|
)
|
|
counts["component_correct"] += best
|
|
counts["component_total"] += width
|
|
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"]),
|
|
}
|
|
|
|
|
|
def slot_type_counts(
|
|
predictions: dict[str, torch.Tensor], targets: list[torch.Tensor]
|
|
) -> tuple[int, int, int, int]:
|
|
matching = hungarian_matching_loss(predictions, targets)
|
|
predicted_types = predictions["type_logits"].argmax(-1)
|
|
type_correct = 0
|
|
type_total = predicted_types.numel()
|
|
null_correct = 0
|
|
null_total = 0
|
|
for batch_index, (slots, target_indices) in enumerate(matching.assignments):
|
|
expected = torch.full_like(predicted_types[batch_index], NULL_TYPE)
|
|
if len(slots):
|
|
expected[slots] = targets[batch_index].to(expected.device)[target_indices, 0]
|
|
type_correct += int((predicted_types[batch_index] == expected).sum())
|
|
null_positions = expected == NULL_TYPE
|
|
null_correct += int((predicted_types[batch_index][null_positions] == NULL_TYPE).sum())
|
|
null_total += int(null_positions.sum())
|
|
return null_correct, null_total, type_correct, type_total
|
|
|
|
|
|
@torch.no_grad()
|
|
def evaluate(
|
|
model: EvidenceExtractorModel, loader: DataLoader, device: torch.device
|
|
) -> dict[str, Any]:
|
|
model.eval()
|
|
totals = {
|
|
"visual_fact": empty_counts(),
|
|
"text_all": empty_counts(),
|
|
"text_fact": empty_counts(),
|
|
"text_rule": empty_counts(),
|
|
}
|
|
null_correct = 0
|
|
null_total = 0
|
|
type_correct = 0
|
|
type_total = 0
|
|
for batch in loader:
|
|
outputs = model(batch["images"].to(device), batch["rules_text"])
|
|
requests = (
|
|
("visual_fact", "visual", "visual_targets", FACT_TYPE),
|
|
("text_all", "text", "text_targets", None),
|
|
("text_fact", "text", "text_targets", FACT_TYPE),
|
|
("text_rule", "text", "text_targets", RULE_TYPE),
|
|
)
|
|
for name, modality, target_name, state_type in requests:
|
|
counts = slot_counts(outputs[modality], batch[target_name], state_type)
|
|
for key, value in counts.items():
|
|
totals[name][key] += value
|
|
for modality, target_name in (
|
|
("visual", "visual_targets"),
|
|
("text", "text_targets"),
|
|
):
|
|
counts = slot_type_counts(outputs[modality], batch[target_name])
|
|
null_correct += counts[0]
|
|
null_total += counts[1]
|
|
type_correct += counts[2]
|
|
type_total += counts[3]
|
|
result = {name: summarize(counts) for name, counts in totals.items()}
|
|
result["null_slot_accuracy"] = null_correct / max(1, null_total)
|
|
result["slot_type_accuracy"] = type_correct / max(1, type_total)
|
|
return result
|
|
|
|
|
|
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", "siglip2"), 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"], strict=False)
|
|
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()
|