| |
| """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: |
| 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}) |
| 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) |
|
|
