#!/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)