data_curation/prepare_lightning_opd.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Prepare Lightning OPD parquet from student rollout data.

Phase 1 – tokenize (CPU-friendly):
    Reads student rollout parquet, builds prompt via chat template,
    tokenizes responses, truncates to --max-response-len, writes intermediate
    parquet WITHOUT teacher logprobs.

Phase 2 – precompute teacher logprobs (requires GPU / teacher sglang server):
    Reads the intermediate parquet produced in Phase 1, sends each
    (prompt + response) sequence to the teacher sglang server, stores
    per-token response logprobs back into the metadata, writes the final
    parquet.

Usage (Phase 1, CPU node):
    python3 data_curation/prepare_lightning_opd.py \\
        --tokenizer-path checkpoints/sft \\
        --input-parquet data/rollouts/rollouts.parquet \\
        --output-dir data/lightning_opd

Usage (Phase 2, GPU node with teacher sglang running):
    python3 data_curation/prepare_lightning_opd.py \\
        --tokenizer-path checkpoints/sft \\
        --input-parquet data/rollouts/rollouts.parquet \\
        --output-dir data/lightning_opd \\
        --compute-teacher-logprobs \\
        --teacher-url http://127.0.0.1:13141/generate
"""

import argparse
import asyncio
from pathlib import Path

import aiohttp
import pandas as pd
from transformers import AutoTokenizer
from tqdm import tqdm


def parse_args():
    parser = argparse.ArgumentParser(
        description="Prepare Lightning OPD parquet data (tokenize + optional teacher logprobs)."
    )
    parser.add_argument(
        "--tokenizer-path", type=str, required=True,
        help="Path to HuggingFace tokenizer (e.g. the student SFT checkpoint).",
    )
    parser.add_argument(
        "--input-parquet", type=str, required=True,
        help="Path to student rollout parquet. Expected columns: messages (list[dict]), tokens (int).",
    )
    parser.add_argument(
        "--output-dir", type=str, required=True,
        help="Directory where intermediate and final parquet files are written.",
    )
    parser.add_argument(
        "--max-response-len", type=int, default=4096,
        help="Maximum response token length; longer responses are truncated (default: 4096).",
    )
    parser.add_argument(
        "--compute-teacher-logprobs", action="store_true",
        help="Run Phase 2: compute teacher logprobs via a running sglang server.",
    )
    parser.add_argument(
        "--teacher-url", type=str, default="http://127.0.0.1:13141/generate",
        help="Teacher sglang server URL (default: http://127.0.0.1:13141/generate).",
    )
    parser.add_argument(
        "--concurrency", type=int, default=64,
        help="Number of concurrent requests to teacher sglang server (default: 64).",
    )
    return parser.parse_args()


# ── Phase 1: tokenize ────────────────────────────────────────────────────────

def phase1_tokenize(args, intermediate_path: Path):
    print(f"[Phase 1] Loading tokenizer from {args.tokenizer_path}")
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True)

    print(f"[Phase 1] Loading input parquet: {args.input_parquet}")
    df = pd.read_parquet(args.input_parquet)
    print(f"[Phase 1] Total rows: {len(df)}")

    rows_out = []
    truncated = 0
    skipped = 0

    for row in tqdm(df.itertuples(), total=len(df), desc="Tokenizing"):
        messages = row.messages

        user_messages = [m for m in messages if m["role"] != "assistant"]
        prompt_str = tokenizer.apply_chat_template(
            user_messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
        )

        assistant_msg = None
        for msg in messages:
            if msg["role"] == "assistant":
                assistant_msg = msg["content"]
                break
        if assistant_msg is None:
            skipped += 1
            continue

        response_ids = tokenizer.encode(assistant_msg, add_special_tokens=False)

        if len(response_ids) > args.max_response_len:
            truncated += 1
            response_ids = response_ids[:args.max_response_len]
            assistant_msg = tokenizer.decode(response_ids, skip_special_tokens=False)

        rows_out.append({
            "prompt": prompt_str,
            "label": "0",
            "metadata": {
                "is_lightning_opd": True,
                "response_tokens": response_ids,
                "loss_mask": [1] * len(response_ids),
                "response": assistant_msg,
            },
        })

    print(f"[Phase 1] Rows written: {len(rows_out)}, "
          f"truncated to {args.max_response_len}: {truncated}, skipped: {skipped}")
    df_out = pd.DataFrame(rows_out)
    intermediate_path.parent.mkdir(parents=True, exist_ok=True)
    df_out.to_parquet(intermediate_path, index=False)
    print(f"[Phase 1] Saved to {intermediate_path}")


# ── Phase 2: precompute teacher logprobs ─────────────────────────────────────

async def _fetch_logprobs(
    session: aiohttp.ClientSession,
    teacher_url: str,
    full_ids: list[int],
    response_len: int,
) -> list[float]:
    """Call teacher sglang server and return per-token logprobs for the response portion."""
    payload = {
        "input_ids": full_ids,
        "sampling_params": {
            "temperature": 0,
            "max_new_tokens": 0,
            "skip_special_tokens": False,
        },
        "return_logprob": True,
        "logprob_start_len": 0,
    }
    async with session.post(teacher_url, json=payload) as resp:
        resp.raise_for_status()
        ret = await resp.json()

    all_lps = ret["meta_info"]["input_token_logprobs"]
    response_lps = [float(item[0]) for item in all_lps[1:]][-response_len:]
    assert len(response_lps) == response_len, (
        f"Expected {response_len} logprobs, got {len(response_lps)}"
    )
    return response_lps


async def _process_all(args, tokenizer, rows: list[dict]) -> list[list[float]]:
    """Process all rows concurrently with a live progress bar, preserving order."""
    semaphore = asyncio.Semaphore(args.concurrency)
    connector = aiohttp.TCPConnector(limit=args.concurrency)
    results = [None] * len(rows)

    async def bounded_fetch(idx: int, full_ids: list[int], response_len: int):
        async with semaphore:
            result = await _fetch_logprobs(session, args.teacher_url, full_ids, response_len)
        results[idx] = result
        pbar.update(1)

    async with aiohttp.ClientSession(connector=connector) as session:
        with tqdm(total=len(rows), desc="[Phase 2] Teacher logprobs") as pbar:
            tasks = []
            for idx, row in enumerate(rows):
                meta = row["metadata"]
                prompt_ids = tokenizer.encode(row["prompt"], add_special_tokens=False)
                response_ids = [int(x) for x in meta["response_tokens"]]
                full_ids = prompt_ids + response_ids
                tasks.append(bounded_fetch(idx, full_ids, len(response_ids)))
            await asyncio.gather(*tasks)

    return results


def phase2_logprobs(args, intermediate_path: Path, output_path: Path):
    print(f"[Phase 2] Loading intermediate parquet: {intermediate_path}")
    df = pd.read_parquet(intermediate_path)
    rows = df.to_dict(orient="records")
    print(f"[Phase 2] Total rows: {len(rows)}")

    print(f"[Phase 2] Loading tokenizer from {args.tokenizer_path}")
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True)

    print(f"[Phase 2] Computing teacher logprobs via {args.teacher_url} "
          f"(concurrency={args.concurrency})")
    all_logprobs = asyncio.run(_process_all(args, tokenizer, rows))

    for row, lps in zip(rows, all_logprobs):
        row["metadata"]["teacher_log_probs"] = lps

    df_out = pd.DataFrame(rows)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    df_out.to_parquet(output_path, index=False)
    print(f"[Phase 2] Saved to {output_path}")

    # Sanity check
    df_check = pd.read_parquet(output_path)
    row0 = df_check.iloc[0]
    meta = row0["metadata"]
    print("\n[Phase 2] Sanity check row 0:")
    print(f"  prompt[:80]:              {row0['prompt'][:80]}")
    print(f"  label:                    {row0['label']}")
    print(f"  len(response_tokens):     {len(meta['response_tokens'])}")
    print(f"  len(teacher_log_probs):   {len(meta['teacher_log_probs'])}")
    print(f"  teacher_log_probs[:5]:    {meta['teacher_log_probs'][:5]}")


# ── Entry point ───────────────────────────────────────────────────────────────

def main():
    args = parse_args()

    output_dir = Path(args.output_dir)
    input_stem = Path(args.input_parquet).stem
    intermediate_path = output_dir / f"{input_stem}-lightning-opd.parquet"
    output_path = output_dir / f"{input_stem}-lightning-opd-precomputed.parquet"

    if args.compute_teacher_logprobs:
        if not intermediate_path.exists():
            print("[INFO] Intermediate parquet not found, running Phase 1 first.")
            phase1_tokenize(args, intermediate_path)
        phase2_logprobs(args, intermediate_path, output_path)
    else:
        phase1_tokenize(args, intermediate_path)
        print(f"\n[INFO] To add teacher logprobs, re-run with --compute-teacher-logprobs "
              f"after starting the teacher sglang server.")


if __name__ == "__main__":
    main()