bibraeval 0.1.0__py3-none-any.whl

bibraeval/__init__.py

@@ -0,0 +1,22 @@
+"""BIBRA evaluation package."""
+
+from .aggregator import Aggregator, IntermediateResults
+from .comparator import Comparator
+from .data_ingester import Record, RecordCollection
+from .evaluator import Evaluator
+from .field_metrics import Exact, FieldMetric, Levenshtein
+from .metric_conf import FieldMetricConfig, FieldMetricMapping
+
+__all__ = [
+    "Aggregator",
+    "Comparator",
+    "Evaluator",
+    "Exact",
+    "FieldMetric",
+    "FieldMetricConfig",
+    "FieldMetricMapping",
+    "IntermediateResults",
+    "Levenshtein",
+    "Record",
+    "RecordCollection",
+]

bibraeval/aggregator.py

@@ -0,0 +1,218 @@
+"""Aggregation of field-level comparisons into macro-averaged metrics."""
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+import polars as pl
+
+from bibraeval.comparator import Comparator
+
+AggregationMode = Literal["doc-avg", "field-avg", "micro-avg"]
+
+METRICS = ("f1", "prec", "rec")
+_PRIMARY_KEYS = ("doc_id", "field_name")
+_UNGROUPABLE = ("gt_value", "pred_value")
+
+
+@dataclass(frozen=True)
+class IntermediateResults:
+    """Per-group metrics and the grouping metadata used to produce them.
+
+    ``data`` contains grouping columns, precision/recall/F1, and support
+    counts. ``group_by`` records the grouping columns; ``mode`` identifies
+    document-, field-, or micro-averaging.
+    """
+
+    data: pl.DataFrame
+    group_by: tuple[str, ...]
+    mode: AggregationMode
+
+    @property
+    def strata(self) -> list[str]:
+        """Grouping columns other than the built-in document/field keys."""
+        return [col for col in self.group_by if col not in _PRIMARY_KEYS]
+
+    @property
+    def weighted_summary(self) -> bool:
+        """Whether field weights are applied when producing calling the
+        summarise_results method."""
+        return "field_name" in self.group_by
+
+
+class Aggregator:
+    """Aggregate a comparator's field-level scores into precision/recall/F1.
+
+    Cell scores are computed at initialization. Add any custom strata columns to
+    ``comparator.comparison_matrix`` before constructing the aggregator.
+    """
+
+    def __init__(self, comparator: Comparator):
+        self.comparator = comparator
+        self.intermediate_results: IntermediateResults | None = None
+        self.comparison_matrix = self.comparator.compute_cell_agreement()
+        self.cells = self._compute_cell_scores(self.comparison_matrix)
+
+    @staticmethod
+    def _compute_cell_scores(cells: pl.DataFrame) -> pl.DataFrame:
+        """Mask cell agreement by presence for precision and recall."""
+        agreement = pl.col("cell_agreement")
+        return cells.with_columns(
+            pl.when(~pl.col("pred_present"))
+            .then(None)
+            .when(pl.col("gt_present"))
+            .then(agreement)
+            .otherwise(0.0)
+            .cast(pl.Float64)
+            .alias("prec_score"),
+            pl.when(~pl.col("gt_present"))
+            .then(None)
+            .when(pl.col("pred_present"))
+            .then(agreement)
+            .otherwise(0.0)
+            .cast(pl.Float64)
+            .alias("rec_score"),
+        )
+
+    def compute_intermediate_results(
+        self, group_by: str | Sequence[str] | None
+    ) -> IntermediateResults:
+        """Compute precision, recall and f1 per group.
+
+        For macro averaging `group_by` must contain exactly one of `doc_id` or
+        `field_name`; other columns may be included as strata.
+        When grouping by `field_name`, field weights are carried to the summary
+        step; otherwise they are applied as weighted means within each group.
+
+        Set ``group_by=None`` for micro-averaging. For a stratified micro
+        average, group only by the desired (additional) stratification columns.
+        Value columns cannot be used as grouping keys.
+
+        Raises:
+            ValueError: If a grouping column is missing, a value column is
+                requested, or both primary keys are included.
+        """
+        if group_by is None:
+            group_cols = ()
+        elif isinstance(group_by, str):
+            group_cols = (group_by,)
+        else:
+            group_cols = tuple(group_by)
+        self._validate_group_by(group_cols)
+
+        weighted_summary = "field_name" in group_cols
+        aggs: list[pl.Expr] = []
+        for name, score in (("prec", "prec_score"), ("rec", "rec_score")):
+            aggs.append(pl.col(score).count().alias(f"support_{name}"))
+            if weighted_summary:
+                aggs.append(pl.col(score).mean().alias(name))
+            else:
+                aggs.append(_weighted_mean(pl.col(score), pl.col("weight")).alias(name))
+        if weighted_summary:
+            aggs.append(pl.col("weight").first())
+
+        if group_cols:
+            data = self.cells.group_by(group_cols, maintain_order=True).agg(aggs)
+        else:
+            data = self.cells.select(aggs)
+
+        data = data.with_columns(_f1(pl.col("prec"), pl.col("rec")).alias("f1"))
+        if group_cols:
+            data = data.sort(group_cols)
+
+        self.intermediate_results = IntermediateResults(
+            data=data, group_by=group_cols, mode=_mode(group_cols)
+        )
+        return self.intermediate_results
+
+    def _validate_group_by(self, group_cols: tuple[str, ...]) -> None:
+        if all(col in group_cols for col in _PRIMARY_KEYS):
+            raise ValueError("group_by cannot include both 'doc_id' and 'field_name'.")
+        missing = [col for col in group_cols if col not in self.cells.columns]
+        if missing:
+            raise ValueError(f"group_by columns not found: {missing}")
+        invalid = [col for col in group_cols if col in _UNGROUPABLE]
+        if invalid:
+            raise ValueError(f"Cannot group by value columns: {invalid}")
+
+    def summarise_results(
+        self, intermediate: IntermediateResults | None = None
+    ) -> pl.DataFrame:
+        """Summarize intermediate results while preserving strata.
+
+        Macro modes average the intermediate results; micro mode returns its
+        pooled result without another aggregation. The long-form output has one
+        row per metric and stratum, with ``metric``, ``mode``, ``value``, and
+        ``support`` columns in addition to any strata columns.
+
+        Raises:
+            ValueError: If no intermediate results are available and none are
+                passed explicitly.
+        """
+        if intermediate is None:
+            intermediate = self.intermediate_results
+        if intermediate is None:
+            raise ValueError(
+                "No intermediate results; call compute_intermediate_results first."
+            )
+
+        strata = intermediate.strata
+        index = strata + (["weight"] if intermediate.weighted_summary else [])
+        long_index = (
+            [*index, "support_prec", "support_rec"]
+            if intermediate.mode in ["micro-avg", "field-avg"]
+            else index
+        )
+        long = intermediate.data.select(*long_index, *METRICS).unpivot(
+            on=list(METRICS), index=long_index, variable_name="metric"
+        )
+        value = (
+            _weighted_mean(pl.col("value"), pl.col("weight"))
+            if intermediate.weighted_summary
+            else pl.col("value").mean()
+        )
+        support = (
+            pl.when(pl.col("metric") == "prec")
+            .then(pl.col("support_prec"))
+            .when(pl.col("metric") == "rec")
+            .then(pl.col("support_rec"))
+            .otherwise(pl.max_horizontal("support_prec", "support_rec"))
+            .mean()
+            if intermediate.mode in ["micro-avg", "field-avg"]
+            else pl.col("value").count()
+        )
+        return (
+            long.group_by([*strata, "metric"])
+            .agg(value.alias("value"), support.alias("support"))
+            .with_columns(pl.lit(intermediate.mode).alias("mode"))
+            .select(*strata, "metric", "mode", "value", "support")
+            .sort([*strata, "metric"])
+        )
+
+
+def _weighted_mean(score: pl.Expr, weight: pl.Expr) -> pl.Expr:
+    """Weighted mean over non-null scores; null if no weight remains."""
+    weight_sum = weight.filter(score.is_not_null()).sum()
+    return (
+        pl.when(weight_sum > 0)
+        .then((score * weight).sum() / weight_sum)
+        .otherwise(None)
+    )
+
+
+def _f1(prec: pl.Expr, rec: pl.Expr) -> pl.Expr:
+    return (
+        pl.when(prec.is_null() & rec.is_null())
+        .then(None)
+        .when(prec.is_null() | rec.is_null() | ((prec + rec) == 0))
+        .then(0.0)
+        .otherwise(2 * prec * rec / (prec + rec))
+    )
+
+
+def _mode(group_cols: tuple[str, ...]) -> AggregationMode:
+    if "doc_id" in group_cols:
+        return "doc-avg"
+    if "field_name" in group_cols:
+        return "field-avg"
+    return "micro-avg"

