Keep all turns and fix system prompt in ShareGPT conversations parsing#345
Merged
yghstill merged 1 commit intoJun 20, 2026
Merged
Conversation
TextDataset._prepare_messages only read the first two entries of a ShareGPT-style "conversations" record and assigned roles positionally, silently dropping every later turn of a multi-turn conversation and mislabeling system-led ones. The branch also guarded on data["system"] but read data["system_prompt"], raising KeyError on any record that carries a system field. Iterate over all turns, carry each turn's role/content through the existing normalization, prepend the optional dataset-level system prompt (deduplicated against a leading system turn), and skip turns with no text payload. Harden the normalization loop so a turn missing a role no longer raises. Add CPU-only unit tests covering the multi-turn, system-field, dedup and malformed-turn cases plus the unchanged messages/input-output branches.
yghstill
approved these changes
Jun 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Fix ShareGPT
conversationshandling inTextDataset._prepare_messages(
angelslim/data/text_dataset.py). The branch had two correctness bugs:KeyError: 'system_prompt'— it guarded ondata["system"]but thenread
data["system_prompt"]:Any ShareGPT record that carries a
systemfield crashes during dataloading. (The sibling
messagesandinput/outputbranches correctlyguard and read
system_prompt; only this branch was inconsistent.)Multi-turn conversations were truncated to the first two turns and roles
were assigned positionally:
A standard multi-turn ShareGPT record (
[{"from": "human", ...}, {"from": "gpt", ...}, {"from": "human", ...}, ...]) lost every turn after the second,and the actual
fromof each turn was ignored. This defeats the downstreamloop, which already scans for the last assistant turn and treats everything
before it as prompt context — i.e. the pipeline is built for multi-turn data.
How
existing role-normalization loop (which already maps
from→role,value→content, andhuman→user/gpt→assistant).data["system"], withdata["system_prompt"]accepted as a fallback), deduplicated against aconversation that already begins with a
systemturn so it is not doubled.turn missing a role no longer raises. The
messagesandinput/outputbranches are left behavior-identical (the
.getis a no-op for them).Tests
tests/test_text_dataset_messages.py— CPU-only, no GPU / model weights / heavystack (the
torch/transformers/datasets/pyarrowimports are stubbed;_prepare_messagesis pure-dict). Covers: system-field no-crash regression,all-turns-preserved, leading-system-turn dedup, malformed-turn skip, and the
unchanged
messages/input-outputbranches. The fourconversationscasesfail on
mainand pass here; the two non-conversationscases pass on both.Note on reuse
angelslim/compressor/speculative/train/data/data_utils.pyalready has aconvert_sharegpt_datahelper with the samehuman→user/gpt→assistantmapping. It is intentionally not reused here because it lives in the
speculative.train.datasubpackage and pullstransformers.image_utilsandtraining-only deps at import — importing it into the core data loader would add
a heavier cross-package dependency than the few lines here. It also indexes
role_mapping[message["from"]]directly, which raises on asystemturn,whereas the normalization in
text_dataset.pyhandles system turns.Reviewer check: does keeping all turns + prepending the system prompt match
the intended ShareGPT calibration format, or should the loader continue to
support only single-turn
conversations? The downstream_load_dataloop(scanning for the last assistant turn) suggests multi-turn is intended.