| """Build v6 SFT mix: FULL forensic set + chat/persona/tool/distill/SOP extras |
| + raw TinyStories replay for fluency. Memory-frugal (streams the 2.1GB txt).""" |
| import json, random |
| from collections import Counter |
| from pathlib import Path |
| from data.tokenizer import load_tokenizer |
|
|
| rng = random.Random(20260802) |
| SEQ = 256 |
| OUT = Path("data/sft_mix_v6.jsonl") |
|
|
| def load(p): |
| return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()] |
|
|
| def dedupe(rows): |
| seen, out = set(), [] |
| for r in rows: |
| k = (r.get("persona", "analyst"), r.get("user", "")[:180]) |
| if k in seen: |
| continue |
| seen.add(k); out.append(r) |
| return out |
|
|
| def reservoir_sample(path, n, chunk=10000): |
| rng2 = random.Random(7) |
| keep = [] |
| with open(path, encoding="utf-8") as f: |
| seen = 0 |
| for line in f: |
| s = line.strip() |
| if not s: |
| continue |
| seen += 1 |
| if len(keep) < n: |
| keep.append(s) |
| else: |
| j = rng2.randrange(seen) |
| if j < n: |
| keep[j] = s |
| return keep |
|
|
| def visible(row, tok, u_id, a_id, eot): |
| if "raw" in row: |
| return True |
| if not row.get("user") or not row.get("assistant"): |
| return False |
| p = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>", "none": ""}.get(row.get("persona", "analyst"), "<|analyst|>") |
| p_ids = tok.encode(p).ids if row.get("persona", "analyst") != "none" else [] |
| ids = p_ids + [u_id] + tok.encode(row["user"]).ids + [a_id] + tok.encode(row["assistant"]).ids + [eot] |
| return len(ids) <= SEQ |
|
|
| def main(): |
| tok = load_tokenizer("data/tokenizer.json") |
| u_id = tok.token_to_id("<|user|>"); a_id = tok.token_to_id("<|assistant|>"); eot = tok.token_to_id("<|endoftext|>") |
| v3 = dedupe(load("data/sft_mix_v3.jsonl")) |
| truth = [r for r in v3 if r.get("user", "").startswith("Answer truthfully:")] |
| chatish = [r for r in v3 if r.get("persona") == "analyst" and "raw" not in r and len(r.get("user", "")) < 90] |
|
|
| mix = [] |
| mix += load("data/general_chat.jsonl") |
| mix += load("data/persona_dialogue.jsonl") |
| mix += load("data/tool_use.jsonl") |
| mix += rng.sample(truth, 40) |
| mix += rng.sample(chatish, 80) |
| mix += rng.sample(load("data/sft_distill_mix.jsonl"), 160) |
| mix += rng.sample(load("data/sft_sop_mix.jsonl"), 120) |
| mix += dedupe(load("data/sft_forensic.jsonl")) |
|
|
| clean = [r for r in dedupe(mix) if visible(r, tok, u_id, a_id, eot)] |
| for ln in reservoir_sample("data/TinyStoriesV2-GPT4-train.txt", 700): |
| clean.append({"raw": ln, "persona": "none"}) |
| rng.shuffle(clean) |
| with open(OUT, "w", encoding="utf-8") as f: |
| for r in clean: |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") |
| print("total", len(clean), dict(Counter(r.get("persona", "?") for r in clean)), flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
