pyt-crf 0.1.1__tar.gz

pyt_crf-0.1.1/PKG-INFO

@@ -0,0 +1,84 @@
+Metadata-Version: 2.3
+Name: pyt-crf
+Version: 0.1.1
+Summary: A modern linear-chain CRF (Conditional Random Field) layer for PyTorch
+Author: id4thomas
+Author-email: id4thomas <id2thomas@gmail.com>
+Requires-Dist: torch>=2.10
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# pyt-crf
+
+A modern linear-chain CRF (Conditional Random Field) layer for PyTorch.
+
+- `forward()` returns the **negative log-likelihood** directly (ready for `.backward()`)
+- `decode()` runs batched Viterbi decoding
+- Batch-first (`batch, seq_len, num_tags`) by default
+
+## Requirements
+
+- Python >= 3.12
+- PyTorch >= 2.10
+
+## Build
+
+```bash
+conda activate py312
+
+# build sdist + wheel into dist/ (uses the uv_build backend)
+uv build
+
+# or, without uv:
+pip install build
+python -m build
+```
+
+## Install
+
+```bash
+# from the built wheel
+pip install dist/pyt_crf-0.1.0-py3-none-any.whl
+
+# or editable, for development
+pip install -e .
+```
+
+## Quick usage
+
+```python
+import torch
+from pyt_crf import CRF
+
+crf = CRF(num_tags=5)
+emissions = torch.randn(2, 7, 5)          # (batch, seq_len, num_tags)
+tags = torch.randint(5, (2, 7))           # (batch, seq_len)
+mask = torch.ones(2, 7, dtype=torch.bool) # True = real token, False = padding
+
+loss = crf(emissions, tags, mask)  # scalar NLL loss
+loss.backward()
+
+best_paths = crf.decode(emissions, mask)  # List[List[int]]
+```
+
+> [!WARNING]
+> **Run the CRF in float32.** The forward/Viterbi recursions accumulate
+> `logsumexp`/max scores over the whole sequence, which is numerically unstable
+> in bf16/fp16. Keep the CRF parameters in fp32 and cast emissions before the
+> layer — `crf(emissions.float(), ...)` — even when the backbone runs in
+> bf16/autocast.
+
+## Examples
+
+- [`examples/conll2003`](examples/conll2003) — NER fine-tuning on
+  [tner/conll2003](https://huggingface.co/datasets/tner/conll2003) with a
+  Gemma3-based backbone + CRF head.
+
+## Changelog
+
+- [0.1.0](docs/changelogs/0_1_0.md) — initial release; full list of differences from pytorch-crf.
+
+## Acknowledgements
+
+- This package mostly follows implementation of [kmkurn/pytorch-crf](https://github.com/kmkurn/pytorch-crf).
+

pyt_crf-0.1.1/README.md

@@ -0,0 +1,74 @@
+# pyt-crf
+
+A modern linear-chain CRF (Conditional Random Field) layer for PyTorch.
+
+- `forward()` returns the **negative log-likelihood** directly (ready for `.backward()`)
+- `decode()` runs batched Viterbi decoding
+- Batch-first (`batch, seq_len, num_tags`) by default
+
+## Requirements
+
+- Python >= 3.12
+- PyTorch >= 2.10
+
+## Build
+
+```bash
+conda activate py312
+
+# build sdist + wheel into dist/ (uses the uv_build backend)
+uv build
+
+# or, without uv:
+pip install build
+python -m build
+```
+
+## Install
+
+```bash
+# from the built wheel
+pip install dist/pyt_crf-0.1.0-py3-none-any.whl
+
+# or editable, for development
+pip install -e .
+```
+
+## Quick usage
+
+```python
+import torch
+from pyt_crf import CRF
+
+crf = CRF(num_tags=5)
+emissions = torch.randn(2, 7, 5)          # (batch, seq_len, num_tags)
+tags = torch.randint(5, (2, 7))           # (batch, seq_len)
+mask = torch.ones(2, 7, dtype=torch.bool) # True = real token, False = padding
+
+loss = crf(emissions, tags, mask)  # scalar NLL loss
+loss.backward()
+
+best_paths = crf.decode(emissions, mask)  # List[List[int]]
+```
+
+> [!WARNING]
+> **Run the CRF in float32.** The forward/Viterbi recursions accumulate
+> `logsumexp`/max scores over the whole sequence, which is numerically unstable
+> in bf16/fp16. Keep the CRF parameters in fp32 and cast emissions before the
+> layer — `crf(emissions.float(), ...)` — even when the backbone runs in
+> bf16/autocast.
+
+## Examples
+
+- [`examples/conll2003`](examples/conll2003) — NER fine-tuning on
+  [tner/conll2003](https://huggingface.co/datasets/tner/conll2003) with a
+  Gemma3-based backbone + CRF head.
+
+## Changelog
+
+- [0.1.0](docs/changelogs/0_1_0.md) — initial release; full list of differences from pytorch-crf.
+
+## Acknowledgements
+
+- This package mostly follows implementation of [kmkurn/pytorch-crf](https://github.com/kmkurn/pytorch-crf).
+

pyt_crf-0.1.1/pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "pyt-crf"
+version = "0.1.1"
+description = "A modern linear-chain CRF (Conditional Random Field) layer for PyTorch"
+readme = "README.md"
+authors = [
+    { name = "id4thomas", email = "id2thomas@gmail.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "torch>=2.10",
+]
+
+[build-system]
+requires = ["uv_build>=0.10.12,<0.11.0"]
+build-backend = "uv_build"

pyt_crf-0.1.1/src/pyt_crf/__init__.py

@@ -0,0 +1,6 @@
+"""pyt-crf: a modern linear-chain CRF layer for PyTorch."""
+
+from pyt_crf.crf import CRF
+
+__version__ = "0.1.0"
+__all__ = ["CRF", "__version__"]

pyt_crf-0.1.1/src/pyt_crf/crf.py

@@ -0,0 +1,328 @@
+"""Linear-chain Conditional Random Field (CRF) layer for PyTorch.
+
+This module implements a first-order linear-chain CRF commonly stacked on top
+of sequence encoders (BiLSTM, Transformer, ...) for tagging tasks such as NER
+and POS tagging.
+
+Design notes (how this differs from the classic ``pytorch-crf`` package):
+
+* Batch-first ``(batch, seq_len, num_tags)`` is the default and the internal
+  canonical layout, matching modern PyTorch conventions.
+* ``forward`` returns the **negative** log-likelihood, i.e. a loss you can
+  backpropagate directly without negating it yourself.
+* The gold-path score is computed fully vectorized (no Python loop over time).
+* The Viterbi backtrace is vectorized over the batch dimension instead of
+  looping over each sequence in Python.
+"""
+
+from typing import List, Optional
+
+import torch
+from torch import Tensor, nn
+
+__all__ = ["CRF"]
+
+# Valid values for the `reduction` argument of CRF.forward.
+_REDUCTIONS = ("none", "sum", "mean", "token_mean")
+
+
+class CRF(nn.Module):
+    """Linear-chain CRF layer.
+
+    The layer learns three sets of scores:
+
+    * ``start_transitions[j]``  -- score for a sequence that begins with tag ``j``.
+    * ``end_transitions[j]``    -- score for a sequence that ends with tag ``j``.
+    * ``transitions[i, j]``     -- score for moving from tag ``i`` to tag ``j``.
+
+    Together with the emission scores produced by an upstream encoder, these
+    define a globally normalized distribution over whole tag sequences.
+
+    Args:
+        num_tags: Number of distinct tags. Must be positive.
+        batch_first: If ``True`` (default), inputs are ``(batch, seq_len, ...)``;
+            otherwise ``(seq_len, batch, ...)``.
+
+    Example:
+        >>> crf = CRF(num_tags=5)
+        >>> emissions = torch.randn(2, 7, 5)   # (batch, seq_len, num_tags)
+        >>> tags = torch.randint(5, (2, 7))    # (batch, seq_len)
+        >>> loss = crf(emissions, tags)        # scalar loss, ready for .backward()
+        >>> best_paths = crf.decode(emissions) # List[List[int]] of best tag ids
+    """
+
+    def __init__(self, num_tags: int, batch_first: bool = True) -> None:
+        # A CRF over zero tags is meaningless; fail fast with a clear message.
+        if num_tags <= 0:
+            raise ValueError(f"num_tags must be positive, got {num_tags}")
+        super().__init__()
+        # Remember the tag-set size for validation and for __repr__.
+        self.num_tags = num_tags
+        # Remember the expected input layout; internally we always work batch-first.
+        self.batch_first = batch_first
+        # Score of starting a sequence with each tag: shape (num_tags,).
+        self.start_transitions = nn.Parameter(torch.empty(num_tags))
+        # Score of ending a sequence with each tag: shape (num_tags,).
+        self.end_transitions = nn.Parameter(torch.empty(num_tags))
+        # transitions[i, j] = score of moving from tag i to tag j: (num_tags, num_tags).
+        self.transitions = nn.Parameter(torch.empty(num_tags, num_tags))
+        # Fill all three parameter tensors with small random values.
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """(Re-)initialize all transition scores uniformly in [-0.1, 0.1]."""
+        # Small symmetric init keeps early training close to "emissions only".
+        for param in (self.start_transitions, self.end_transitions, self.transitions):
+            nn.init.uniform_(param, -0.1, 0.1)
+
+    def extra_repr(self) -> str:
+        # Shown inside the standard nn.Module repr, e.g. CRF(num_tags=5, batch_first=True).
+        return f"num_tags={self.num_tags}, batch_first={self.batch_first}"
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+    def forward(
+        self,
+        emissions: Tensor,
+        tags: Tensor,
+        mask: Optional[Tensor] = None,
+        reduction: str = "mean",
+    ) -> Tensor:
+        """Compute the negative log-likelihood of ``tags`` given ``emissions``.
+
+        Args:
+            emissions: Emission score tensor of shape ``(batch, seq_len, num_tags)``
+                (or ``(seq_len, batch, num_tags)`` when ``batch_first=False``).
+            tags: Gold tag indices, shape ``(batch, seq_len)``.
+            mask: Optional boolean tensor of the same shape as ``tags``; ``True``
+                marks real tokens, ``False`` marks padding. The first timestep of
+                every sequence must be ``True``. Defaults to all-``True``.
+                ``False`` positions may also appear mid-sequence (e.g. sub-word
+                continuations labeled ``-100``): they are skipped entirely, with
+                a single transition bridging each gap.
+            reduction: One of ``"none"`` (per-sequence losses), ``"sum"``,
+                ``"mean"`` (default, mean over sequences), or ``"token_mean"``
+                (sum divided by the number of unmasked tokens).
+
+        Returns:
+            The NLL loss: a scalar unless ``reduction="none"``, in which case a
+            ``(batch,)`` tensor.
+        """
+        # Reject typos early instead of silently falling through.
+        if reduction not in _REDUCTIONS:
+            raise ValueError(f"reduction must be one of {_REDUCTIONS}, got {reduction!r}")
+        # Normalize layout/dtype and run all shape checks in one place.
+        emissions, tags, mask = self._canonicalize(emissions, tags, mask)
+        # Unnormalized log-score of the gold tag path: shape (batch,).
+        gold_score = self._joint_score(emissions, tags, mask)
+        # log of the partition function Z (sum over ALL tag paths): shape (batch,).
+        log_z = self._log_partition(emissions, mask)
+        # NLL = -log p(tags | emissions) = log Z - score(gold path).
+        nll = log_z - gold_score
+        # Apply the requested reduction over the batch.
+        if reduction == "none":
+            return nll
+        if reduction == "sum":
+            return nll.sum()
+        if reduction == "mean":
+            return nll.mean()
+        # "token_mean": normalize by the total number of real (unmasked) tokens.
+        return nll.sum() / mask.sum()
+
+    @torch.no_grad()
+    def decode(
+        self,
+        emissions: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> List[List[int]]:
+        """Find the highest-scoring tag sequence for each batch element (Viterbi).
+
+        Args:
+            emissions: Emission scores, same layout as in :meth:`forward`.
+            mask: Optional boolean padding mask, same semantics as in :meth:`forward`.
+
+        Returns:
+            A list of tag-id lists; entry ``b`` covers exactly the timesteps of
+            sequence ``b`` where ``mask`` is ``True``.
+        """
+        # Reuse the same canonicalization path; `tags=None` skips tag checks.
+        emissions, _, mask = self._canonicalize(emissions, None, mask)
+        # Run batched Viterbi and unpad the results into plain Python lists.
+        return self._viterbi(emissions, mask)
+
+    # ------------------------------------------------------------------
+    # Input handling
+    # ------------------------------------------------------------------
+    def _canonicalize(
+        self,
+        emissions: Tensor,
+        tags: Optional[Tensor],
+        mask: Optional[Tensor],
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        """Validate inputs and return them in batch-first layout with a bool mask."""
+        # Emissions must be a 3-D tensor: two sequence dims plus the tag dim.
+        if emissions.dim() != 3:
+            raise ValueError(f"emissions must be 3-dimensional, got {emissions.dim()}")
+        # The size of the tag dimension must match this layer's tag set.
+        if emissions.size(2) != self.num_tags:
+            raise ValueError(
+                f"expected {self.num_tags} tags in the last dimension of emissions, "
+                f"got {emissions.size(2)}"
+            )
+        # Bring everything into (batch, seq_len, ...) order if it isn't already.
+        if not self.batch_first:
+            emissions = emissions.transpose(0, 1)
+            if tags is not None:
+                tags = tags.transpose(0, 1)
+            if mask is not None:
+                mask = mask.transpose(0, 1)
+        # Tags, when given, must line up with emissions on (batch, seq_len).
+        if tags is not None and tags.shape != emissions.shape[:2]:
+            raise ValueError(
+                f"tags shape {tuple(tags.shape)} does not match emissions "
+                f"batch/sequence dims {tuple(emissions.shape[:2])}"
+            )
+        if mask is None:
+            # No mask means every position is a real token.
+            mask = emissions.new_ones(emissions.shape[:2], dtype=torch.bool)
+        else:
+            # Accept 0/1 float or int masks too by casting to bool.
+            mask = mask.to(dtype=torch.bool)
+            # The mask must also line up with emissions on (batch, seq_len).
+            if mask.shape != emissions.shape[:2]:
+                raise ValueError(
+                    f"mask shape {tuple(mask.shape)} does not match emissions "
+                    f"batch/sequence dims {tuple(emissions.shape[:2])}"
+                )
+            # Every sequence must contain at least its first token; the recursions
+            # below are seeded from timestep 0 and cannot represent empty sequences.
+            if not bool(mask[:, 0].all()):
+                raise ValueError("mask must be True at the first timestep of every sequence")
+        # When tags are absent (decode path) return a dummy to keep the tuple shape.
+        if tags is None:
+            tags = emissions.new_zeros(emissions.shape[:2], dtype=torch.long)
+        return emissions, tags, mask
+
+    @staticmethod
+    def _last_valid_index(mask: Tensor) -> Tensor:
+        """Return the index of the last ``True`` timestep for each sequence.
+
+        Computed as an argmax over actual positions rather than ``sum(mask) - 1``,
+        so it stays correct even if a mask has ``False`` holes in the middle.
+        """
+        # positions[b, t] = t for every batch element: shape (batch, seq_len).
+        positions = torch.arange(mask.size(1), device=mask.device).expand_as(mask)
+        # Zero out masked positions, then take the max position that survived.
+        # (Position 0 is always valid, so the zero fallback is never wrong.)
+        return torch.where(mask, positions, positions.new_zeros(())).amax(dim=1)
+
+    # ------------------------------------------------------------------
+    # Score computation
+    # ------------------------------------------------------------------
+    def _joint_score(self, emissions: Tensor, tags: Tensor, mask: Tensor) -> Tensor:
+        """Score of the gold tag path, fully vectorized. Returns shape ``(batch,)``.
+
+        The path is defined over the *unmasked* timesteps only:
+
+        score = start(first tag) + sum emit(t, tags[t]) over unmasked t
+                + sum trans(previous unmasked tag, tags[t]) + end(last tag)
+
+        Transitions bridge over masked holes (matching the forward algorithm,
+        which freezes its recursion at masked steps), so tags at masked
+        positions never contribute — not even as transition sources.
+        """
+        # Cast the mask once; used to zero out contributions from padding.
+        mask_f = mask.to(emissions.dtype)
+        # Emission score of the gold tag at every timestep: gather along the tag
+        # dimension, giving shape (batch, seq_len); masked steps are zeroed.
+        emit_scores = emissions.gather(2, tags.unsqueeze(2)).squeeze(2) * mask_f
+        # For every position t, find the most recent unmasked position <= t:
+        # masked positions carry -1, and a running max propagates the latest
+        # unmasked index forward. Shape (batch, seq_len).
+        positions = torch.arange(mask.size(1), device=mask.device).expand_as(mask)
+        last_seen = torch.where(mask, positions, positions.new_full((), -1)).cummax(dim=1).values
+        # Shift right by one step to get the unmasked position strictly BEFORE t
+        # (-1 when there is none, i.e. at t = 0).
+        prev_pos = torch.cat(
+            [last_seen.new_full((mask.size(0), 1), -1), last_seen[:, :-1]], dim=1
+        )
+        # A transition is scored for every unmasked step that has an unmasked
+        # predecessor: (previous unmasked tag) -> (current tag).
+        pair_ok = (mask & (prev_pos >= 0)).to(emissions.dtype)
+        prev_tags = tags.gather(1, prev_pos.clamp(min=0))
+        trans_scores = self.transitions[prev_tags, tags] * pair_ok
+        # Score for starting with the first gold tag (step 0 is always unmasked).
+        start_scores = self.start_transitions[tags[:, 0]]
+        # Find the gold tag at each sequence's last real timestep...
+        last_tags = tags.gather(1, self._last_valid_index(mask).unsqueeze(1)).squeeze(1)
+        # ...and add the score for ending the sequence with that tag.
+        end_scores = self.end_transitions[last_tags]
+        # Total path score: sum every component over time, per batch element.
+        return start_scores + emit_scores.sum(dim=1) + trans_scores.sum(dim=1) + end_scores
+
+    def _log_partition(self, emissions: Tensor, mask: Tensor) -> Tensor:
+        """log Z via the forward algorithm in log space. Returns shape ``(batch,)``."""
+        seq_len = emissions.size(1)
+        # alpha[b, j] = logsumexp of scores of all partial paths that end in tag j
+        # at the current timestep. Seed with the start scores + first emissions.
+        alpha = self.start_transitions + emissions[:, 0]  # (batch, num_tags)
+        # Walk the sequence left to right, updating alpha one step at a time.
+        for t in range(1, seq_len):
+            # Combine every (previous tag i -> current tag j) hypothesis:
+            #   alpha as (batch, i, 1) + transitions (i, j) + emissions (batch, 1, j)
+            step_scores = (
+                alpha.unsqueeze(2) + self.transitions + emissions[:, t].unsqueeze(1)
+            )
+            # Marginalize over the previous tag i in log space: (batch, num_tags).
+            next_alpha = torch.logsumexp(step_scores, dim=1)
+            # Sequences already past their end keep their old alpha (frozen).
+            alpha = torch.where(mask[:, t].unsqueeze(1), next_alpha, alpha)
+        # Close every path with its end-transition score, then sum over final tags.
+        return torch.logsumexp(alpha + self.end_transitions, dim=1)
+
+    # ------------------------------------------------------------------
+    # Viterbi decoding
+    # ------------------------------------------------------------------
+    def _viterbi(self, emissions: Tensor, mask: Tensor) -> List[List[int]]:
+        """Batched Viterbi: max-product forward pass + vectorized backtrace."""
+        batch_size, seq_len, _ = emissions.shape
+        # score[b, j] = best score of any partial path ending in tag j right now.
+        score = self.start_transitions + emissions[:, 0]  # (batch, num_tags)
+        # backpointers[t-1][b, j] = best previous tag when arriving at tag j at step t.
+        backpointers: List[Tensor] = []
+        # Forward pass: identical shape dance to _log_partition, but with max.
+        for t in range(1, seq_len):
+            # All (previous tag i -> current tag j) scores: (batch, i, j).
+            step_scores = (
+                score.unsqueeze(2) + self.transitions + emissions[:, t].unsqueeze(1)
+            )
+            # Keep only the best previous tag for each current tag.
+            best_scores, best_prev = step_scores.max(dim=1)  # both (batch, num_tags)
+            # Frozen (padded) sequences carry their old scores forward.
+            score = torch.where(mask[:, t].unsqueeze(1), best_scores, score)
+            # Record where each tag's best score came from, for the backtrace.
+            backpointers.append(best_prev)
+        # Add the cost of ending on each tag before picking the winner.
+        score = score + self.end_transitions
+        # Best final tag per sequence: shape (batch,).
+        current = score.argmax(dim=1)
+        # Where each sequence actually ends (last unmasked timestep).
+        last_idx = self._last_valid_index(mask)
+        # Buffer for the decoded path of every sequence: (batch, seq_len).
+        paths = emissions.new_zeros((batch_size, seq_len), dtype=torch.long)
+        # The final position of every path holds the best final tag.
+        paths[:, seq_len - 1] = current
+        # Backtrace right to left, for the whole batch at once.
+        for t in range(seq_len - 2, -1, -1):
+            # Look up the predecessor of the current tag: backpointers[t] maps
+            # tags at step t+1 back to tags at step t. Shape (batch,).
+            prev = backpointers[t].gather(1, current.unsqueeze(1)).squeeze(1)
+            # Only step back where step t+1 was a real, in-range token; otherwise
+            # (padding or a masked hole) the chain just carries the current tag.
+            follow = mask[:, t + 1] & (t + 1 <= last_idx)
+            current = torch.where(follow, prev, current)
+            # Record the tag chosen for timestep t.
+            paths[:, t] = current
+        # Strip padding (and masked holes): keep only timesteps where mask is True.
+        return [paths[b, mask[b]].tolist() for b in range(batch_size)]