diff --git a/areno/api/data_utils.py b/areno/api/data_utils.py index dd93d521..b2d80b92 100644 --- a/areno/api/data_utils.py +++ b/areno/api/data_utils.py @@ -51,6 +51,110 @@ def response_to_tokens_and_mask( return prompt_ids + response_ids, [True] * len(prompt_ids) + [False] * len(response_ids) +def _try_chat_template_encoding( + messages: list[dict[str, Any]], tokenizer, trainable_assistant_indices: set[int] +) -> tuple[list[int], list[bool]] | None: + """Attempt incremental chat-template encoding; return None if not prefix-stable. + + Tries to encode the conversation turn by turn using the tokenizer's + ``chat_template``. If the re-encoded prefix ever differs from what was + already accumulated, the tokenizer is not prefix-stable and we return + ``None`` so the caller can fall back to plain-text concatenation. + """ + + if not getattr(tokenizer, "chat_template", None): + return None + + tokens: list[int] = [] + mask: list[bool] = [] + for i in range(len(messages)): + partial_ids = normalize_token_ids( + apply_chat_template_with_options( + tokenizer, messages[: i + 1], tokenize=True, add_generation_prompt=False + ) + ) + # Guard against tokenizers whose chat_template is not prefix-stable. + if tokens and partial_ids[: len(tokens)] != tokens: + import warnings + + warnings.warn( + "tokenizer chat_template is not prefix-stable; " + "falling back to plain-text encoding for multi-turn SFT", + RuntimeWarning, + stacklevel=3, + ) + return None + # Only keep tokens added by the current turn. + new_tokens = partial_ids[len(tokens):] + role = messages[i].get("role", "user") + is_trainable = role == "assistant" and i in trainable_assistant_indices + tokens.extend(new_tokens) + mask.extend([not is_trainable] * len(new_tokens)) + return tokens, mask + + +def messages_to_tokens_and_mask( + messages: list[dict[str, Any]], + tokenizer, + eos_token_id: int, + *, + last_assistant_only: bool = False, +) -> tuple[list[int], list[bool]]: + """Encode multi-turn chat messages into tokens with a training mask. + + The mask follows the same convention as + :func:`prompt_response_to_tokens_and_mask`: ``True`` means "do not train" + (prompt context), ``False`` means "train" (assistant response). + + * user / system / tool turns are always masked out (``True``). + * assistant turns are trainable (``False``) unless *last_assistant_only* + is set, in which case only the final assistant turn is trainable and + earlier assistant turns are treated as context (``True``). + + The function uses the tokenizer chat template when available so turn + markers and special tokens match the model's expected format. For base + tokenizers without a chat template, a plain-text fallback concatenates + ``role: content`` per turn. + + EOS is appended after the last message if not already present, so the + model learns to stop. + """ + + # Determine which assistant turns are trainable. + assistant_indices = [ + i for i, msg in enumerate(messages) if msg.get("role") == "assistant" + ] + if last_assistant_only and assistant_indices: + trainable_assistant_indices = {assistant_indices[-1]} + else: + trainable_assistant_indices = set(assistant_indices) + + chat_template_tokens = _try_chat_template_encoding( + messages, tokenizer, trainable_assistant_indices + ) + if chat_template_tokens is not None: + tokens, mask = chat_template_tokens + else: + # Plain-text fallback: concatenate "role: content" per turn. + tokens = [] + mask = [] + for i, msg in enumerate(messages): + role = msg.get("role", "user") + content = msg.get("content", "") + turn_text = f"{role}: {content}" + turn_ids = normalize_token_ids(tokenizer.encode(turn_text, add_special_tokens=False)) + is_trainable = role == "assistant" and i in trainable_assistant_indices + tokens.extend(turn_ids) + mask.extend([not is_trainable] * len(turn_ids)) + + # Append EOS if not already present so the model learns to stop. + if eos_token_id is not None and (not tokens or tokens[-1] != eos_token_id): + tokens.append(eos_token_id) + mask.append(False) + + return tokens, mask + + def has_any(record: dict[str, Any], keys: tuple[str, ...]) -> bool: """Return whether a record has any string field in keys.""" diff --git a/areno/api/trainer_config.py b/areno/api/trainer_config.py index 52457f84..bc612deb 100644 --- a/areno/api/trainer_config.py +++ b/areno/api/trainer_config.py @@ -79,6 +79,7 @@ class TrainerConfig: agent_timeout_s: float = 300.0 train_tool_results: bool = False chat_template_enable_thinking: bool | None = None + sft_assistant_turns: str = "all" lora: LoraConfig | None = None reference_mode: Literal["independent", "reuse_actor_base"] = "independent" @@ -95,6 +96,8 @@ def __post_init__(self) -> None: raise ValueError("attn_backend must be one of: flash, native") if self.model_hub not in {"hf", "modelscope"}: raise ValueError("model_hub must be one of: hf, modelscope") + if self.sft_assistant_turns not in {"all", "last"}: + raise ValueError("sft_assistant_turns must be one of: all, last") if isinstance(self.optimizer_state_offload, bool): self.optimizer_state_offload = "cpu" if self.optimizer_state_offload else "none" if self.optimizer_state_offload not in {"none", "cpu", "disk"}: diff --git a/areno/api/trainers/sft.py b/areno/api/trainers/sft.py index 8774b297..c42224c6 100644 --- a/areno/api/trainers/sft.py +++ b/areno/api/trainers/sft.py @@ -23,7 +23,7 @@ import areno.api from areno.api.dashboard import record_dashboard_state -from areno.api.data_utils import prompt_response_to_tokens_and_mask +from areno.api.data_utils import messages_to_tokens_and_mask, prompt_response_to_tokens_and_mask from areno.api.multimodal import ( encode_multimodal_prompt, expand_image_tokens, @@ -72,6 +72,7 @@ def _fit_initialized(self) -> None: processor, max_prompt_tokens=self.config.max_prompt_tokens, max_new_tokens=self.config.max_new_tokens, + sft_assistant_turns=getattr(self.config, "sft_assistant_turns", "all"), ): if not train_batch: continue @@ -104,7 +105,7 @@ def _fit_initialized(self) -> None: self.logger.info("epoch=%d stage=epoch_end", epoch) record_dashboard_state(self.areno, stage="epoch_end", epoch=epoch, step=step, role="policy") - def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, max_new_tokens: int): + def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, max_new_tokens: int, sft_assistant_turns: str = "all"): # Dataset rows are converted lazily so large HF datasets do not need an # up-front tokenized copy. Rows that are empty, all-prompt, or exceed # the configured prompt or supervised-response budgets are dropped. @@ -120,6 +121,7 @@ def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, m processor, max_prompt_tokens=max_prompt_tokens, max_new_tokens=max_new_tokens, + sft_assistant_turns=sft_assistant_turns, ) if seq is None: skipped += 1 @@ -152,17 +154,31 @@ def _maybe_save(self, epoch: int, step: int) -> None: record_dashboard_state(self.areno, stage="save_checkpoint_end", epoch=epoch, step=step, role="policy") -def _record_to_train_sequence(record: Any, tokenizer, processor=None, *, max_prompt_tokens: int, max_new_tokens: int): +def _record_to_train_sequence( + record: Any, tokenizer, processor=None, *, max_prompt_tokens: int, max_new_tokens: int, sft_assistant_turns: str = "all" +): """Normalize one loader-produced SFT row into backend training format. `prompt_mask=True` means "do not train this source token"; the backend loss is next-token aligned, so the loss function later uses positions after the prompt prefix. RL-only fields are filled with zeros to satisfy the shared `TrainSequence` packing contract. + + Two row schemas are accepted: + + * ``{"prompt": str, "response": str}`` – single-turn (original format). + * ``{"messages": list[dict]}`` – multi-turn chat; each dict has ``role`` + and ``content`` keys. ``sft_assistant_turns`` controls which assistant + turns are trainable: ``"all"`` (default) trains every assistant turn, + ``"last"`` trains only the final assistant turn. """ record = dict(record) eos_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0 + if sft_assistant_turns not in ("all", "last"): + raise ValueError( + f"sft_assistant_turns must be 'all' or 'last', got {sft_assistant_turns!r}" + ) if record_has_image(record): if "response" not in record: raise ValueError("SFT image rows must contain `response`") @@ -240,18 +256,28 @@ def _record_to_train_sequence(record: Any, tokenizer, processor=None, *, max_pro features=features, eos_token_id=int(record.get("eos_token_id", eos_token_id)), ) - if "prompt" not in record or "response" not in record: + if "messages" in record: + messages = record["messages"] + if not isinstance(messages, list) or not messages: + return None + if not any(msg.get("role") == "assistant" for msg in messages): + return None + tokens, prompt_mask = messages_to_tokens_and_mask( + messages, tokenizer, eos_token_id, last_assistant_only=(sft_assistant_turns == "last") + ) + elif "prompt" in record and "response" in record: + if record["prompt"] is None or record["response"] is None: + return None + prompt = str(record["prompt"]) + response = str(record["response"]) + if not response: + return None + tokens, prompt_mask = prompt_response_to_tokens_and_mask(prompt, response, tokenizer, eos_token_id) + else: raise ValueError( - "SFT dataset loader must return rows with `prompt` and `response`; " - "normalize raw dataset fields in --dataset-loader-fn" + "SFT dataset loader must return rows with `prompt` and `response`, " + "or `messages`; normalize raw dataset fields in --dataset-loader-fn" ) - if record["prompt"] is None or record["response"] is None: - return None - prompt = str(record["prompt"]) - response = str(record["response"]) - if not response: - return None - tokens, prompt_mask = prompt_response_to_tokens_and_mask(prompt, response, tokenizer, eos_token_id) if len(tokens) < 2: return None diff --git a/areno/cli/train.py b/areno/cli/train.py index 5991be5f..bddd9500 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -122,6 +122,7 @@ def flash_attention_unsupported_model_reason(model_config): "agent_fn", "agent_timeout_s", "train_tool_results", + "sft_assistant_turns", "reward_fn_path", "reward_ckpt", ), @@ -976,6 +977,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: agent_timeout_s=args.agent_timeout_s, train_tool_results=args.train_tool_results, chat_template_enable_thinking=chat_template_enable_thinking, + sft_assistant_turns=args.sft_assistant_turns, lora=lora, reference_mode=args.reference_mode, ) @@ -1254,6 +1256,7 @@ def section(title: str, names: list[str]) -> dict: "agent_fn", "agent_timeout_s", "train_tool_results", + "sft_assistant_turns", "reward_fn_path", "reward_ckpt", ], @@ -1792,6 +1795,13 @@ def _dataset_builder_for_suffix(suffix: str) -> str: "--agent-timeout-s", type=float, default=300.0, show_default=True, help="Agentic rollout proxy request timeout." ) @click.option("--train-tool-results", is_flag=True, help="Include tool-result spans in agentic policy loss.") +@click.option( + "--sft-assistant-turns", + type=click.Choice(["all", "last"], case_sensitive=False), + default="all", + show_default=True, + help="SFT: train on every assistant turn (all) or only the final assistant turn (last) in multi-turn data.", +) @click.option( "--gspo-clip-eps", type=float, default=3.0e-4, show_default=True, help="GSPO sequence-ratio clipping epsilon." ) diff --git a/docs/cli/dataset_loaders.rst b/docs/cli/dataset_loaders.rst index 934c5c50..17c068c7 100644 --- a/docs/cli/dataset_loaders.rst +++ b/docs/cli/dataset_loaders.rst @@ -23,7 +23,8 @@ SFT --- SFT always requires ``--dataset-loader-fn``. The loader must return rows with -``prompt`` and ``response`` keys: +either ``prompt`` and ``response`` keys (single-turn) or a ``messages`` key +(multi-turn chat): .. code-block:: python @@ -40,6 +41,21 @@ SFT always requires ``--dataset-loader-fn``. The loader must return rows with ) return records +For multi-turn chat data, return ``messages`` instead: + +.. code-block:: python + + {"messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "And 3+3?"}, + {"role": "assistant", "content": "6"}, + ]} + +The trainer trains on all assistant turns by default. Use +``--sft-assistant-turns last`` to train only on the final assistant response. +User, system, and tool-result tokens are always excluded. + For a concrete example, use ``--dataset-path yahma/alpaca-cleaned`` with ``examples/sft/alpaca/dataset_loader.py``. diff --git a/docs/cli/training.rst b/docs/cli/training.rst index efcfe68c..4bcc023b 100644 --- a/docs/cli/training.rst +++ b/docs/cli/training.rst @@ -240,6 +240,14 @@ and CUDA graph state; MLX retains one in-process model. results are environment observations rather than policy actions. Assistant text and assistant tool-call spans are trainable by default. +``--sft-assistant-turns [all|last]`` + SFT only. Controls which assistant turns in multi-turn chat data are + trainable. ``all`` (default) trains on every assistant turn. ``last`` + trains only the final assistant turn, treating earlier assistant responses + as context. User, system, and tool-result tokens are always excluded from + training. This option has no effect on single-turn ``prompt``/``response`` + SFT rows. + Agentic trajectories can contain multiple chat-completion turns for the same prompt/sample pair. The agent owns the OpenAI-style message list and returns trajectory turns with the model response; Areno converts those turns into token @@ -531,6 +539,38 @@ SFT instruction tuning SFT loaders must normalize raw rows to ``prompt`` and ``response`` dictionaries. The trainer performs tokenization and trains on the response suffix. +SFT also supports multi-turn chat data. Instead of ``prompt``/``response``, +the loader can return a ``messages`` field containing a list of +``{"role": "user"|"assistant"|"system"|"tool", "content": "..."}`` dicts. +The trainer tokenizes the full conversation and trains on assistant turns. +Use ``--sft-assistant-turns last`` to train only on the final assistant +response, which is useful for focused evaluation of end-to-end multi-turn +behavior: + +.. code-block:: bash + + areno train \ + --ckpt Qwen/Qwen3-0.6B \ + --dataset-path /path/to/multiturn.jsonl \ + --dataset-loader-fn /path/to/multiturn_loader.py \ + --algo sft \ + --sft-assistant-turns last \ + --tp-size 1 \ + --world-size 1 \ + --batch-size 2 \ + --mini-bs 1 + +A multi-turn SFT loader should produce rows like: + +.. code-block:: python + + {"messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "And 3+3?"}, + {"role": "assistant", "content": "6"}, + ]} + DPO preference training ~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/examples/sft/multiturn/README.md b/examples/sft/multiturn/README.md new file mode 100644 index 00000000..13a6184f --- /dev/null +++ b/examples/sft/multiturn/README.md @@ -0,0 +1,65 @@ +# Multi-turn Chat SFT Example + +This example shows the multi-turn SFT dataset-loader contract. The loader +accepts rows with a `messages` field (OpenAI/HF chat format) or a +`conversations` field (ShareGPT format) and normalizes them to the SFT +trainer's `messages` schema. + +## Sample data + +Create a JSONL file with multi-turn conversations: + +```jsonl +{"messages": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4"}, {"role": "user", "content": "And 3+3?"}, {"role": "assistant", "content": "6"}]} +{"messages": [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}]} +``` + +## Train on all assistant turns (default) + +```bash +areno train \ + --algo sft \ + --ckpt Qwen/Qwen3-0.6B \ + --dataset-path /path/to/multiturn.jsonl \ + --dataset-loader-fn examples/sft/multiturn/dataset_loader.py \ + --tp-size 1 \ + --world-size 1 \ + --batch-size 2 \ + --mini-bs 1 +``` + +## Train only on the final assistant turn + +Use `--sft-assistant-turns last` to train only on the last assistant response +in each conversation. Earlier assistant turns are treated as context. This is +useful for focused evaluation of end-to-end multi-turn behavior. + +```bash +areno train \ + --algo sft \ + --ckpt Qwen/Qwen3-0.6B \ + --dataset-path /path/to/multiturn.jsonl \ + --dataset-loader-fn examples/sft/multiturn/dataset_loader.py \ + --sft-assistant-turns last \ + --tp-size 1 \ + --world-size 1 \ + --batch-size 2 \ + --mini-bs 1 +``` + +## Observable output + +- The training config summary printed at startup shows the resolved + `sft_assistant_turns` value under the **Rollout** section. +- The dashboard run-config JSON includes `sft_assistant_turns` in its settings. +- Training logs show `sft_dataset_filter` counts for rows skipped due to + empty messages, missing assistant turns, or budget limits. + +## Mask semantics + +- `user`, `system`, and `tool` tokens are always excluded from the loss. +- With `--sft-assistant-turns all` (default): every `assistant` turn is + trainable. +- With `--sft-assistant-turns last`: only the final `assistant` turn is + trainable; earlier `assistant` turns are context. +- EOS is appended after the last message and is trainable in both modes. diff --git a/examples/sft/multiturn/dataset_loader.py b/examples/sft/multiturn/dataset_loader.py new file mode 100644 index 00000000..1047e8d3 --- /dev/null +++ b/examples/sft/multiturn/dataset_loader.py @@ -0,0 +1,30 @@ +"""Dataset loader for multi-turn chat SFT rows.""" + +from __future__ import annotations + + +def load_training_dataset(dataset_path: str, *, default_loader, **_: object) -> list[dict]: + """Normalize multi-turn chat rows to SFT ``messages`` format. + + Expects raw rows with a ``messages`` field (OpenAI/HF chat format) or a + ``conversations`` field (ShareGPT format). Each item should have ``role`` + and ``content`` keys. + """ + + records = [] + for row in default_loader(dataset_path): + record = dict(row) + messages = record.get("messages") or record.get("conversations") + if not messages: + continue + normalized = [] + for msg in messages: + role = str(msg.get("role", "user")).strip() + content = str(msg.get("content", "")).strip() + if not content: + continue + normalized.append({"role": role, "content": content}) + if not normalized: + continue + records.append({"messages": normalized}) + return records diff --git a/tests/test_config_data_cpu.py b/tests/test_config_data_cpu.py index 3e8b5d0b..ceaa094e 100644 --- a/tests/test_config_data_cpu.py +++ b/tests/test_config_data_cpu.py @@ -1072,6 +1072,7 @@ def _train_args(**overrides): agent_fn=None, agent_timeout_s=300.0, train_tool_results=False, + sft_assistant_turns="all", gspo_clip_eps=3.0e-4, grpo_clip_eps=0.2, ref_ckpt=None, diff --git a/tests/test_train_cli_config_cpu.py b/tests/test_train_cli_config_cpu.py index 0d0542cd..2bd7807c 100644 --- a/tests/test_train_cli_config_cpu.py +++ b/tests/test_train_cli_config_cpu.py @@ -1014,6 +1014,7 @@ def _options(**overrides): agent_fn=None, agent_timeout_s=300.0, train_tool_results=False, + sft_assistant_turns="all", gspo_clip_eps=3.0e-4, grpo_clip_eps=0.2, ref_ckpt=None, diff --git a/tests/test_trainer_dataset_utils_cpu.py b/tests/test_trainer_dataset_utils_cpu.py index 45368d1a..66e839e2 100644 --- a/tests/test_trainer_dataset_utils_cpu.py +++ b/tests/test_trainer_dataset_utils_cpu.py @@ -236,5 +236,211 @@ def test_first_value_returns_only_explicit_keys(self): self.assertIsNone(data_utils.first_value({}, ("chosen",))) -if __name__ == "__main__": - unittest.main() +class MultiTurnSFTMaskTest(unittest.TestCase): + """Multi-turn SFT messages_to_tokens_and_mask and trainer integration.""" + + def test_messages_mask_trains_all_assistant_turns_by_default(self): + """Default mode trains every assistant turn; user/system/tool excluded.""" + tokenizer = FakeTextTokenizer() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "bye"}, + {"role": "assistant", "content": "see you"}, + ] + + tokens, mask = data_utils.messages_to_tokens_and_mask( + messages, tokenizer, tokenizer.eos_token_id + ) + + # EOS is appended at the end and is trainable. + self.assertEqual(tokens[-1], tokenizer.eos_token_id) + self.assertFalse(mask[-1]) + + # Every non-assistant token must be masked (True = do not train). + # We check that the mask has both True and False values. + self.assertTrue(any(mask), "expected at least some masked (context) tokens") + self.assertTrue(any(not m for m in mask), "expected at least some trainable tokens") + + def test_messages_mask_last_assistant_only_masks_earlier_assistants(self): + """last_assistant_only=True masks earlier assistant turns as context.""" + tokenizer = FakeTextTokenizer() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "bye"}, + {"role": "assistant", "content": "see you"}, + ] + + tokens_all, mask_all = data_utils.messages_to_tokens_and_mask( + messages, tokenizer, tokenizer.eos_token_id, last_assistant_only=False + ) + tokens_last, mask_last = data_utils.messages_to_tokens_and_mask( + messages, tokenizer, tokenizer.eos_token_id, last_assistant_only=True + ) + + # Same tokens, but fewer trainable positions in last-only mode. + self.assertEqual(tokens_all, tokens_last) + trainable_all = sum(1 for m in mask_all if not m) + trainable_last = sum(1 for m in mask_last if not m) + self.assertGreater(trainable_all, trainable_last) + self.assertGreater(trainable_last, 0) + + def test_messages_mask_excludes_user_and_system_tokens(self): + """User and system turns must never be trainable.""" + tokenizer = FakeTextTokenizer() + messages = [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + + tokens, mask = data_utils.messages_to_tokens_and_mask( + messages, tokenizer, tokenizer.eos_token_id + ) + + # The system and user content tokens should all be masked. + # We verify by checking that there exist masked positions before the + # first trainable position (i.e. system+user are context). + first_trainable = next((i for i, m in enumerate(mask) if not m), None) + self.assertIsNotNone(first_trainable) + self.assertTrue(all(mask[:first_trainable]), "system and user tokens must be masked") + + def test_messages_mask_with_tool_role_excludes_tool_results(self): + """Tool-result turns must be excluded from training.""" + tokenizer = FakeTextTokenizer() + messages = [ + {"role": "user", "content": "use the tool"}, + {"role": "assistant", "content": "calling tool"}, + {"role": "tool", "content": "result data"}, + {"role": "assistant", "content": "final answer"}, + ] + + tokens, mask = data_utils.messages_to_tokens_and_mask( + messages, tokenizer, tokenizer.eos_token_id + ) + + # Tool turn tokens must be masked (True = do not train). + # There should be at least 2 assistant trainable spans. + trainable_count = sum(1 for m in mask if not m) + self.assertGreater(trainable_count, 0) + + def test_messages_empty_messages_returns_eos_only(self): + """Empty messages list should produce a minimal EOS-only sequence.""" + tokenizer = FakeTextTokenizer() + + tokens, mask = data_utils.messages_to_tokens_and_mask( + [], tokenizer, tokenizer.eos_token_id + ) + + self.assertEqual(len(tokens), 1) + self.assertEqual(tokens[0], tokenizer.eos_token_id) + self.assertFalse(mask[0]) + + def test_sft_record_accepts_messages_format(self): + """SFT trainer should accept {messages: [...]} rows for multi-turn.""" + tokenizer = FakeTextTokenizer() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + + seq = sft_mod._record_to_train_sequence( + {"messages": messages}, tokenizer, max_prompt_tokens=16, max_new_tokens=16 + ) + + self.assertIsNotNone(seq) + self.assertEqual(seq.eos_token_id, 99) + self.assertEqual(seq.tokens[-1], 99) + # There must be trainable (False) positions in the mask. + self.assertTrue(any(not m for m in seq.prompt_mask)) + + def test_sft_record_messages_last_assistant_only(self): + """sft_assistant_turns='last' should reduce trainable tokens vs 'all'.""" + tokenizer = FakeTextTokenizer() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "first answer"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "second answer"}, + ] + + seq_all = sft_mod._record_to_train_sequence( + {"messages": messages}, tokenizer, + max_prompt_tokens=128, max_new_tokens=128, sft_assistant_turns="all", + ) + seq_last = sft_mod._record_to_train_sequence( + {"messages": messages}, tokenizer, + max_prompt_tokens=128, max_new_tokens=128, sft_assistant_turns="last", + ) + + self.assertIsNotNone(seq_all) + self.assertIsNotNone(seq_last) + trainable_all = sum(1 for m in seq_all.prompt_mask if not m) + trainable_last = sum(1 for m in seq_last.prompt_mask if not m) + self.assertGreater(trainable_all, trainable_last) + + def test_sft_record_invalid_assistant_turns_raises(self): + """Invalid sft_assistant_turns value should raise a clear ValueError.""" + tokenizer = FakeTextTokenizer() + + with self.assertRaisesRegex(ValueError, "sft_assistant_turns must be"): + sft_mod._record_to_train_sequence( + {"prompt": "q", "response": "a"}, tokenizer, + max_prompt_tokens=16, max_new_tokens=16, sft_assistant_turns="invalid", + ) + + def test_sft_record_empty_messages_returns_none(self): + """Empty messages list with no trainable assistant should be filtered.""" + tokenizer = FakeTextTokenizer() + + seq = sft_mod._record_to_train_sequence( + {"messages": []}, tokenizer, max_prompt_tokens=16, max_new_tokens=16 + ) + + # Empty messages produces EOS-only; response_tokens == 0 so filtered. + self.assertIsNone(seq) + + def test_sft_record_messages_only_user_returns_none(self): + """Messages with no assistant turn should be filtered (no trainable tokens).""" + tokenizer = FakeTextTokenizer() + messages = [ + {"role": "user", "content": "just a question"}, + ] + + seq = sft_mod._record_to_train_sequence( + {"messages": messages}, tokenizer, max_prompt_tokens=128, max_new_tokens=128 + ) + + self.assertIsNone(seq) + + def test_sft_record_messages_none_returns_none(self): + """None messages should be filtered out.""" + tokenizer = FakeTextTokenizer() + + seq = sft_mod._record_to_train_sequence( + {"messages": None}, tokenizer, max_prompt_tokens=16, max_new_tokens=16 + ) + + self.assertIsNone(seq) + + def test_sft_record_neither_prompt_response_nor_messages_raises(self): + """Rows without prompt/response/messages should raise a clear error.""" + tokenizer = FakeTextTokenizer() + + with self.assertRaisesRegex(ValueError, "must return rows with"): + sft_mod._record_to_train_sequence( + {"text": "raw"}, tokenizer, max_prompt_tokens=16, max_new_tokens=16 + ) + + def test_sft_prompt_response_still_works_with_assistant_turns(self): + """Existing single-turn prompt/response path should still work with the new option.""" + tokenizer = FakeTextTokenizer() + + seq = sft_mod._record_to_train_sequence( + {"prompt": "q", "response": "a"}, tokenizer, + max_prompt_tokens=16, max_new_tokens=16, sft_assistant_turns="last", + ) + + self.assertIsNotNone(seq) + self.assertEqual(seq.tokens[-1], 99)