scripts/convert_data.py

#!/usr/bin/env python3
"""
convert_data.py: turn local ChatML datasets into tokenized train.pt/val.pt
for train_chat.py.

Accepts .jsonl, .json, and .parquet files (auto-detected by extension), and
these record shapes (one conversation per record):

    {"messages": [{"role": "user", "content": ...}, ...]}   # HF ChatML style
    [{"role": ..., "content": ...}, ...]                     # bare message list
    {"conversation": [...]} or {"chat": [...]}               # aliases
    {"instruction": ..., "input": ..., "output": ...}        # alpaca style (converted)
    {"data": [...]} / {"conversations": [...]}               # JSON containers of the above

Parquet rows may store the messages column as a list of dicts or as a JSON
string (both work).

Usage:
    python3 convert_data.py --model-path . --input my_data.jsonl \
        --output data/my_data
    python3 convert_data.py --model-path . --input a.jsonl b.jsonl \
        --output data/mixed
    python3 convert_data.py --model-path . --input ./folder \
        --output data/folder

Then train:
    python3 train_chat.py --model-path . --data-dir data/my_data \
        --output-dir my_checkpoints --lr 7e-5 --epochs 3
"""

import argparse
import json
import sys
from pathlib import Path

import numpy as np
from transformers import AutoTokenizer

sys.path.insert(0, str(Path(__file__).resolve().parent))

from prepare_data import (  # noqa: E402
    add_chat_tokens,
    build_conv_segments,
    save_dataset,
    tokenize_and_split,
)

SUPPORTED_EXTS = {".jsonl", ".json", ".parquet"}
CONTAINER_KEYS = ("conversations", "data", "rows", "examples")
MESSAGE_KEYS = ("messages", "conversation", "chat")


# ---------------------------------------------------------------------------
# File reading
# ---------------------------------------------------------------------------

def collect_files(paths):
    """Expand --input args (files and/or dirs) into a sorted file list."""
    files = []
    for p in paths:
        p = Path(p)
        if p.is_dir():
            files.extend(f for f in sorted(p.iterdir())
                         if f.suffix.lower() in SUPPORTED_EXTS)
        elif p.suffix.lower() in SUPPORTED_EXTS:
            files.append(p)
        else:
            print(f"[!] Skipping unsupported file: {p} (want .jsonl/.json/.parquet)")
    return files


def iter_records(path):
    """Yield one raw record (dict or list) per conversation from a file."""
    ext = path.suffix.lower()
    if ext == ".jsonl":
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                yield json.loads(line)
    elif ext == ".json":
        with open(path) as f:
            obj = json.load(f)
        if isinstance(obj, list):
            yield from obj
        elif isinstance(obj, dict):
            for key in CONTAINER_KEYS:
                if isinstance(obj.get(key), list):
                    yield from obj[key]
                    return
            yield obj  # single-conversation file
    elif ext == ".parquet":
        import pandas as pd  # lazy: only needed for parquet
        df = pd.read_parquet(path)
        for _, row in df.iterrows():
            yield dict(row)
    else:
        raise ValueError(f"Unsupported extension: {path}")


def normalize_record(rec):
    """Turn one record into a list of {role, content} messages, or None."""
    if isinstance(rec, list):
        msgs = [m for m in rec
                if isinstance(m, dict) and m.get("content")]
        return msgs or None

    if not isinstance(rec, dict):
        return None

    # ChatML-style keys (value may be a list of dicts, a numpy array of dicts
    # from parquet round-trips, or a JSON string)
    for key in MESSAGE_KEYS:
        v = rec.get(key)
        if isinstance(v, str):
            try:
                v = json.loads(v)
            except json.JSONDecodeError:
                continue
        if isinstance(v, (list, np.ndarray)):
            msgs = [m for m in v
                    if isinstance(m, dict) and m.get("content")]
            if msgs:
                return msgs

    # Alpaca-style record: instruction / input / output
    if rec.get("instruction") and rec.get("output"):
        user = rec["instruction"]
        if rec.get("input"):
            user += f"\n\n{rec['input']}"
        return [{"role": "user", "content": user},
                {"role": "assistant", "content": rec["output"]}]

    return None


def load_convs_from_files(files):
    """Build (roles, segs) conversations from all files."""
    convs = []
    for path in files:
        n_before = len(convs)
        for rec in iter_records(path):
            msgs = normalize_record(rec)
            if msgs is None:
                continue
            segs, roles = build_conv_segments(msgs)
            if not segs:
                continue
            convs.append((roles, segs))
        print(f"[*] {path.name}: {len(convs) - n_before} conversations")
    return convs


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Convert local ChatML data (jsonl/json/parquet) to tokenized .pt")
    parser.add_argument("--model-path", default=".",
                        help="Dir with tokenizer.json (the release dir works)")
    parser.add_argument("--input", nargs="+", required=True,
                        help="File(s) or dir(s): .jsonl, .json, .parquet")
    parser.add_argument("--output", default="data/converted",
                        help="Output dir (train.pt, val.pt, tokenizer/, stats.json)")
    parser.add_argument("--val-size", type=int, default=500, help="Held-out examples")
    parser.add_argument("--max-len", type=int, default=1024, help="Max tokens per example")
    parser.add_argument("--max-resp-tokens", type=int, default=256,
                        help="Cap on target response tokens (keeps <|im_end|>)")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)

    print(f"[*] Loading tokenizer from {args.model_path}")
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    add_chat_tokens(tokenizer)
    tokenizer.save_pretrained(out / "tokenizer")

    files = collect_files(args.input)
    if not files:
        print("[!] No .jsonl/.json/.parquet files found in the inputs.")
        sys.exit(1)
    print(f"[*] Files: {', '.join(f.name for f in files)}")

    convs = load_convs_from_files(files)
    if not convs:
        print("[!] No conversations parsed (check the record shapes in the docstring).")
        sys.exit(1)
    print(f"[*] Total: {len(convs)} conversations")

    # Auto-scale the val split: never take everything for small datasets
    val_size = min(args.val_size, max(1, len(convs) // 10))
    if val_size != args.val_size:
        print(f"[*] Small dataset: using val_size={val_size}")

    train_examples, val_examples, skipped, dupes = tokenize_and_split(
        convs, set(), tokenizer, val_size, args.max_len,
        args.max_resp_tokens, args.seed)

    save_dataset(out, train_examples, val_examples, tokenizer, skipped, dupes)


if __name__ == "__main__":
    main()