File size: 9,986 Bytes
685e018 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 | """Render the general-chat corpus into the adaptive training layout.
The corpus (``scripts/gen_chat_dataset.py``) stores conversations as system prompt plus turns,
each assistant turn carrying its own notes. This module turns them into the layout the hybrid
objective trains on: one example per assistant turn, whose prefix holds the system prompt, the
merged ledger of notes whose messages have fallen out of the visible window, and the last
``keep_messages`` messages verbatim.
Dropping the older messages is the point. With the full transcript in the prefix the model can
re-read instead of remember and the thinking block stops being memory, which is what
:func:`diffusion_lm.claims.to_examples` established for the claims corpus. The window and the
ledger merge are imported from that module rather than reimplemented, so the prefix a training
example sees is byte-identical to the one the playground builds at inference.
"""
from __future__ import annotations
import argparse
import json
import random
from collections import Counter
from pathlib import Path
from diffusion_lm.claims import (
IM_END,
KEEP_MESSAGES,
chat_prefix,
ledger_line,
ledger_notes,
merge_notes,
)
from diffusion_lm.reasoning import ReasoningExample
NOTE_JOIN = '; '
def load(paths: list[Path]) -> list[dict]:
"""Read consolidated corpus files, tagging each conversation with its origin file."""
conversations = []
for path in paths:
for line in path.open(encoding='utf-8'):
if not line.strip():
continue
conversation = json.loads(line)
conversation.setdefault('source', path.stem)
conversations.append(conversation)
return conversations
def _messages(conversation: dict) -> list[dict[str, str]]:
"""Corpus turns in the message shape the ledger helpers expect.
A turn's notes collapse into one ``note`` string joined by ``NOTE_JOIN``, which is the
separator :func:`diffusion_lm.claims.merge_notes` splits on, so a multi-fact turn still
contributes one ledger entry per fact.
"""
messages = []
for turn in conversation.get('turns') or []:
message = {'role': turn['role'], 'content': turn['content']}
notes = [str(note).strip() for note in (turn.get('thinking') or []) if str(note).strip()]
if turn['role'] == 'assistant' and notes:
message['note'] = NOTE_JOIN.join(notes)
messages.append(message)
return messages
def to_examples(
conversation: dict, *, keep_messages: int = KEEP_MESSAGES
) -> list[ReasoningExample]:
"""One example per assistant turn, each thinking note becoming its own block.
A turn the corpus marked as needing no notes yields an empty chain, which is what teaches
the controller to answer without opening a thinking block; the encoder accepts it.
"""
messages = _messages(conversation)
turns = conversation.get('turns') or []
reference = str(conversation.get('reference') or '')
last = max((i for i, m in enumerate(messages) if m['role'] == 'assistant'), default=-1)
examples = []
for index, message in enumerate(messages):
if message['role'] != 'assistant' or index == 0:
continue
history = messages[:index]
older = merge_notes(ledger_notes(history, keep_messages))
window = history[max(0, index - keep_messages):]
notes = [str(n).strip() for n in (turns[index].get('thinking') or []) if str(n).strip()]
examples.append(ReasoningExample(
problem=chat_prefix(window, system=conversation['system'], extra=ledger_line(older)),
steps=tuple(notes),
answer=message['content'] + IM_END,
expected_answer=reference if index == last else '',
))
return examples
def _document(conversation: dict, index: int) -> str:
"""Split key. Two conversations built from one passage share its facts.
Splitting by example would leak within a conversation as well, so the whole conversation
travels together and grounded slices travel with their source item.
"""
return str(conversation.get('source_id') or f'{conversation.get("source", "")}-{index}')
def _apply_caps(
conversations: list[dict], caps: dict[str, int], seed: int
) -> list[dict]:
"""Drop conversations so a source contributes at most ``caps[source]`` of them.
Capping is by CONVERSATION but the reason is examples: a source's weight in the mix is its
turn count, not its row count, and the two differ by an order of magnitude (CoQA yields 11.9
examples per conversation against 1.07 for a single-question source). Sampling is seeded and
whole conversations travel together, so the split stays document-clean.
"""
if not caps:
return conversations
rng = random.Random(seed)
by_source: dict[str, list[int]] = {}
for index, conversation in enumerate(conversations):
by_source.setdefault(conversation.get('source', ''), []).append(index)
dropped: set[int] = set()
for source, limit in caps.items():
indices = by_source.get(source)
if indices is None:
raise ValueError(f'no conversations carry source {source!r}')
if len(indices) <= limit:
print(f'cap {source}={limit}: {len(indices)} present, nothing dropped')
continue
dropped |= set(indices) - set(rng.sample(indices, limit))
print(f'cap {source}={limit}: dropped {len(indices) - limit:,} of {len(indices):,}')
return [c for index, c in enumerate(conversations) if index not in dropped]
def _parse_caps(pairs: list[str]) -> dict[str, int]:
caps = {}
for pair in pairs:
source, _, count = pair.partition('=')
if not count.isdigit():
raise ValueError(f'--cap expects SOURCE=N, got {pair!r}')
caps[source] = int(count)
return caps
def prepare(args: argparse.Namespace) -> None:
import numpy as np
from diffusion_lm.reasoning import ExampleEncoder, LayoutSpec, _write_packed, size_token_ids
from diffusion_lm.tokenizer import load_tokenizer
tokenizer = load_tokenizer(args.tokenizer)
spec = LayoutSpec(seq_len=args.seq_len, block=min(args.sizes), max_slots=args.max_slots,
sizes=tuple(sorted(args.sizes)))
encoder = ExampleEncoder(tokenizer, spec)
conversations = _apply_caps(load(args.inputs), _parse_caps(args.cap or []), args.seed)
documents = sorted({_document(c, i) for i, c in enumerate(conversations)})
rng = random.Random(args.seed)
rng.shuffle(documents)
held = set(documents[:max(1, round(len(documents) * args.val_fraction))])
split: dict[str, list[tuple]] = {'train': [], 'validation': []}
dropped = 0
blocks: Counter[int] = Counter()
empty_chains = 0
for index, conversation in enumerate(conversations):
bucket = 'validation' if _document(conversation, index) in held else 'train'
for example in to_examples(conversation, keep_messages=args.keep_messages):
encoded = encoder.encode_adaptive(example)
if encoded is None:
dropped += 1
continue
blocks.update(encoded.block_sizes)
empty_chains += not encoded.block_sizes
split[bucket].append((encoded.tokens, encoded.regions))
args.output_dir.mkdir(parents=True, exist_ok=True)
for name, rows in split.items():
if not rows:
raise ValueError(f'no examples in the {name} split')
_write_packed(
args.output_dir / f'{name}-adaptive.bin',
np.stack([tokens for tokens, _ in rows]),
np.stack([regions for _, regions in rows]),
layout='adaptive', spec=spec, tokenizer_path=args.tokenizer, tokenizer=tokenizer,
extra_metadata={
'sizes': list(spec.sizes),
# reasoning_train resolves the adaptive control ids from the pack, not the
# tokenizer, and refuses a pack without them.
'size_token_ids': size_token_ids(tokenizer, spec.sizes),
'source': 'general-chat',
},
)
print(f'{name}: {len(rows):,} examples -> {args.output_dir}')
total = sum(len(rows) for rows in split.values())
print(f'{len(conversations):,} conversations, {len(documents):,} documents, '
f'{dropped:,} dropped at encode ({dropped / max(1, dropped + total):.1%})')
print(f'examples answering with no thinking block: {empty_chains:,} '
f'({empty_chains / max(1, total):.1%})')
print('block sizes: ' + ', '.join(f'{size}:{count:,}' for size, count in sorted(blocks.items())))
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
sub = parser.add_subparsers(dest='command', required=True)
prep = sub.add_parser('prepare', help='pack the corpus into the adaptive layout')
prep.add_argument('--inputs', type=Path, nargs='+', required=True)
prep.add_argument('--tokenizer', type=Path,
default=Path('artifacts/tokenizer-qwen3-adaptive.json'))
prep.add_argument('--output-dir', type=Path, required=True)
prep.add_argument('--seq-len', type=int, default=2048)
prep.add_argument('--sizes', type=int, nargs='+', default=[32, 64, 128])
prep.add_argument('--max-slots', type=int, default=40)
prep.add_argument('--keep-messages', type=int, default=KEEP_MESSAGES)
prep.add_argument('--cap', nargs='*', metavar='SOURCE=N',
help='keep at most N conversations from a source, e.g. ground-coqa=2273; '
'weight in the mix is examples, and sources differ ~10x in examples '
'per conversation')
prep.add_argument('--val-fraction', type=float, default=0.02)
prep.add_argument('--seed', type=int, default=1337)
args = parser.parse_args()
prepare(args)
if __name__ == '__main__':
main()
|