zsgdp/verify/table_structure.py

"""Table structure similarity against ground-truth tables.

Definitions (pinned):

- A TableRecord is a dict with keys: rows (list[list[str]] of normalized cell
  strings), page_num (int|None), table_id (str|None), bbox (xyxy|None).
- Predictions are matched to truths greedily, page-aware. The match score
  blends shape similarity (row/col count) with cell-content overlap, so a
  prediction with the same dimensions and overlapping headers wins over one
  with extra phantom rows.
- For each matched pair we compute:
  - shape_similarity = 1 - mean(|rows_p - rows_t|/max_rows, |cols_p - cols_t|/max_cols)
    bounded to [0, 1].
  - cell_content_f1 = F1 over the multiset of normalized cell strings.
    Multiset (not set) so duplicate "0.00" cells in financial tables count.
  - score = 0.5 * shape_similarity + 0.5 * cell_content_f1.
- Document-level aggregates:
  - mean_table_score = mean(score) over matched pairs.
  - table_match_rate = matched_pairs / max(len(predictions), len(truths)).
  - table_count_delta = len(predictions) - len(truths).
- Edge cases: empty/empty -> mean_table_score=1.0 (vacuous); one side empty
  -> 0.0 with all unmatched recorded as fp/fn.
"""

from __future__ import annotations

import re
from typing import Any, Iterable

TableRecord = dict[str, Any]


def compute_table_structure_score(
    predictions: Iterable[TableRecord],
    truths: Iterable[TableRecord],
) -> dict[str, Any]:
    pred_list = [item for item in (_normalize(record) for record in predictions) if item is not None]
    truth_list = [item for item in (_normalize(record) for record in truths) if item is not None]

    if not pred_list and not truth_list:
        return {
            "prediction_count": 0,
            "truth_count": 0,
            "matched_pair_count": 0,
            "table_match_rate": 1.0,
            "mean_table_score": 1.0,
            "mean_shape_similarity": 1.0,
            "mean_cell_content_f1": 1.0,
            "table_count_delta": 0,
            "matches": [],
            "unmatched_predictions": [],
            "unmatched_truths": [],
        }

    pairs: list[tuple[float, int, int, dict[str, float]]] = []
    for pred_index, prediction in enumerate(pred_list):
        for truth_index, truth in enumerate(truth_list):
            if prediction["page_num"] is not None and truth["page_num"] is not None:
                if prediction["page_num"] != truth["page_num"]:
                    continue
            scores = _score_pair(prediction, truth)
            if scores["match_score"] <= 0.0:
                continue
            pairs.append((scores["match_score"], pred_index, truth_index, scores))

    pairs.sort(key=lambda item: item[0], reverse=True)
    pred_taken = [False] * len(pred_list)
    truth_taken = [False] * len(truth_list)
    matches: list[dict[str, Any]] = []
    for _match_score, pred_index, truth_index, scores in pairs:
        if pred_taken[pred_index] or truth_taken[truth_index]:
            continue
        pred_taken[pred_index] = True
        truth_taken[truth_index] = True
        matches.append(
            {
                "prediction_index": pred_index,
                "truth_index": truth_index,
                "page_num": pred_list[pred_index]["page_num"] or truth_list[truth_index]["page_num"],
                "score": scores["score"],
                "shape_similarity": scores["shape_similarity"],
                "cell_content_f1": scores["cell_content_f1"],
                "predicted_shape": [len(pred_list[pred_index]["rows"]), _max_cols(pred_list[pred_index]["rows"])],
                "truth_shape": [len(truth_list[truth_index]["rows"]), _max_cols(truth_list[truth_index]["rows"])],
            }
        )

    matched_pair_count = len(matches)
    denominator = max(len(pred_list), len(truth_list))
    mean_table_score = (sum(match["score"] for match in matches) / denominator) if denominator else 1.0
    mean_shape = (sum(match["shape_similarity"] for match in matches) / matched_pair_count) if matched_pair_count else 0.0
    mean_cell = (sum(match["cell_content_f1"] for match in matches) / matched_pair_count) if matched_pair_count else 0.0

    return {
        "prediction_count": len(pred_list),
        "truth_count": len(truth_list),
        "matched_pair_count": matched_pair_count,
        "table_match_rate": (matched_pair_count / denominator) if denominator else 1.0,
        "mean_table_score": mean_table_score,
        "mean_shape_similarity": mean_shape,
        "mean_cell_content_f1": mean_cell,
        "table_count_delta": len(pred_list) - len(truth_list),
        "matches": matches,
        "unmatched_predictions": [index for index, taken in enumerate(pred_taken) if not taken],
        "unmatched_truths": [index for index, taken in enumerate(truth_taken) if not taken],
    }


