zeyuzy/DLLM-Planing-Task

1"""2Generic fixed-length character-level preprocessing for reasoning tasks.3 4This script builds the shared vocabulary, serializes each example as5[quiz_padded][response_padded], writes flattened uint16 train/val binaries, and6stores metadata needed by AR/Tom-CAT evaluation and training entrypoints.7"""8 9import argparse10import json11import os12import pickle13import random14import string15 16import numpy as np17 18def parse_args():19    parser = argparse.ArgumentParser(description='Generic data preparation for fixed-length reasoning tasks')20 21    # Input/output paths.22    parser.add_argument('--data_path', type=str, default='data/3sat7_train.jsonl', help='Path to cd5 JSONL file')23    parser.add_argument('--out_dir', type=str, default='data/sat/3sat7/k8', help='Output directory for processed data')24 25    # Dataset schema.26    parser.add_argument('--input_key', type=str, default='input', help='JSON key for the prompt/quiz')27    parser.add_argument('--output_key', type=str, default='output', help='JSON key for the completion/response')28    parser.add_argument('--meta_name', type=str, default='meta.pkl', help='Name of the output metadata file')29 30    # Train/validation split.31    parser.add_argument('--val_ratio', type=float, default=0.1, help='Ratio of data to use for validation (e.g., 0.1 for 10%)')32    parser.add_argument('--seed', type=int, default=42, help='Random seed for shuffling and splitting')33 34    # Vocabulary control.35    parser.add_argument('--custom_vocab', type=str, default='', help='Comma-separated custom characters. If empty, auto-scans the dataset.')36 37    return parser.parse_args()38 39def main():40    args = parse_args()41    os.makedirs(args.out_dir, exist_ok=True)42    random.seed(args.seed)43 44    # 1. Read the JSONL source file.45    print(f"Loading data from {args.data_path}...")46    data = []47    with open(args.data_path, 'r', encoding='utf-8') as f:48        for line in f:49            if line.strip():50                data.append(json.loads(line.strip()))51    52    print(f"Loaded {len(data)} samples.")53 54    # 2. Build the vocabulary.55    special_tokens = ["<PAD>", "<SEP>", "<EOS>", "<MASK>", "$"]56 57    # Global base character set shared across common reasoning tasks.58    global_base_chars = [59        "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",60        ",", "+", "-", "/", "=", "*",61    ] + list(string.ascii_lowercase)62    63    if args.custom_vocab:64        print("Using custom vocabulary...")65        base_chars = [c.strip() for c in args.custom_vocab.split(',') if c.strip()]66    else:67        print("Using Global Unified Vocabulary...")68        base_chars = global_base_chars.copy()69 70        # Scan the dataset to append any unseen characters beyond the shared base set.71        all_chars = set()72        for sample in data:73            all_chars.update(list(str(sample.get(args.input_key, ''))))74            all_chars.update(list(str(sample.get(args.output_key, ''))))75        76        unseen_chars = set(all_chars) - set(base_chars) - set(special_tokens)77        if unseen_chars:78            print(f"⚠️ Notice: Found new characters not in global vocab: {unseen_chars}")79            base_chars.extend(sorted(list(unseen_chars)))80 81    # Keep special tokens unique and place them at the end of the vocabulary.82    chars = [c for c in base_chars if c not in special_tokens] + special_tokens83    84    stoi = {ch: i for i, ch in enumerate(chars)}85    itos = {i: ch for i, ch in enumerate(chars)}86    vocab_size = len(chars)87    if vocab_size >= 65536:88        raise ValueError(f"vocab_size={vocab_size} exceeds uint16 capacity")89    print(f"Vocab size: {vocab_size}")90 91    def encode(s):92        return [stoi[c] for c in str(s)]93 94    # 3. Compute the maximum raw input/output lengths.95    max_quiz_len = 096    max_response_len = 097    for sample in data:98        quiz = str(sample.get(args.input_key, ''))99        response = str(sample.get(args.output_key, ''))100        max_quiz_len = max(max_quiz_len, len(quiz))101        max_response_len = max(max_response_len, len(response))102 103    quiz_size = max_quiz_len + 1      # +1 for <SEP>104    response_size = max_response_len + 1  # +1 for <EOS>105    data_size = quiz_size + response_size106 107    print(f"max_quiz_len={max_quiz_len}, max_response_len={max_response_len}")108    print(f"quiz_size={quiz_size}, response_size={response_size}, data_size={data_size}")109 110    # 4. Dedup by (input, output) FIRST, so the same problem can't land in both111    #    train and val (a content-level leak that inflates val/test metrics). Then112    #    shuffle (fixed seed) and split into disjoint train/validation sets.113    _seen = set()114    _deduped = []115    for _s in data:116        _key = (str(_s.get(args.input_key, '')), str(_s.get(args.output_key, '')))117        if _key in _seen:118            continue119        _seen.add(_key)120        _deduped.append(_s)121    if len(_deduped) != len(data):122        print(f"Dedup: {len(data)} -> {len(_deduped)} samples "123              f"({len(data) - len(_deduped)} duplicates removed before split)")124    data = _deduped125 126    random.shuffle(data)127    num_val = int(len(data) * args.val_ratio)128    val_samples = data[:num_val]129    train_samples = data[num_val:]130    print(f"Split: {len(train_samples)} train samples, {len(val_samples)} val samples.")131 132    # 5. Encode samples into the fixed-length serialization:133    #    [quiz chars + PAD ... + SEP][response chars + PAD ... + EOS]134    def process_samples(samples, dataset_name):135        processed_seqs = []136        for idx, sample in enumerate(samples):137            quiz = str(sample.get(args.input_key, ''))138            response = str(sample.get(args.output_key, ''))139            140            quiz_encoded = encode(quiz)141            response_encoded = encode(response)142 143            # Pad the quiz and response to their dataset-wide maximum lengths.144            quiz_padded = quiz_encoded + [stoi["<PAD>"]] * (max_quiz_len - len(quiz_encoded)) + [stoi["<SEP>"]]145            response_padded = response_encoded + [stoi["<PAD>"]] * (max_response_len - len(response_encoded)) + [stoi["<EOS>"]]146            147            seq = quiz_padded + response_padded148            149            if len(seq) != data_size:150                print(f"[{dataset_name}] Skipping invalid sequence at index {idx}: seq_len={len(seq)}, expected={data_size}")151                continue152                153            processed_seqs.extend(seq)154        return processed_seqs155 156    train_data = process_samples(train_samples, "Train")157    val_data = process_samples(val_samples, "Val")158 159    print(f"Raw train tokens: {len(train_data)}")160    print(f"Raw val tokens: {len(val_data)}")161 162    # Print a few decoded examples to verify the packing protocol.163    print("\n" + "="*60)164    print("VERIFICATION: Checking a few processed training samples...")165    print("="*60)166    num_examples_to_print = 3167    if len(train_data) >= data_size * num_examples_to_print:168        for i in range(num_examples_to_print):169            start_idx = i * data_size170            end_idx = start_idx + data_size171            sample_seq = train_data[start_idx:end_idx]172 173            decoded_seq = [itos[token_id] for token_id in sample_seq]174 175            quiz_part = "".join(decoded_seq[:quiz_size])176            resp_part = "".join(decoded_seq[quiz_size:])177 178            print(f"--- Example {i+1} ---")179            print(f"Padded Quiz     (len={len(quiz_part)}): {quiz_part}")180            print(f"Padded Response (len={len(resp_part)}): {resp_part}")181            print()182    else:183        print("Not enough data to print examples.")184    print("="*60 + "\n")185 186    # 6. Truncate any trailing partial example if earlier skips broke alignment.187    def truncate_to_block(data_list, block_size, name):188        remainder = len(data_list) % block_size189        if remainder != 0:190            print(f"Truncating {name} data by {remainder} tokens to align with block size.")191            return data_list[:-remainder]192        return data_list193 194    train_data = truncate_to_block(train_data, data_size, "train")195    val_data = truncate_to_block(val_data, data_size, "val")196 197    # 7. Convert to uint16 arrays and sanity-check the vocabulary range.198    train_bin = np.array(train_data, dtype=np.uint16)199    val_bin = np.array(val_data, dtype=np.uint16)200 201    assert train_bin.max() < vocab_size, f"Dirty data detected! Max token {train_bin.max()} >= vocab_size {vocab_size}"202    if len(val_bin) > 0:203        assert val_bin.max() < vocab_size, f"Dirty data detected! Max token {val_bin.max()} >= vocab_size {vocab_size}"204 205    # Save flattened binary files.206    train_bin.tofile(os.path.join(args.out_dir, 'train.bin'))207    val_bin.tofile(os.path.join(args.out_dir, 'val.bin'))208 209    # 8. Save metadata describing the serialization protocol.210    meta = {211        'format_version': 'fixed_length_char_v1',212        'vocab_size': vocab_size,213        'stoi': stoi,214        'itos': itos,215        'block_size': data_size - 1,216        'quiz_size': quiz_size,217        'response_size': response_size,218        'data_size': data_size,219        'max_quiz_len': max_quiz_len,220        'max_response_len': max_response_len,221        'max_input_len': max_quiz_len,222        'max_output_len': max_response_len,223        'input_key': args.input_key,224        'output_key': args.output_key,225        'special_tokens': special_tokens,226        'pad_token': '<PAD>',227        'sep_token': '<SEP>',228        'eos_token': '<EOS>',229        'mask_token': '<MASK>',230        'dollar_token': '$',231        'tokenizer_type': 'char',232        'serialization': 'quiz_pad_sep + response_pad_eos',233        'dtype': 'uint16',234        'data_path': args.data_path,235        'val_ratio': args.val_ratio,236        'seed': args.seed,237    }238 239    meta_path = os.path.join(args.out_dir, args.meta_name)240    with open(meta_path, 'wb') as f:241        pickle.dump(meta, f)242 243    print(f"✅ Data successfully prepared in '{args.out_dir}'.")244    print(f"   Saved train.bin, val.bin, and {args.meta_name}.")245 246if __name__ == "__main__":247    main()248