Skip to content

dew.data.chat

Supervised fine-tuning data: conversations with a role on every token.

A ChatMessages source reads conversations from a parquet file, a JSONL file or a Hub dataset id, and renders each with the tokenizer’s chat template. Every token gets the role of the message that wrote it, so LMObjective with loss_role=Role.ASSISTANT trains on assistant tokens only.

Packing is the token pipeline’s plan over the whole corpus (PackedWindows) with text_roles as one more per-token field. A window carries text, text_roles, text_segment_ids, text_positions and the identical text_roles_segment_ids, text_roles_positions, all aligned. The training stream’s position is a global window count that resumes on any process count.

Conversations are structured. A Message carries what the Hugging Face chat-template contract reads: a role; content that is a string, a list of typed parts, or nothing; an assistant’s tool_calls; a tool response’s tool_call_id and name; and any further keys the template wants (reasoning_content, thinking) untouched. A Conversation adds the tools schemas the template renders into its system block. The stored messages retain typed content. Text-tokenizer rendering joins all-text parts into strings and refuses other part types, which require a processor. Tool-call fields and message metadata reach the template unchanged.

NameSummary
ROLES_KEYBatch key the chat pipeline packs [B, seq_len + 1] int8 roles under.
RoleWho wrote a token of a rendered conversation.
ContentPartOne typed block of a message’s content.
ToolCallOne function call an assistant message asks for.
RESERVEDThe message keys the boundary types; every other key rides in extra.
MessageOne turn, in the shape the chat template reads.
ConversationThe messages of one row and the tool schemas its template renders.
render_promptTokenizes a text conversation with the next assistant header.
render_conversationToken ids and per-token roles for one conversation.
ConversationSourceReads the conversations at path by index, with their tool schemas beside them where the rows carry any.
RenderConversationTurns source rows into rendered ids and per-token roles, for map_with_index.
ChatMessagesRenders conversations with the tokenizer’s chat template and packs them.

attribute source

ROLES_KEY = 'text_roles'

Batch key the chat pipeline packs [B, seq_len + 1] int8 roles under.

class source

class Role(int, Enum)

Who wrote a token of a rendered conversation.

The values are the text_roles column: 0 pads, the rest name the message whose span holds the token. An assistant turn’s opening header belongs to no span; it reads as padding, so only the completion counts. DEVELOPER is the Harmony instruction role; templates without it render the message to nothing, which the renderer refuses.

dataclass source

class ContentPart(type: str, fields: Mapping[str, object])

One typed block of a message’s content.

type names the block; fields is the rest of it verbatim, text for a text block, a url or a path for media. Text rendering accepts text parts only; media stays available in the stored conversation.

def parse(part: Mapping[str, object], where: str) -> ContentPart
def as_template() -> Mapping[str, object]

dataclass source

class ToolCall(
name: str,
arguments: Mapping[str, object],
id: str | None = None,
type: str = 'function',
)

One function call an assistant message asks for.

arguments is the parsed object. A source may hold it as a JSON string, since parquet cannot carry a struct whose fields differ per call, while the Hugging Face contract hands templates a mapping. The string parses here, and a string that is not a JSON object fails. id links the call to the tool_call_id of its response where the template uses ids. type is the contract’s "function". A source may write the call flat ({name, arguments}) or nested under function; both read the same.

def parse(call: Mapping[str, object], where: str) -> ToolCall
def as_template() -> Mapping[str, object]

attribute source

RESERVED = ('role', 'content', 'tool_calls', 'tool_call_id', 'name')

The message keys the boundary types; every other key rides in extra.

dataclass source

class Message(
role: Role,
content: str | tuple[ContentPart, ...] | None,
tool_calls: tuple[ToolCall, ...] = (),
tool_call_id: str | None = None,
name: str | None = None,
extra: Mapping[str, object] = dict(),
)

One turn, in the shape the chat template reads.

content is the text, the typed parts, or None when the turn is only its tool calls. None reaches the template as None, the wire form tool loops send. tool_calls belong to assistant turns and tool_call_id to tool responses; name is the function a response answers, or a participant’s name on other roles. extra holds every further key verbatim, so reasoning_content or thinking reach the template that reads them. A value of None on an optional key reads as absent, which is how parquet spells a field a row does not have.

def parse(message: Mapping[str, object], where: str) -> Message
def as_template() -> Mapping[str, object]

The structured HF message, retaining content parts and metadata.

Conversation.text_rows adapts these fields for text tokenizers.

dataclass source

