Function tokenize_sft_batch

src/art/preprocessing/tokenize.py:366–432

Tokenize a single batch of trajectories for SFT. Args: trajectory_batch: List of trajectories in this batch learning_rate: Learning rate for this batch tokenizer: Tokenizer to use for encoding instruction_part: Instruction template part (e.g., "<|im_start|>user")

(
    trajectory_batch: list[Trajectory],
    learning_rate: float,
    tokenizer: PreTrainedTokenizerBase,
    instruction_part: str,
    response_part: str,
)

def tokenize_sft_batch(
    trajectory_batch: list[Trajectory],
    learning_rate: float,
    tokenizer: PreTrainedTokenizerBase,
    instruction_part: str,
    response_part: str,
) -> SFTBatch:
    """Tokenize a single batch of trajectories for SFT.

    Args:
        trajectory_batch: List of trajectories in this batch
        learning_rate: Learning rate for this batch
        tokenizer: Tokenizer to use for encoding
        instruction_part: Instruction template part (e.g., "<|im_start|>user")
        response_part: Response template part (e.g., "<|im_start|>assistant")

    Returns:
        SFTBatch object for this batch
    """
    import unsloth  # noqa: F401 - Must be imported first to set UNSLOTH_IS_PRESENT env var
    from unsloth_zoo.dataset_utils import train_on_responses_only

    train_on_responses_only_fn = train_on_responses_only(
        trainer=None,
        instruction_part=instruction_part,
        response_part=response_part,
        force_match=False,
        tokenizer=tokenizer,
        return_function=True,
    )
    # Tokenize all trajectories (no padding — each keeps its natural length)
    trajectory_tensors = []
    num_trainable_tokens = 0
    for trajectory in trajectory_batch:
        messages = trajectory.messages_and_choices
        tools = trajectory.tools

        # Single-step tokenization: apply_chat_template with tokenize=True
        input_ids = cast(
            list[int],
            tokenizer.apply_chat_template(
                cast(Any, messages),
                tools=cast(Any, tools),
                tokenize=True,
                add_generation_prompt=False,
            ),
        )

        attention_mask = [1] * len(input_ids)

        labels = train_on_responses_only_fn({"input_ids": [input_ids]})["labels"][0]

        trajectory_tensors.append(
            {
                "input_ids": torch.tensor([input_ids], dtype=torch.long),
                "attention_mask": torch.tensor([attention_mask], dtype=torch.long),
                "labels": torch.tensor([labels], dtype=torch.long),
            }