Skip to content

Make transform_tokens an explicit contract for pipeline steps (add alias to LogProbParser) #297

Description

@chicham

Problem

The token-mode prediction path introduced by #240 (BaseDetector.predict_token_proba) routes each non-final pipeline step through transform_tokens if it exists, falling back to transform via except AttributeError. The fallback exists only because LogProbParser (landed with the earlier parser work, #278-era) predates the token-mode design and never declared its token behavior.

Two failure modes of the implicit dispatch:

  • except AttributeError cannot distinguish "step has no transform_tokens" from "step's transform_tokens raised AttributeError internally" — the latter silently reroutes to transform and produces wrong-shaped data far downstream.
  • A future transformer that reduces over the token axis in transform (like EntropyTransformer) but forgets to implement transform_tokens gets silently sequence-reduced in the token path: a shape/semantics bug with no error.

Proposal: structural Protocol + beartype

Express the token-mode capability as a typed, structural contract instead of attribute introspection:

from typing import Protocol, runtime_checkable

import numpy as np
from beartype.door import is_bearable


@runtime_checkable
class TokenTransformer(Protocol):
    """A pipeline step that can emit per-token (non-reduced) features."""

    def transform_tokens(self, x: np.ndarray) -> np.ndarray: ...

and in BaseDetector.predict_token_proba:

for _, transformer in self.steps[:-1]:
    if is_bearable(transformer, TokenTransformer):
        raw_output = transformer.transform_tokens(raw_output)
    else:
        raw_output = transformer.transform(raw_output)

Benefits:

  • The capability is a named, documented type — the contract lives in the code, not in a comment; third-party transformers opt in structurally by just implementing the method.
  • No hasattr/getattr/try-except AttributeError introspection; an AttributeError raised inside an implementation propagates with its true traceback.
  • beartype is already a project dependency (used in entropy_methods), so no new requirement.

Complementary: give LogProbParser the one-line alias transform_tokens = transform (its (n, max_tokens, k) output is already per-token, the two modes coincide), so all first-party steps satisfy TokenTransformer and the else branch remains only as the documented extension point for sequence-only steps.

Scope note

Flagged during review of #291 but the root cause belongs to the parser PR, which landed before the token-mode design existed — hence a separate issue rather than a #291 change request.

Refs: #240, #277, #278, #291

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or requesttaskConcrete implementation issue, child of an epic

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions