File size: 2,953 Bytes
00bb63e | 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 | """Render the compaction dataset for ReAligned-Qwen3.5-4B.
Two things this has to get right, both learned on sibling runs in this project.
**Thinking is closed in the prompt.** Qwen3.5's template opens `<think>` on every generation.
Rendered naively the target would begin `\n</think>\n\n`, teaching the model to emit a closing think
tag as the first thing in its answer — the broken hybrid-thinking gate. This template does honour
`enable_thinking=False`, which emits the empty block as part of the prompt, where the trainer masks
it. The summariser should answer directly anyway: egirl calls it for a summary, not a deliberation.
**Over-length rows are dropped, not truncated.** The trainer truncates prompts from the right,
which here removes the most recent messages in the window — the current state, the part a summary
most needs. A row that cannot fit is worth less than a row that teaches summarising from a
mutilated transcript.
"""
import argparse
import json
import statistics as st
from transformers import AutoTokenizer
MODEL = "Lazarus-Ai/ReAligned-Qwen3.5-4B"
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--max-tokens", type=int, default=16384)
args = ap.parse_args()
tok = AutoTokenizer.from_pretrained(MODEL)
for split in ("train", "val"):
rows = [json.loads(line) for line in open(f"{split}.jsonl")]
out, lens, dropped = [], [], 0
for r in rows:
prompt = tok.apply_chat_template(
[{"role": "system", "content": r["system"]},
{"role": "user", "content": r["prompt"]}],
add_generation_prompt=True, enable_thinking=False, tokenize=False)
response = r["chosen"] + tok.eos_token
n = len(tok(prompt)["input_ids"]) + len(
tok(response, add_special_tokens=False)["input_ids"])
lens.append(n)
if n > args.max_tokens:
dropped += 1
continue
out.append({"prompt": prompt, "chosen": response,
"source_group": r["source_group"]})
with open(f"compaction_{split}.jsonl", "w") as f:
for o in out:
f.write(json.dumps(o) + "\n")
lens.sort()
print(f"{split}: {len(rows)} -> {len(out)} kept, {dropped} over {args.max_tokens}")
print(f" tokens mean {st.mean(lens):.0f} median {lens[len(lens)//2]} "
f"p95 {lens[int(.95*len(lens))]} max {max(lens)}")
for cap in (8192, 12288, 16384, 24576):
print(f" cap {cap:6d} would drop {sum(1 for x in lens if x > cap):4d}/{len(lens)}")
print(f" total tokens: {sum(x for x in lens if x <= args.max_tokens):,}")
sample = json.loads(open("compaction_train.jsonl").readline())
print("\nprompt tail:", repr(sample["prompt"][-70:]))
print("response head:", repr(sample["chosen"][:90]))
if __name__ == "__main__":
main()
|