class Conversation(
messages: tuple[Message, ...],
tools: tuple[Mapping[str, object], ...] = (),
)

The messages of one row and the tool schemas its template renders.

tools are the JSON schemas apply_chat_template takes; a source may hold them as a JSON string. With none, the template renders its plain system block.

def parse(
messages: object,
tools: object = None,
where: str = 'conversation',
) -> Conversation

One row’s messages and tool schemas, as a parquet column holds them.

Each is a list, or the JSON text a column of varying schema carries.

def rows() -> list[Mapping[str, object]]
def text_rows(where: str) -> list[Mapping[str, object]]

The text-tokenizer input, with all-text parts joined in order.

The stored messages and rows retain their structured content. Media needs a processor to expand its payload into model inputs; a text tokenizer cannot do that, even if its template emits a marker.

function source

def render_prompt(
tokenizer: PreTrainedTokenizerBase,
conversation: Conversation,
where: str,
thinking: bool | None = None,
) -> list[int]

Tokenizes a text conversation with the next assistant header.

SFT and prompt sampling use the same message conversion and tool schemas. All-text parts concatenate in order; nontext parts require a processor. thinking is _token_ids’s.

function source

def render_conversation(
tokenizer: PreTrainedTokenizerBase,
conversation: Conversation,
where: str,
) -> tuple[np.ndarray, np.ndarray]

Token ids and per-token roles for one conversation.

Message k’s span comes from prefix rendering. An assistant turn is exact: messages[:k] with the generation prompt against messages[:k+1] without it, and both must be prefixes of the whole render. The completion the loss counts is then the completion the model will see. A template that renders the turn differently once later messages follow it, or whose generation prompt is not the turn’s header, fails here with the message index rather than mis-masking it.

Every other turn spans from where the previous turn’s render ended to where its own render stops agreeing with the whole. That is what a template needs when it re-segments a run of tool responses into one block. A turn that renders to no tokens, or rewrites an earlier turn’s tokens, fails. What a template emits before any message, a tools block for one, has no message of its own and counts as the first message’s span.

where names the tokenizer and the row for those refusals. The arrays are int32 ids and int8 roles.

class source

class ConversationSource(
path: str,
*,
column: str = 'messages',
split: str = 'train',
options: HFOptions | None = None,
)

Reads the conversations at path by index, with their tool schemas beside them where the rows carry any.

Three things hold conversations and one reader takes all three: a parquet file, a .jsonl file, and a Hub dataset id resolved through HFOptions. The suffix decides which. chat.jsonl is lines, chat.parquet is a table, an existing file without either suffix is a table too, and anything else is a repo id at split.

One record is one conversation, a list of messages in the verl layout, and its tool schemas, a list or a JSON string. The rows are read once and come back as plain dicts, which pickle across to grain workers. Other columns are not read.

class source

class RenderConversation(tokenizer: str)

Turns source rows into rendered ids and per-token roles, for map_with_index.

Holds only the tokenizer path, so grain workers unpickle the name and load their own copy. Failures name the tokenizer and the row.

dataclass source

class ChatMessages(
tokenizer: str,
path: str | None = None,
val_path: str | None = None,
column: str = 'messages',
split: str = 'train',
val_split: str | None = None,
options: HubOptions = HFOptions(),
seq_len: int = 256,
val_batches: int | None = 4,
packing_bins: int = 8,
*,
seed: int = 0,
loading: Loading = Loading(),
)

Renders conversations with the tokenizer’s chat template and packs them.

path names the conversations: a parquet file, a .jsonl file, or a Hub dataset id read at split through options, which is the same value the hf provider forwards to datasets.load_dataset. Whichever it is, the rows carry lists of messages under column (or prompt, which is what the verl layout calls it) and their tool schemas under tools where they have any. tokenizer is the hub name or local path whose chat template renders them.

Each conversation, in chunks when it outgrows the window, is one element the packing plan adds to the first window with room. Every window carries text_roles beside the ids, so the loss can count one role’s targets. The plan runs over the whole corpus in row order, ahead of the shard, as PackedTokens plans its documents, so records is the windows of a pass exactly and a saved position is a global window count.

val_path is a second source of the same three kinds, read at val_split and scored as one pass; None trains without validation.

column: str

The column of conversations; prompt is read where a row has that.

split: str

Which split path is read at, when it names a Hub dataset.

val_split: str | None

Which split val_path is read at; None reads split.

options: HubOptions

What datasets.load_dataset takes beside the id and the split.

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset