File size: 1,741 Bytes
b5989f0 | 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 | #!/usr/bin/env python3
"""Build long-coherent (query, reasoning) pairs for teacher-forced 4000+ training.
Dolphin-R1 reasoning-deepseek, keep items whose reference reasoning is >=4000 tokens (genuine
long coherent CoT). Save query + reasoning text (capped) -> /workspace/long_tf.jsonl."""
import sys, json
from datasets import load_dataset
from transformers import AutoTokenizer
WANT = int(sys.argv[1]) if len(sys.argv) > 1 else 3000
MIN_RSN_TOK = 4000
CAP_TOK = 6144
tok = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")
ds = load_dataset("cognitivecomputations/dolphin-r1", "reasoning-deepseek", split="train", streaming=True)
out, seen = [], 0
for r in ds:
rsn = r.get("reasoning") or ""
if not isinstance(rsn, str) or len(rsn) < MIN_RSN_TOK * 3: # cheap char prefilter
continue
m = r.get("messages"); q = None
if isinstance(m, list):
for t in m:
if (t.get("role") or "").lower() == "user": q = t.get("content"); break
if not q: continue
qn = len(tok(q, add_special_tokens=True).input_ids)
if not (8 <= qn <= 160): continue
a_ids = tok(rsn, add_special_tokens=False).input_ids
if len(a_ids) < MIN_RSN_TOK: continue
a_ids = a_ids[:CAP_TOK]
out.append({"q": q, "a_ids": a_ids}) # pre-tokenized answer (capped)
seen += 1
if len(out) >= WANT: break
if seen % 500 == 0: print(f" kept {len(out)}", flush=True)
import statistics
lens = [len(o["a_ids"]) for o in out]
with open("/workspace/long_tf.jsonl", "w") as f:
for o in out: f.write(json.dumps(o) + "\n")
print(f"LONG POOL: {len(out)} | ans_len min/med/max = {min(lens)}/{int(statistics.median(lens))}/{max(lens)}", flush=True)
print("PREPLONG_DONE", flush=True)
|