mirror of https://github.com/razor-ai/soup.git
114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
"""Data collators for sample packing (v0.33.0 #47).
|
|
|
|
When ``training.packing=true`` and ``training.packing_cross_doc_attn_mask=true``,
|
|
multiple short documents are packed into a single sequence; the default causal
|
|
mask leaks attention across document boundaries. ``CrossDocCollator`` builds an
|
|
explicit block-diagonal causal mask via :func:`utils.cross_doc_attn.build_cross_doc_mask`
|
|
and injects it as ``attention_mask`` on the batch, preferred over TRL's
|
|
``packing_strategy="attention_free"`` flag (which is best-effort across TRL
|
|
versions).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
|
|
class CrossDocCollator:
|
|
"""Wraps a base data collator and overrides ``attention_mask`` with a
|
|
block-diagonal causal mask derived from per-doc boundaries.
|
|
|
|
The base collator produces the standard ``input_ids`` / ``labels`` tensors;
|
|
this wrapper looks for a ``doc_lengths`` field in the underlying examples
|
|
(per-doc token counts after tokenization) and computes the block-diagonal
|
|
mask. If ``doc_lengths`` is missing, the original ``attention_mask`` is
|
|
preserved so the wrapper degrades gracefully.
|
|
|
|
Usage::
|
|
|
|
from transformers import DataCollatorForLanguageModeling
|
|
from soup_cli.data.collators import CrossDocCollator
|
|
|
|
base = DataCollatorForLanguageModeling(tokenizer, mlm=False)
|
|
collator = CrossDocCollator(base_collator=base)
|
|
trainer = SFTTrainer(..., data_collator=collator)
|
|
"""
|
|
|
|
def __init__(
|
|
self, base_collator: Any, doc_lengths_key: str = "doc_lengths",
|
|
) -> None:
|
|
if base_collator is None:
|
|
raise ValueError("CrossDocCollator requires a base_collator")
|
|
self._base = base_collator
|
|
self._key = doc_lengths_key
|
|
|
|
def __call__(self, features: list[dict]) -> dict:
|
|
# Pull the per-example doc_lengths out before delegating, since the
|
|
# base collator may strip unknown keys.
|
|
per_example_lengths: list[Optional[list[int]]] = []
|
|
cleaned: list[dict] = []
|
|
for example in features:
|
|
lengths = example.pop(self._key, None) if isinstance(example, dict) else None
|
|
per_example_lengths.append(lengths)
|
|
cleaned.append(example)
|
|
|
|
batch = self._base(cleaned)
|
|
|
|
# Try to build the cross-doc mask. If anything goes wrong (no
|
|
# doc_lengths, mismatched shapes, no numpy/torch), preserve the
|
|
# base ``attention_mask`` and continue.
|
|
try:
|
|
self._inject_block_diag_mask(batch, per_example_lengths)
|
|
except Exception: # noqa: BLE001 — degrade rather than crash training
|
|
pass
|
|
return batch
|
|
|
|
def _inject_block_diag_mask(
|
|
self, batch: dict, per_example_lengths: list[Optional[list[int]]],
|
|
) -> None:
|
|
from soup_cli.utils.cross_doc_attn import (
|
|
build_cross_doc_mask,
|
|
compute_doc_boundaries,
|
|
)
|
|
|
|
if "input_ids" not in batch:
|
|
return
|
|
input_ids = batch["input_ids"]
|
|
if not hasattr(input_ids, "shape") or len(input_ids.shape) != 2:
|
|
return
|
|
batch_size, seq_length = input_ids.shape
|
|
|
|
if not any(lengths for lengths in per_example_lengths):
|
|
return # Nothing to do — no doc_lengths supplied
|
|
|
|
try:
|
|
import numpy as np
|
|
import torch
|
|
except ImportError:
|
|
return
|
|
|
|
masks = []
|
|
for lengths in per_example_lengths:
|
|
if lengths and sum(lengths) <= seq_length:
|
|
# Pad final segment if doc lengths sum to < seq_length
|
|
total = sum(lengths)
|
|
padded_lengths = list(lengths)
|
|
if total < seq_length:
|
|
padded_lengths.append(seq_length - total)
|
|
boundaries = compute_doc_boundaries(padded_lengths)
|
|
mask = build_cross_doc_mask(boundaries, seq_length)
|
|
else:
|
|
# Fallback: lower-triangular causal (no doc separation).
|
|
mask = np.tril(np.ones((seq_length, seq_length), dtype=np.uint8))
|
|
masks.append(mask)
|
|
|
|
if len(masks) != batch_size:
|
|
return
|
|
|
|
attn_tensor = torch.from_numpy(np.stack(masks))
|
|
# HF expects ``attention_mask`` shape (batch, seq) for standard
|
|
# masking; the (batch, seq, seq) block-diag mask is consumed by
|
|
# FlashAttn-2 / SDPA when passed as ``attn_mask`` instead. We expose
|
|
# both names so downstream code can pick whichever it needs.
|
|
batch["cross_doc_attn_mask"] = attn_tensor
|