fix: CompletionsDataset.process crashes when mask_prompt is enabled

[iamadalek]

Overview

CompletionsDataset.process() crashes when mask_prompt=True because a single message dict is passed to apply_chat_template instead of a list.

Problem

When training with the completions format and mask_prompt enabled, the process method builds a two-element messages list, then passes messages[0] (a bare dict) to tokenizer.apply_chat_template(). This method expects a list of message dicts, causing a crash.

The analogous code path in ChatDataset.process() correctly passes a list slice (messages[:-1]), so this appears to be an oversight specific to the completions path.

Requirements

• The mask-prompt code path in CompletionsDataset.process() must pass a properly-typed argument (list of message dicts) to apply_chat_template
• The prompt masking offset calculation must match the behavior of ChatDataset.process() (i.e., tokenize only the user prompt portion with add_generation_prompt=True)
• Existing non-mask-prompt behavior must remain unchanged

Acceptance Criteria

• Training with completions format and mask_prompt=True no longer crashes
• The prompt mask offset correctly covers only the prompt tokens (not the completion)
• Training with completions format and mask_prompt=False is unaffected
• ChatDataset with mask_prompt=True continues to work correctly

Constraints

• Fix should be minimal and not change any unrelated behavior
• Must be backwards-compatible with existing training configs

Reference

• ChatDataset.process() handles the same pattern correctly and serves as the reference implementation

[iamadalek]

Design Session

Complexity triage: solo — lead assessment only, no specialist dispatch.

Decision

Change messages[0] to messages[:1] on line 119 of mlx_lm/tuner/datasets.py in CompletionsDataset.process(). This mirrors the pattern used by ChatDataset.process() on line 69, which correctly passes a list slice (messages[:-1]) to apply_chat_template().

Rationale

• tokenizer.apply_chat_template() expects a List[Dict], not a bare Dict
• ChatDataset.process() already demonstrates the correct pattern on line 69: messages[:-1] produces a list slice
• CompletionsDataset always has exactly 2 messages (user + assistant), so messages[:1] is equivalent to [messages[0]] — the user prompt wrapped in a list
• The add_generation_prompt=True parameter is already correctly set (matches ChatDataset behavior for assistant-role last messages)
• Using messages[:1] (slice) over [messages[0]] (wrap) for consistency with ChatDataset's slicing idiom

Specialist Input

Specialist	Finding	Confidence
Lead (solo assessment)	Single-character fix (messages[0] → messages[:1]), bug cause and fix pattern are both evident from the code. ChatDataset serves as a verified reference.	High

Conflicts Resolved

N/A — solo assessment.

Action Items (Inline)

• Change messages[0] to messages[:1] on line 119 of mlx_lm/tuner/datasets.py
• Add test covering CompletionsDataset.process() with mask_prompt=True
• Verify mask_prompt=False path is unaffected
• Verify ChatDataset.process() with mask_prompt=True still works

Action Items (Separate Issue Required)

None identified.

Open Questions

None.

---

Solo design assessment via /coding-workflows:prepare-issue (complexity: low)

[iamadalek]

Implementation Plan

Prior Design Decisions

From design session comment

• Change messages[0] to messages[:1] on line 119 of mlx_lm/tuner/datasets.py
• Use slice syntax ([:1]) over wrap syntax ([messages[0]]) for consistency with ChatDataset's idiom
• Fix must be minimal, no unrelated behavior changes

Specialist Review

Specialist	Recommendation	Key Input
Tuner (fine-tuning specialist)	Approve	Core fix is correct. Suggested aligning add_generation_prompt logic, but this is unnecessary since CompletionsDataset always has an assistant message last — hardcoded True is already equivalent to the conditional.

Approach

Single-line fix: change messages[0] (bare dict) to messages[:1] (list slice) in CompletionsDataset.process() at line 119 of mlx_lm/tuner/datasets.py. This makes the argument type match what tokenizer.apply_chat_template() expects (a List[Dict]) and mirrors the pattern used by ChatDataset.process() at line 69.

Components

1. Bug fix: messages[0] → messages[:1] on line 119
2. Test coverage: Add mask_prompt=True test cases for both CompletionsDataset and ChatDataset

Files to Modify

• mlx_lm/tuner/datasets.py — Change messages[0] to messages[:1] on line 119
• tests/test_datsets.py — Add test_completions_mask_prompt and test_chat_mask_prompt test methods

Testing Strategy

