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