bibraeval/comparator.py

@@ -0,0 +1,178 @@
+from typing import Any
+
+import polars as pl
+
+from bibraeval.data_ingester import MetadataT, RecordCollection
+from bibraeval.field_metrics import Exact, FieldMetric
+from bibraeval.metric_conf import FieldMetricMapping
+
+
+class Comparator:
+    """Match two record collections and score their metadata fields.
+
+    The comparison matrix contains one row per document and metadata field.
+    Document ID sets must match unless ``drop_mismatched_doc_ids`` is enabled;
+    in that case only IDs present in both collections are compared.
+    """
+
+    def __init__(
+        self,
+        ground_truth: RecordCollection[MetadataT],
+        predictions: RecordCollection[MetadataT],
+        metric_schema: FieldMetricMapping | dict[str, Any] | None = None,
+        drop_mismatched_doc_ids: bool = False,
+    ):
+        """Create a comparator for validated record collections.
+
+        Args:
+            ground_truth: Reference records keyed by document ID.
+            predictions: Predicted records keyed by document ID.
+            metric_schema: Per-field metric configuration. Unlisted fields are
+                omitted from scoring when a schema is supplied.
+            drop_mismatched_doc_ids: If true, compare only IDs in both
+                collections; otherwise a mismatch raises ``ValueError``.
+        """
+        self.ground_truth = ground_truth
+        self.predictions = predictions
+        self.drop_mismatched_doc_ids = drop_mismatched_doc_ids
+        if isinstance(metric_schema, dict):
+            metric_schema = FieldMetricMapping.from_dict(metric_schema)
+        self.metric_schema = metric_schema
+
+        if not self.drop_mismatched_doc_ids:
+            self._check_doc_id_match()
+        self.comparison_matrix = self._fuse_records()
+
+    def _check_doc_id_match(self) -> bool:
+        """Check if two record collections contain the same doc_ids."""
+        gt_doc_ids = set(self.ground_truth.records)
+        pred_doc_ids = set(self.predictions.records)
+        if gt_doc_ids == pred_doc_ids:
+            return True
+
+        mismatches_gt_not_pred = gt_doc_ids - pred_doc_ids
+        mismatches_pred_not_gold = pred_doc_ids - gt_doc_ids
+        n_mismatches = len(gt_doc_ids ^ pred_doc_ids)
+        raise ValueError(
+            f"Record collections have {n_mismatches} mismatched doc_ids: "
+            f"ground_truth not in predictions: {mismatches_gt_not_pred}, "
+            f"predictions not in ground_truth: {mismatches_pred_not_gold}"
+            f"Use drop_mismatched_doc_ids=True to ignore mismatched doc_ids."
+        )
+
+    def _fuse_records(self) -> pl.DataFrame:
+        """Fuse records into one long Polars DataFrame by doc_id and field.
+
+        The ``gt_present`` and ``pred_present`` flags are false for ``None``,
+        empty strings, and empty lists. Other values, including ``0`` and
+        ``False``, count as present.
+
+        Columns returned:
+        - `doc_id`
+        - `field_name`
+        - `gt_present`
+        - `pred_present`
+        - `gt_value` (may be nested for list fields)
+        - `pred_value` (may be nested for list fields)
+        """
+        fused_data = []
+        for doc_id in sorted(self.ground_truth.records):
+            if self.drop_mismatched_doc_ids and doc_id not in self.predictions.records:
+                continue
+            gt_record = self.ground_truth.records[doc_id]
+            pred_record = self.predictions.records[doc_id]
+            gt_metadata = gt_record.metadata.model_dump(by_alias=True)
+            pred_metadata = pred_record.metadata.model_dump(by_alias=True)
+            field_names = sorted(gt_metadata.keys() | pred_metadata.keys())
+
+            for field_name in field_names:
+                gt_value = gt_metadata.get(field_name)
+                pred_value = pred_metadata.get(field_name)
+                fused_data.append(
+                    {
+                        "doc_id": doc_id,
+                        "field_name": field_name,
+                        # custom test for presence of value, exclusing empty
+                        # strings and empty lists
+                        "gt_present": self._is_present(gt_value),
+                        "pred_present": self._is_present(pred_value),
+                        "gt_value": gt_value,
+                        "pred_value": pred_value,
+                    }
+                )
+
+        return pl.DataFrame(
+            fused_data,
+            schema_overrides={"gt_value": pl.Object, "pred_value": pl.Object},
+        )
+
+    @staticmethod
+    def _is_present(value: object) -> bool:
+        return value is not None and value != "" and value != []
+
+    def compute_cell_agreement(
+        self,
+        comparison_matrix: pl.DataFrame | None = None,
+        metric: FieldMetric | None = None,
+    ) -> pl.DataFrame:
+        """Compute cell-level agreement between gold and predicted values.
+
+        With a `metric_schema`, only fields listed in the schema are scored, each
+        with its configured metric and weight. Otherwise all fields are scored
+        with `metric` (default `Exact()`) and weight 1.0.
+
+        Args:
+            comparison_matrix: Matrix to score; defaults to this comparator's
+                matrix. It must contain ``field_name``, ``gt_value``, and
+                ``pred_value`` columns. Additional columns may be present to
+                facilitate stratified analysis.
+            metric: Metric used when no metric schema was supplied.
+
+        Raises:
+            ValueError: If both ``metric`` and a configured metric schema are
+                supplied.
+        """
+        if comparison_matrix is None:
+            comparison_matrix = self.comparison_matrix
+
+        if self.metric_schema is not None:
+            if metric is not None:
+                raise ValueError("Pass either metric or metric_schema, not both.")
+            return self._compute_schema_agreement(comparison_matrix, self.metric_schema)
+
+        if metric is None:
+            metric = Exact()
+
+        scores = [
+            metric.score(gt_value, pred_value)
+            for gt_value, pred_value in comparison_matrix.select(
+                "gt_value", "pred_value"
+            ).iter_rows()
+        ]
+        return comparison_matrix.with_columns(
+            pl.Series("cell_agreement", scores, dtype=pl.Float64),
+            pl.lit(1.0, dtype=pl.Float64).alias("weight"),
+        )
+
+    def _compute_schema_agreement(
+        self, comparison_matrix: pl.DataFrame, schema: FieldMetricMapping
+    ) -> pl.DataFrame:
+        comparison_matrix = comparison_matrix.filter(
+            pl.col("field_name").is_in(list(schema.fields))
+        )
+        metrics = {name: schema.build_metric(name) for name in schema.fields}
+
+        scores = [
+            metrics[field_name].score(gt_value, pred_value)
+            for field_name, gt_value, pred_value in comparison_matrix.select(
+                "field_name", "gt_value", "pred_value"
+            ).iter_rows()
+        ]
+        weights = [
+            schema.fields[field_name].weight
+            for field_name in comparison_matrix["field_name"]
+        ]
+        return comparison_matrix.with_columns(
+            pl.Series("cell_agreement", scores, dtype=pl.Float64),
+            pl.Series("weight", weights, dtype=pl.Float64),
+        )

