Selecting Tokens and Samples

Picking out token positions and sample subsets using INIF’s selectors.

INIF’s selectors are intentionally tiny — methods on Sample / InifDocument that return a TokenSelection (or list of samples) without mutating the input. Compose them rather than building a query DSL.

TokenSelection

Every per-sample selector returns a TokenSelection:

from inif import TokenSelection

@dataclass
class TokenSelection:
    sample_id: str | int
    tokens: list[TokenOrSeqRef]
    positions: list[int]

positions always reflects the actual sample-local indices (sorted, deduplicated). tokens is the parallel list of TokenOrSeqRef objects pulled from those positions.

Token selectors (on Sample)

By position

sample.select_by_position(5)               # one position
sample.select_by_position([3, 5, 7])       # explicit list
sample.select_by_position(slice(0, 10))    # slice
sample.select_by_position(slice(-5, None)) # last 5

Slices are resolved against len(sample.tokens) via slice.indices(n), so negative starts and step-2 slices behave like Python’s normal slicing. Out of range integers are silently dropped (they just don’t appear in tokens).

By annotation

selection = sample.select_by_annotation("generated")

Pulls every position covered by any annotation whose name matches — positions are flattened across all the ranges of every matching annotation, so it works whether the annotation has one range or twenty.

By sequence id

selection = sample.select_by_sequence_id("seq_0")

Returns the sequence-ref tokens that point at seq_0. Useful for finding where a shared run is referenced in a sample without expanding it.

By span

selection = sample.select_by_span("answer_span")

Unions the positions across every Span whose name matches.

Sample selectors (on InifDocument)

By score

failures = doc.filter_samples_by_score(
    scorer="exact_match", predicate=lambda v: v == 0
)
correct = doc.filter_samples_by_score(
    scorer="exact_match", predicate=lambda v: v == 1
)

Returns a list[Sample] (not a new document). The predicate sees the raw SampleScore.value, which can be str | int | float | bool | list | dict depending on the scorer.

By arbitrary predicate

InifDocument.filter_samples(predicate) returns the matching samples:

long_samples = doc.filter_samples(lambda s: len(s.tokens) > 1000)

For a self-contained sub-document — pruned sequences, deep-copied content — use subset instead:

mini_doc = doc.subset(lambda s: s.id in {"sample_3", "sample_7"})

subset rebuilds the sequence list to contain only sequences referenced by the kept samples, so the resulting document is independent and minimal.

Composing selectors

Methods return plain lists / TokenSelections — chain them however you want:

failures = doc.filter_samples_by_score(
    scorer="exact_match", predicate=lambda v: v == 0
)
generated_tokens = [s.select_by_annotation("generated") for s in failures]

Or pull positions across the whole document and feed them to nnsight / nnterp:

all_positions = []
for sample in doc.samples:
    sel = sample.select_by_annotation("user")
    for pos in sel.positions:
        all_positions.append((sample.id, pos))

What about flat-array scans?

When you need to scan tokens across samples — counting matches over an entire document, computing per-token statistics, looking for regex hits at the corpus level — the Sample-by-Sample API gets slow because every loop reconstructs Python objects. For those workloads, see FlatTokenStore, which exposes flat parallel arrays of token ids, strings, and sequence provenance, plus regex helpers that operate directly on those arrays.

FlatTokenStore is an analysis representation, not a replacement file format — call FlatTokenStore.from_document(doc) to materialise one.