• Unit: test_completions_mask_prompt — Create CompletionsDataset with mask_prompt=True, call process(), verify no crash, verify offset > 0 and offset < len(tokens)
• Unit: test_chat_mask_prompt — Create ChatDataset with mask_prompt=True, call process(), verify offset > 0 and offset < len(tokens) (regression guard)
• Existing: test_completions — Verify mask_prompt=False path is unaffected (already tested, no changes needed)
• Existing: test_chat — Verify mask_prompt=False path is unaffected (already tested, no changes needed)

Risks & Mitigations

Risk	Mitigation
Offset value incorrect after fix	Tests assert offset > 0 and offset < len(tokens)
Regression in mask_prompt=False path	Existing tests cover this; no code changes on that path

---

Plan created with specialist input via /coding-workflows:plan-issue

[iamadalek]

Plan Review

Reviewers

Specialist	Verdict	Key Concern
Tuner (adversarial)	Approve	No blocking issues. Recommends adding explicit mask_prompt=False offset assertion and noting the pre-existing test filename typo.

Blocking Issues

None.

Important Issues

• Add explicit regression assertion in existing test_completions and test_chat that verifies offset == 0 when mask_prompt=False -- raised by tuner specialist
• Note that tests/test_datsets.py has a pre-existing filename typo (missing 'a') — plan should use the existing filename as-is to stay minimal -- raised by tuner specialist

Minor Suggestions

• Consider a token-decoding validation test that checks semantic content of tokens[:offset] vs tokens[offset:] -- raised by tuner specialist
• Consider edge case test with empty completion string -- raised by tuner specialist

Overall Verdict: Needs Revision

The core fix and approach are correct. Two important issues need addressing: (1) the new tests should also verify offset == 0 for the mask_prompt=False path as an explicit regression guard, and (2) plan should acknowledge the pre-existing test file typo.

---

Adversarial review via /coding-workflows:review-plan

[iamadalek]

Implementation Plan (Revised)

Issues Addressed

Issue	Severity	Resolution
Missing mask_prompt=False offset assertion	Important	New tests verify offset == 0 for both CompletionsDataset and ChatDataset when mask_prompt=False
Test filename typo	Important	Using existing tests/test_datsets.py as-is (pre-existing, out of scope for this fix)

Prior Design Decisions

From design session comment

• Change messages[0] to messages[:1] on line 119 of mlx_lm/tuner/datasets.py
• Use slice syntax ([:1]) for consistency with ChatDataset's idiom
• Fix must be minimal, no unrelated behavior changes

Approach

Single-line fix: change messages[0] (bare dict) to messages[:1] (list slice) in CompletionsDataset.process() at line 119 of mlx_lm/tuner/datasets.py. This makes the argument type match what tokenizer.apply_chat_template() expects (a List[Dict]) and mirrors the pattern used by ChatDataset.process() at line 69.

Components

1. Bug fix: messages[0] → messages[:1] on line 119
2. Test coverage: Add mask_prompt=True and mask_prompt=False test cases for both dataset types

Files to Modify

• mlx_lm/tuner/datasets.py — Change messages[0] to messages[:1] on line 119
• tests/test_datsets.py — Add test_completions_mask_prompt and test_chat_mask_prompt test methods

Testing Strategy

• Unit: test_completions_mask_prompt — Create CompletionsDataset with mask_prompt=True, call process(), verify no crash, verify offset > 0 and offset < len(tokens). Then create with mask_prompt=False, call process(), verify offset == 0.
• Unit: test_chat_mask_prompt — Create ChatDataset with mask_prompt=True, call process(), verify offset > 0 and offset < len(tokens). Then create with mask_prompt=False, call process(), verify offset == 0.
• Existing: test_completions — Unchanged, covers basic mask_prompt=False loading path
• Existing: test_chat — Unchanged, covers basic mask_prompt=False loading path

Risks & Mitigations

Risk	Mitigation
Offset value incorrect after fix	Tests assert offset > 0 and offset < len(tokens) for mask_prompt=True
Regression in mask_prompt=False path	New tests explicitly assert offset == 0; existing tests unchanged
Test filename typo causes confusion	Documented as pre-existing, not addressed in this PR

---

Revised plan addressing review feedback
Review: #1 (comment)

[iamadalek]

Session Checkpoint

Branch: fix/completions-mask-prompt-type-error
Last commit: 7fc4f20 fix: pass list slice to apply_chat_template in CompletionsDataset

Completed:

• Bug fix: messages[0] → messages[:1] on line 119 of mlx_lm/tuner/datasets.py
• Test: test_completions_mask_prompt (mask_prompt=True and False)
• Test: test_chat_mask_prompt (mask_prompt=True and False)
• Full suite: 159 tests passing
• PR created: fix: CompletionsDataset mask_prompt passes wrong type to apply_chat_template #2