bibraeval/data_ingester.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, TypeVar
+
+from pydantic import BaseModel, Field
+
+MetadataT = TypeVar("MetadataT", bound=BaseModel)
+
+
+class Record[MetadataT: BaseModel](BaseModel):
+    """A record with a stable document identifier and validated metadata.
+
+    ``metadata`` is parsed with the caller-supplied Pydantic model; ``doc_id``
+    is kept separately so records can be matched independently of row order.
+    """
+
+    doc_id: str
+    metadata: MetadataT
+
+    @classmethod
+    def from_payload(
+        cls,
+        doc_id: str,
+        payload: dict[str, Any],
+        schema: type[MetadataT],
+    ) -> Record[MetadataT]:
+        """Validate a metadata payload and create a record for ``doc_id``.
+
+        Args:
+            doc_id: Stable identifier used to match this record across datasets.
+            payload: Metadata fields accepted by ``schema``.
+            schema: Pydantic model used to validate the metadata.
+        """
+        validated_metadata = schema.model_validate(payload)
+        return cls(doc_id=doc_id, metadata=validated_metadata)
+
+
+class RecordCollection[MetadataT: BaseModel](BaseModel):
+    """A collection of records keyed by document identifier."""
+
+    records: dict[str, Record[MetadataT]] = Field(default_factory=dict)
+
+    @classmethod
+    def create_from_dir(
+        cls,
+        directory: str | Path,
+        schema: type[MetadataT],
+    ) -> RecordCollection[MetadataT]:
+        """Load ``*.json`` records from a directory, using each filename stem
+        as its document ID.
+
+        Files must contain JSON objects. Other files and nested directories are
+        ignored.
+        """
+        collection = cls()
+        for record_path in sorted(Path(directory).glob("*.json")):
+            with record_path.open(encoding="utf-8") as record_file:
+                payload = json.load(record_file)
+            if not isinstance(payload, dict):
+                msg = f"Record file {record_path} must contain a JSON object."
+                raise TypeError(msg)
+            collection.add_payload(
+                doc_id=record_path.stem,
+                payload=payload,
+                schema=schema,
+            )
+        return collection
+
+    @classmethod
+    def create_from_file(
+        cls,
+        file_path: str | Path,
+        schema: type[MetadataT],
+    ) -> RecordCollection[MetadataT]:
+        """Load JSONL records whose objects each contain a non-empty ``doc_id``.
+
+        Blank lines are ignored. Duplicate IDs replace the earlier record.
+        """
+        collection = cls()
+        with Path(file_path).open(encoding="utf-8") as jsonl_file:
+            for line_number, line in enumerate(jsonl_file, start=1):
+                if not line.strip():
+                    continue
+                payload = json.loads(line)
+                if not isinstance(payload, dict):
+                    msg = (
+                        f"Line {line_number} in {file_path} must contain a JSON object."
+                    )
+                    raise TypeError(msg)
+                doc_id = payload.get("doc_id")
+                if not isinstance(doc_id, str) or not doc_id:
+                    msg = f"Line {line_number} in {file_path} must contain a doc_id."
+                    raise ValueError(msg)
+                raw_record = payload.get("record")
+                if not isinstance(raw_record, dict) or not raw_record:
+                    msg = (
+                        f"Line {line_number} in {file_path} "
+                        "must contain a record object."
+                    )
+                    raise TypeError(msg)
+                collection.add_payload(
+                    doc_id=doc_id,
+                    payload=raw_record,
+                    schema=schema,
+                )
+        return collection
+
+    def add_record(self, record: Record[MetadataT]) -> None:
+        """Store a record, replacing any existing record with the same ID."""
+        self.records[record.doc_id] = record
+
+    def get_record(self, doc_id: str) -> Record[MetadataT] | None:
+        """Return a record by ID, or ``None`` when the ID is absent."""
+        return self.records.get(doc_id)
+
+    def add_payload(
+        self,
+        doc_id: str,
+        payload: dict[str, Any],
+        schema: type[MetadataT],
+    ) -> Record[MetadataT]:
+        """Validate payload against metadata schema and store
+        payload, replacing an existing matching ID."""
+        record = Record.from_payload(doc_id=doc_id, payload=payload, schema=schema)
+        self.add_record(record)
+        return record
+
+    def as_list(self) -> list[Record[MetadataT]]:
+        """Return the stored records as a list in insertion order."""
+        return list(self.records.values())