pear-mt 1.0.0__py3-none-any.whl

pear/__init__.py

@@ -0,0 +1,31 @@
+"""PEAR-MT: pairwise QE metrics for machine translation evaluation."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from pear.inference import (
+    build_pairwise_samples,
+    build_reference_anchored_samples,
+    load_metric,
+    score_pairwise,
+    score_reference_anchored,
+)
+from pear.mbr import pear_utility_matrix, select_mbr_hypothesis
+from pear.models import download_model, load_from_checkpoint
+
+try:
+    __version__ = version("pear-mt")
+except PackageNotFoundError:
+    __version__ = "0+unknown"
+
+__all__ = [
+    "__version__",
+    "build_pairwise_samples",
+    "build_reference_anchored_samples",
+    "download_model",
+    "load_from_checkpoint",
+    "load_metric",
+    "pear_utility_matrix",
+    "score_pairwise",
+    "score_reference_anchored",
+    "select_mbr_hypothesis",
+]

pear/commands/__init__.py

@@ -0,0 +1 @@
+"""Console entry points for PEAR."""

pear/commands/main.py

@@ -0,0 +1,214 @@
+"""Command line interface for PEAR."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+from pathlib import Path
+
+from pear import __version__
+from pear.inference import load_metric, score_pairwise, score_reference_anchored
+from pear.mbr import pear_utility_matrix, select_mbr_hypothesis
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run PEAR MT metric inference.")
+    parser.add_argument(
+        "--version", action="version", version=f"%(prog)s {__version__}"
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    score = sub.add_parser("score", help="Score TSV rows with src/mt_0/mt_1 columns.")
+    _add_model_args(score)
+    score.add_argument("--input", type=Path, required=True)
+    score.add_argument("--output", type=Path)
+    score.add_argument("--mode", choices=["pairwise", "reference"], default="pairwise")
+    score.add_argument("--both", action="store_true", help="Score both pair orders.")
+    score.add_argument(
+        "--batch-size",
+        type=int,
+        default=8,
+        help="Number of pairwise examples per PEAR forward batch.",
+    )
+    score.add_argument(
+        "--gpus",
+        type=_parse_gpus,
+        default="auto",
+        help=(
+            "GPU usage: auto uses one available CUDA/MPS device, 0 forces CPU, "
+            "and 1 forces one accelerator device."
+        ),
+    )
+
+    mbr = sub.add_parser("mbr", help="Run PEAR MBR over JSONL n-best lists.")
+    _add_model_args(mbr)
+    mbr.add_argument(
+        "--input",
+        type=Path,
+        required=True,
+        help="JSONL with src and hypotheses fields.",
+    )
+    mbr.add_argument("--output", type=Path)
+    mbr.add_argument("--utility", choices=["full", "half"], default="half")
+    mbr.add_argument(
+        "--batch-size",
+        type=int,
+        default=8,
+        help="Number of pairwise utility comparisons per PEAR forward batch.",
+    )
+    mbr.add_argument(
+        "--gpus",
+        type=_parse_gpus,
+        default="auto",
+        help=(
+            "GPU usage: auto uses one available CUDA/MPS device, 0 forces CPU, "
+            "and 1 forces one accelerator device."
+        ),
+    )
+
+    args = parser.parse_args()
+    if args.command == "score":
+        _score(args)
+    elif args.command == "mbr":
+        _mbr(args)
+
+
+def _add_model_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--model",
+        default="pear",
+        help=(
+            "Model to load: pear, pear-xl, a Hugging Face repo id, or a local "
+            ".ckpt path. Defaults to pear."
+        ),
+    )
+    parser.add_argument(
+        "--hf-model",
+        help=(
+            "Explicit Hugging Face repo id or alias to load. This is a clearer "
+            "alternative to passing the same value through --model."
+        ),
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        help=(
+            "Explicit local Lightning .ckpt path to load. This is a clearer "
+            "alternative to passing the same path through --model."
+        ),
+    )
+    parser.add_argument("--cache-dir", type=Path)
+    parser.add_argument("--local-files-only", action="store_true")
+    parser.add_argument(
+        "--revision",
+        help="Hugging Face checkpoint revision (branch, tag, or commit SHA).",
+    )
+    parser.add_argument(
+        "--encoder-revision",
+        help="Hugging Face base-encoder revision (branch, tag, or commit SHA).",
+    )
+
+
+def _parse_gpus(value: str) -> int | str:
+    if value == "auto":
+        return value
+    try:
+        gpus = int(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError("expected auto, 0, or 1") from exc
+    if gpus not in {0, 1}:
+        raise argparse.ArgumentTypeError("expected auto, 0, or 1")
+    return gpus
+
+
+def _resolve_model_arg(args: argparse.Namespace) -> str | Path:
+    selected = [args.checkpoint is not None, args.hf_model is not None]
+    if sum(selected) > 1:
+        raise SystemExit("Use only one of --checkpoint or --hf-model.")
+    if args.checkpoint is not None:
+        return args.checkpoint
+    if args.hf_model is not None:
+        return args.hf_model
+    return args.model
+
+
+def _score(args: argparse.Namespace) -> None:
+    model = load_metric(
+        _resolve_model_arg(args),
+        cache_dir=args.cache_dir,
+        local_files_only=args.local_files_only,
+        revision=args.revision,
+        encoder_revision=args.encoder_revision,
+    )
+    rows = list(csv.DictReader(args.input.open(newline=""), delimiter="\t"))
+    sources = [row["src"] for row in rows]
+    mode = "both" if args.both else "single"
+    if args.mode == "pairwise":
+        scores = score_pairwise(
+            model,
+            sources,
+            [r["mt_0"] for r in rows],
+            [r["mt_1"] for r in rows],
+            mode=mode,
+            batch_size=args.batch_size,
+            gpus=args.gpus,
+        )
+    else:
+        scores = score_reference_anchored(
+            model,
+            sources,
+            [r["mt"] for r in rows],
+            [r["ref"] for r in rows],
+            mode=mode,
+            batch_size=args.batch_size,
+            gpus=args.gpus,
+        )
+    if isinstance(scores, dict):
+        for key, values in scores.items():
+            for row, value in zip(rows, values):
+                row[f"pear_{key}"] = value
+    else:
+        for row, value in zip(rows, scores):
+            row["pear_score"] = value
+    fieldnames = list(rows[0]) if rows else []
+    output = args.output or args.input.with_suffix(".pear.tsv")
+    with output.open("w", newline="") as out:
+        writer = csv.DictWriter(out, fieldnames=fieldnames, delimiter="\t")
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def _mbr(args: argparse.Namespace) -> None:
+    model = load_metric(
+        _resolve_model_arg(args),
+        cache_dir=args.cache_dir,
+        local_files_only=args.local_files_only,
+        revision=args.revision,
+        encoder_revision=args.encoder_revision,
+    )
+    output = args.output or args.input.with_suffix(".pear-mbr.jsonl")
+    with args.input.open() as src, output.open("w") as out:
+        for line in src:
+            item = json.loads(line)
+            matrix = pear_utility_matrix(
+                model,
+                item["src"],
+                item["hypotheses"],
+                mode=args.utility,
+                batch_size=args.batch_size,
+                gpus=args.gpus,
+            )
+            index, utility = select_mbr_hypothesis(matrix)
+            item.update(
+                {
+                    "selected_index": index,
+                    "selected": item["hypotheses"][index],
+                    "utility": utility,
+                }
+            )
+            out.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+
+if __name__ == "__main__":
+    main()

pear/encoders/__init__.py

@@ -0,0 +1,8 @@
+from .xlmr import XLMREncoder
+from .xlmr_xl import XLMRXLEncoder
+
+str2encoder = {
+    "XLM-RoBERTa": XLMREncoder,
+    "XLM-RoBERTa-XL": XLMRXLEncoder,
+    "InfoXLM": XLMREncoder,
+}

pear/encoders/base.py

@@ -0,0 +1,393 @@
+r"""
+Encoder Model base
+====================
+    Module defining the common interface between all pretrained encoder models.
+"""
+
+import abc
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedTokenizerBase
+
+
+class Encoder(nn.Module, metaclass=abc.ABCMeta):
+    """Base class for an encoder model."""
+
+    tokenizer: PreTrainedTokenizerBase
+    model: nn.Module
+
+    @property
+    @abc.abstractmethod
+    def output_units(self) -> int:
+        """
+        Hidden size of the encoder model.
+
+        Returns:
+            int: Hidden size of the encoder model.
+        """
+        pass
+
+    @property
+    @abc.abstractmethod
+    def max_positions(self) -> int:
+        """
+        Max number of tokens the encoder handles.
+
+        Returns:
+            int: Max number of tokens the encoder handles.
+        """
+        pass
+
+    @property
+    @abc.abstractmethod
+    def num_layers(self) -> int:
+        """
+        Number of model layers available.
+
+        Returns:
+            int: Number of model layers available.
+        """
+        pass
+
+    @property
+    @abc.abstractmethod
+    def size_separator(self) -> int:
+        """Size of the seperator.
+        E.g: For BERT is just 1 (`[SEP]`) but models such as XLM-R and InfoXLM use 2 (`</s></s>`).
+
+        Returns:
+            int: Number of tokens used between two segments.
+        """
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def from_pretrained(
+        cls,
+        pretrained_model: str,
+        load_pretrained_weights: bool = True,
+        local_files_only: bool = False,
+        revision: str | None = None,
+    ) -> "Encoder":
+        """
+        Function that loads a pretrained encoder and the respective tokenizer.
+
+        Args:
+            pretrained_model (str): Name of the pretrained model.
+            load_pretrained_weights (bool): Whether to load pretrained weights.
+            local_files_only (bool): Whether to only use locally cached files.
+            revision (str | None): Hugging Face revision to load.
+
+        Returns:
+            Encoder: Pretrained model from Hugging Face.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+        """
+        raise NotImplementedError
+
+    def freeze(self) -> None:
+        """
+        Freezes the entire encoder.
+        """
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def unfreeze(self) -> None:
+        """
+        Unfreezes the entire encoder.
+        """
+        for param in self.parameters():
+            param.requires_grad = True
+
+    @abc.abstractmethod
+    def freeze_embeddings(self) -> None:
+        """
+        Freezes the embedding layer.
+        """
+        pass
+
+    @abc.abstractmethod
+    def layerwise_lr(self, lr: float, decay: float) -> List:
+        """
+        Calculates the learning rate for each encoder layer by applying a small decay.
+
+        Args:
+            lr (float): Learning rate for the highest encoder layer.
+            decay (float): decay percentage for the lower layers.
+
+        Returns:
+            List: List of model parameters for all layers and the corresponding learning rate.
+        """
+        pass
+
+    @abc.abstractmethod
+    def forward(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
+    ) -> torch.Tensor:
+        """
+        Encoder model forward.
+
+        Args:
+            input_ids (torch.Tensor): Tokenized batch.
+            attention_mask (torch.Tensor): Batch attention mask.
+
+        Returns:
+            torch.Tensor: All layers embeddings.
+        """
+        pass
+
+    def prepare_sample(
+        self,
+        sample: List[str],
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Receives a list of strings and applies tokenization and vectorization.
+
+        Args:
+            sample (List[str]): List with text segments to be tokenized and padded.
+
+        Returns:
+            Dict[str, torch.Tensor]: dict with `'input_ids'` and `'attention_mask'`.
+        """
+        return {
+            k: v
+            for k, v in self.tokenizer(
+                sample,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=self.max_positions,
+            ).items()
+        }
+
+    def pad_tensor(
+        self, tensor: torch.Tensor, length: int, padding_index: int
+    ) -> torch.Tensor:
+        """
+        Pad a tensor to length with padding_index.
+
+        Args:
+            tensor (torch.Tensor): Tensor to pad.
+            length (int): Sequence length after padding.
+            padding_index (int): Index to pad tensor with.
+
+        Returns:
+            torch.Tensor: Input batch
+        """
+        n_padding = length - tensor.shape[0]
+        assert n_padding >= 0
+        if n_padding == 0:
+            return tensor
+        padding = tensor.new(n_padding, *tensor.shape[1:]).fill_(padding_index)
+        return torch.cat((tensor, padding), dim=0)
+
+    def concat_sequences(
+        self,
+        inputs: List[Dict[str, torch.Tensor]],
+        return_span_masks: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Concatenate K tokenized sequences per example into a single encoder input, distributing the length budget across segments before concatenation.
+        Optionally it returns span masks for SRC/MT0/MT1 content tokens (separators and pads are 0).
+        Works for:
+            - K=2 : SRC + TR        (Absolute QE)
+            - K=3 : SRC + TR1 + TR2 (Pairwise)
+
+        Args:
+            inputs (List[Dict[str, torch.Tensor]]): List of model inputs.
+            return_span_masks (bool): Whether to return span masks. Default: `False`.
+
+        Returns:
+            Dict[str, torch.Tensor]: Returns a single model input with all sentences concatenated into a single input.
+        """
+        if len(inputs) == 0:
+            raise ValueError("concat_sequences expects at least one input dict.")
+
+        # 1) Strip padding via attention masks → python lists of ids per segment/example
+        segments: List[List[List[int]]] = []
+        for encoder_input in inputs:
+            ids = encoder_input["input_ids"]  # [B, T]
+            am = encoder_input["attention_mask"].bool()
+            seg_i = [ids[b][am[b]].tolist() for b in range(ids.size(0))]
+            segments.append(seg_i)
+
+        K = len(segments)
+        B = len(segments[0])
+
+        if K != 2 and K != 3:
+            raise ValueError(
+                f"concat_sequences expects K=2 or K=3 segments, got K={K}."
+            )
+
+        # 2) Specials budget (depends on family via size_separator)
+        # BERT: [CLS] s1 [SEP] s2 [SEP] s3 [SEP] → specials = 1 + K
+        # RoBERTa/XLM-R/InfoXLM:
+        #   K>=2: <s> s1 </s></s> s2 </s></s> s3 </s> ... → K + 2 (two double boundaries)
+        if self.size_separator == 1:
+            specials = 1 + K
+        else:
+            specials = 2 if K == 1 else 2 + (K - 1) * self.size_separator
+
+        total_content_budget = max(self.max_positions - specials, 0)
+
+        # Ids
+        sep_id = self.tokenizer.sep_token_id
+        cls_id = self.tokenizer.cls_token_id
+        if cls_id is None or sep_id is None:
+            raise ValueError("Tokenizer must define cls_token_id and sep_token_id.")
+        assert isinstance(cls_id, int) and isinstance(sep_id, int), (
+            f"cls_token_id and sep_token_id must be integers, got "
+            f"{type(cls_id)}: {cls_id!r}, {type(sep_id)}: {sep_id!r}"
+        )
+
+        batch_tokens: List[torch.Tensor] = []
+        # When masks requested, we store per-example seg-ids (0=special/pad, 1=SRC, 2=MT0, 3=MT1)
+        segid_lists: List[List[int]] = [] if return_span_masks else None  # type: ignore
+
+        for b in range(B):
+            # content tokens per segment (strip outer specials)
+            contents = [segments[k][b][1:-1] for k in range(K)]
+            raw_lens = torch.tensor([len(c) for c in contents], dtype=torch.long)
+
+            # 3) Equal-split allocation with round-robin handout
+            q = total_content_budget // K
+            alloc = torch.full((K,), q, dtype=torch.long)
+            alloc = torch.minimum(alloc, raw_lens)  # respect per-segment caps
+            leftover = int(total_content_budget - int(alloc.sum().item()))
+            while leftover > 0:
+                elig = torch.nonzero(raw_lens > alloc, as_tuple=False).flatten()
+                if elig.numel() == 0:
+                    break
+                for t in elig.tolist():
+                    i = int(t)
+                    alloc[i] += 1
+                    leftover -= 1
+                    if leftover == 0:
+                        break
+
+            # truncate contents to allocated budgets
+            contents = [contents[k][: int(alloc[k].item())] for k in range(K)]
+
+            merged: List[int] = []
+            segids: List[int] = [] if return_span_masks else None  # type: ignore
+
+            if self.size_separator == 1:
+                # BERT: [CLS] s1 [SEP] s2 [SEP] (s3 [SEP])*
+                merged.append(cls_id)
+                if segids is not None:
+                    segids.append(0)
+                # s1 (SRC)
+                for t in contents[0]:
+                    merged.append(t)
+                    if segids is not None:
+                        segids.append(1)
+                merged.append(sep_id)
+                if segids is not None:
+                    segids.append(0)
+                # s2 (MT0)
+                for t in contents[1]:
+                    merged.append(t)
+                    if segids is not None:
+                        segids.append(2)
+                merged.append(sep_id)
+                if segids is not None:
+                    segids.append(0)
+                # s3 (MT1) if present
+                if K == 3:
+                    for t in contents[2]:
+                        merged.append(t)
+                        if segids is not None:
+                            segids.append(3)
+                    merged.append(sep_id)
+                    if segids is not None:
+                        segids.append(0)
+            else:
+                # RoBERTa/XLM-R/InfoXLM
+                merged.append(cls_id)
+                if segids is not None:
+                    segids.append(0)
+                # <s> s1 </s></s> s2 </s></s> [ + s3 </s>]
+                # s1 (SRC)
+                for t in contents[0]:
+                    merged.append(t)
+                    if segids is not None:
+                        segids.append(1)
+                # first double boundary
+                merged.append(sep_id)
+                merged.append(sep_id)
+                if segids is not None:
+                    segids.append(0)
+                    segids.append(0)
+                # s2 (MT0)
+                for t in contents[1]:
+                    merged.append(t)
+                    if segids is not None:
+                        segids.append(2)
+                # second double boundary
+                merged.append(sep_id)
+                merged.append(sep_id)
+                if segids is not None:
+                    segids.append(0)
+                    segids.append(0)
+                # s3.. (MT1 if present)
+                if K == 3:
+                    for t in contents[2]:
+                        merged.append(t)
+                        if segids is not None:
+                            segids.append(3)
+                    merged.append(sep_id)
+                    if segids is not None:
+                        segids.append(0)
+
+            assert (
+                len(merged) <= self.max_positions
+            ), f"Merged length {len(merged)} exceeds max_positions {self.max_positions}."
+
+            batch_tokens.append(torch.tensor(merged, dtype=torch.long))
+            if return_span_masks:
+                segid_lists.append(segids)  # type: ignore
+
+        # 4) Pad and build attention mask
+        lengths = [t.size(0) for t in batch_tokens]
+        max_len = max(lengths)
+        pad_raw = self.tokenizer.pad_token_id
+        if not isinstance(pad_raw, int):
+            raise ValueError(
+                f"Tokenizer must have integer pad_token_id; got {pad_raw!r}."
+            )
+        pad_id: int = pad_raw
+
+        padded = [self.pad_tensor(t, max_len, pad_id) for t in batch_tokens]
+        padded = torch.stack(padded, dim=0).contiguous()  # [B, L]
+        lens_t = torch.tensor(lengths, dtype=torch.long)
+        arange = torch.arange(max_len, device=padded.device)
+        attention_mask = (arange[None, :] < lens_t[:, None]).long()  # [B, L]
+
+        out = {"input_ids": padded, "attention_mask": attention_mask}
+
+        if return_span_masks:
+            # Build masks from segids (0=special/pad)
+            src_mask = torch.zeros((B, max_len), dtype=torch.long)
+            mt0_mask = torch.zeros((B, max_len), dtype=torch.long)
+            mt1_mask = torch.zeros((B, max_len), dtype=torch.long) if K == 3 else None
+            for b in range(B):
+                Lb = lengths[b]
+                seg = segid_lists[b]
+                # copy segids into [B,L], pad tail stays 0
+                for i in range(Lb):
+                    sid = seg[i]
+                    if sid == 1:
+                        src_mask[b, i] = 1
+                    elif sid == 2:
+                        mt0_mask[b, i] = 1
+                    elif sid == 3 and K == 3:
+                        mt1_mask[b, i] = 1  # type: ignore
+            out["src_mask"] = src_mask
+            out["mt0_mask"] = mt0_mask
+            if K == 3:
+                out["mt1_mask"] = mt1_mask  # type: ignore
+
+        return out