def markdown_to_rows(markdown: str | None) -> list[list[str]]:
    if not markdown:
        return []
    rows: list[list[str]] = []
    for line in markdown.splitlines():
        stripped = line.strip()
        if not (stripped.startswith("|") and stripped.endswith("|")):
            continue
        cells = [_normalize_cell(cell) for cell in stripped.strip("|").split("|")]
        if cells and all(re.fullmatch(r":?-{3,}:?", cell) for cell in cells if cell):
            continue
        rows.append(cells)
    return rows


def html_to_rows(html: str | None) -> list[list[str]]:
    if not html:
        return []
    rows: list[list[str]] = []
    for tr_match in re.finditer(r"<tr[^>]*>(.*?)</tr>", html, flags=re.IGNORECASE | re.DOTALL):
        tr_body = tr_match.group(1)
        cells = [
            _normalize_cell(_strip_html(cell.group(1)))
            for cell in re.finditer(r"<t[hd][^>]*>(.*?)</t[hd]>", tr_body, flags=re.IGNORECASE | re.DOTALL)
        ]
        if cells:
            rows.append(cells)
    return rows


def _strip_html(text: str) -> str:
    return re.sub(r"<[^>]+>", " ", text)


def _normalize_cell(cell: str) -> str:
    return " ".join(cell.replace("\xa0", " ").split()).strip().lower()


def _normalize(record: TableRecord | None) -> TableRecord | None:
    if not isinstance(record, dict):
        return None
    rows = record.get("rows")
    if rows is None:
        rows = markdown_to_rows(record.get("markdown"))
    if not rows:
        rows = html_to_rows(record.get("html"))
    rows = [[_normalize_cell(cell) for cell in row] for row in rows if row]
    if not rows:
        return None
    page_num = record.get("page_num")
    return {
        "rows": rows,
        "page_num": int(page_num) if isinstance(page_num, (int, float)) else None,
        "table_id": record.get("table_id"),
        "bbox": record.get("bbox"),
    }


def _score_pair(prediction: TableRecord, truth: TableRecord) -> dict[str, float]:
    shape_similarity = _shape_similarity(prediction["rows"], truth["rows"])
    cell_content_f1 = _multiset_f1(_cell_multiset(prediction["rows"]), _cell_multiset(truth["rows"]))
    score = 0.5 * shape_similarity + 0.5 * cell_content_f1
    # Prefer matches that share at least one cell content token to break ties.
    overlap_bonus = 0.0
    if cell_content_f1 > 0:
        overlap_bonus = 0.01
    return {
        "shape_similarity": shape_similarity,
        "cell_content_f1": cell_content_f1,
        "score": score,
        "match_score": score + overlap_bonus,
    }


def _shape_similarity(a: list[list[str]], b: list[list[str]]) -> float:
    rows_a = len(a)
    rows_b = len(b)
    cols_a = _max_cols(a)
    cols_b = _max_cols(b)
    max_rows = max(rows_a, rows_b, 1)
    max_cols = max(cols_a, cols_b, 1)
    row_diff = abs(rows_a - rows_b) / max_rows
    col_diff = abs(cols_a - cols_b) / max_cols
    return max(0.0, 1.0 - 0.5 * (row_diff + col_diff))


def _max_cols(rows: list[list[str]]) -> int:
    return max((len(row) for row in rows), default=0)


def _cell_multiset(rows: list[list[str]]) -> list[str]:
    return [cell for row in rows for cell in row if cell]


def _multiset_f1(predicted: list[str], truth: list[str]) -> float:
    if not predicted and not truth:
        return 1.0
    if not predicted or not truth:
        return 0.0
    truth_remaining = list(truth)
    tp = 0
    for cell in predicted:
        if cell in truth_remaining:
            truth_remaining.remove(cell)
            tp += 1
    fp = len(predicted) - tp
    fn = len(truth_remaining)
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    return (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0