Kalso42/WorldModelForMaze

1import os
2import sys
3import pickle
4import numpy as np
5import re
6import argparse
7from tqdm import tqdm
8
9# Ensure project root is importable when running this script directly
10sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
11from cli_utils import parse_count, format_count
12
13parser = argparse.ArgumentParser(description='Create the multitask dataset based on the given parameters.')
14parser.add_argument('--num_nodes', type=int, default=100, help='Number of nodes in the graph')
15parser.add_argument('--num_train_dataset', type=parse_count, default='10M',
16                    help='Number of training data entries to use (supports K/M/B, default: 50000)')
17parser.add_argument('--num_test_dataset', type=parse_count, default=10000,
18                    help='Number of test data entries to use (supports K/M/B, default: 10000)')
19parser.add_argument('--tasks', type=str, default='H1',
20                    help='Task specification (e.g., A1, A1B1, A3B2, A1D1F1). Default: A1')
21parser.add_argument('--CL', action=argparse.BooleanOptionalAction, default=False,
22                    help='Enable Task C label mode (append node labels after L/R turns) and add _CL_ in filenames')
23parser.add_argument('--path_type', type=str, default='RWs', choices=['RWc', 'RWa', 'RWs'],
24                    help='Path generation type: RWc (random walk with cycles), RWa (random walk acyclic, default), RWs (single source random walk).')
25# Arguments for task tag handling
26parser.add_argument('--no_task_tag', action='store_true', default=False,
27                    help='Data files do not contain task identifiers (A, B, C, etc.). When enabled, task tokens will not be included in vocabulary and data parsing will skip task tags.')
28parser.add_argument('--both', action='store_true', default=False,
29                    help='Process both versions (with and without task tags). When set, --no_task_tag is ignored and two sets of bin files and meta files are produced.')
30parser.add_argument('--num_labels', type=int, default=10,
31                    help='Number of distinct node labels (default: 10). Must match the value used in data generation.')
32parser.add_argument('--num_workers', type=int, default=1,
33                    help='Number of parallel worker processes for encoding (default: 1 = serial). Uses fork; requires Linux/macOS.')
34args = parser.parse_args()
35
36num_nodes = args.num_nodes
37tasks_str = args.tasks
38tasks_tag_base = f"{tasks_str}_CL" if args.CL else tasks_str
39# Add path type tag for filenames
40path_type_tag = args.path_type
41tasks_tag_base = f"{tasks_tag_base}_{path_type_tag}"
42# Include num_labels in tag when non-default (match create_multitask_maze.py)
43if args.num_labels != 10:
44    tasks_tag_base = f"{tasks_tag_base}_L{args.num_labels}"
45
46train_label = format_count(args.num_train_dataset)
47test_label = format_count(args.num_test_dataset)
48num_labels = args.num_labels
49
50def first_existing(paths):
51    for p in paths:
52        if os.path.exists(p):
53            return p
54    return paths[0]
55
56
57def process_data_for_tag_mode(no_task_tag_mode, tasks_tag_suffix=""):
58    """Process data for a specific task tag mode."""
59    # Construct tasks_tag for this mode
60    if no_task_tag_mode:
61        tasks_tag = f"{tasks_tag_base}_NT"
62        if tasks_tag_suffix:
63            tasks_tag = f"{tasks_tag}_{tasks_tag_suffix}"
64    else:
65        tasks_tag = tasks_tag_base
66        if tasks_tag_suffix:
67            tasks_tag = f"{tasks_tag}_{tasks_tag_suffix}"
68
69    # Find input files
70    train_file_path = first_existing([
71        os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/train_{tasks_tag}_{train_label}.txt'),
72        os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/train_{tasks_tag}_{args.num_train_dataset}.txt'),
73        os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/train_{tasks_str}_{train_label}.txt'),
74        os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/train_{tasks_str}_{args.num_train_dataset}.txt'),
75    ])
76    val_file_path = first_existing([
77        os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/test_{tasks_tag}_{test_label}.txt'),
78        os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/test_{tasks_tag}_{args.num_test_dataset}.txt'),
79        os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/test_{tasks_str}_{test_label}.txt'),
80        os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/test_{tasks_str}_{args.num_test_dataset}.txt'),
81    ])
82
83    print(f"\nProcessing mode: {'Without task tags' if no_task_tag_mode else 'With task tags'}")
84    print(f"Training file: {train_file_path}")
85    print(f"Test file: {val_file_path}")
86
87    with open(train_file_path, 'r') as f:
88        train_data = f.read()
89    print(f"length of train dataset in characters: {len(train_data):,}")
90
91    with open(val_file_path, 'r') as f:
92        val_data = f.read()
93    print(f"length of val dataset in characters: {len(val_data):,}")
94
95    all_data = train_data + val_data
96
97    chars = sorted(list(find_characters(all_data)))
98    direction_tokens = ['N', 'S', 'E', 'W', 'L', 'R', 'F', 'T']
99    # Only include task tokens if no_task_tag_mode is False
100    task_tokens = [] if no_task_tag_mode else ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']
101    label_tokens = ([chr(ord('a') + i) for i in range(num_labels)] if num_labels <= 26
102                    else [f'l{i}' for i in range(num_labels)])
103    label_tokens.append('/')  # separator token for neighbor labels
104    special_tokens = [':']
105    # Adjust vocab_size calculation
106    vocab_size = num_nodes + 2 + len(direction_tokens) + len(task_tokens) + len(label_tokens) + len(special_tokens)
107    print("all the unique characters:", ' '.join(chars))
108    print(f"vocab size: {vocab_size:,}")
109    print(f"No task tag mode: {'Enabled' if no_task_tag_mode else 'Disabled'}")
110
111    stoi = {}
112    itos = {}
113
114    for i in range(num_nodes):
115        stoi[str(i)] = i + 2
116        itos[i + 2] = str(i)
117
118    base = 2 + num_nodes
119    for idx, tok in enumerate(direction_tokens):
120        stoi[tok] = base + idx
121        itos[base + idx] = tok
122
123    # Only add task tokens to vocabulary if no_task_tag_mode is False
124    if not no_task_tag_mode:
125        base = 2 + num_nodes + len(direction_tokens)
126        for idx, tok in enumerate(task_tokens):
127            stoi[tok] = base + idx
128            itos[base + idx] = tok
129        base = 2 + num_nodes + len(direction_tokens) + len(task_tokens)
130    else:
131        base = 2 + num_nodes + len(direction_tokens)
132
133    for idx, tok in enumerate(label_tokens):
134        stoi[tok] = base + idx
135        itos[base + idx] = tok
136
137    base = base + len(label_tokens)
138    for idx, tok in enumerate(special_tokens):
139        stoi[tok] = base + idx
140        itos[base + idx] = tok
141
142    stoi['[PAD]'] = 0
143    itos[0] = '[PAD]'
144    stoi['\n'] = 1
145    itos[1] = '\n'
146
147    def encode(s):
148        ss = s.split(" ")
149        return [stoi[ch] for ch in ss]
150
151    def decode(l):
152        return ' '.join(itos[i] for i in l)
153
154    # Calculate block_size with theoretical minimum
155    n = int(num_nodes ** 0.5)  # grid size
156    if no_task_tag_mode:
157        theoretical_min_tokens = num_nodes + 3
158    else:
159        theoretical_min_tokens = num_nodes + 4
160
161    theoretical_min_block_size = (theoretical_min_tokens // 32 + 1) * 32
162
163    nw = args.num_workers
164
165    def get_block_size(s, desc="scan block size"):
166        split_text = s.split('\n')
167        if nw and nw > 1 and len(split_text) > 0:
168            import multiprocessing as mp
169            chunk_size = max(1, min(20000, len(split_text) // (nw * 50) or 1))
170            chunks = _chunk_list(split_text, chunk_size)
171            ctx = mp.get_context('fork')
172            bs = 0
173            with ctx.Pool(processes=nw) as pool:
174                with tqdm(total=len(split_text), desc=desc) as pbar:
175                    for r in pool.imap_unordered(_prep_max_len_batch, chunks):
176                        if r > bs:
177                            bs = r
178                        pbar.update(chunk_size if pbar.n + chunk_size <= len(split_text)
179                                    else len(split_text) - pbar.n)
180            return bs
181        # Serial
182        bs = 0
183        for st in tqdm(split_text, desc=desc):
184            if st != "":
185                enc_str = encode(st) + [1]
186                bs = max(bs, len(enc_str))
187        return bs
188
189    data_block_size = (max(get_block_size(train_data, desc="scan train block size"),
190                           get_block_size(val_data, desc="scan val block size")) // 32 + 1) * 32
191    block_size = max(theoretical_min_block_size, data_block_size)
192    print(
193        f"the block size is {block_size} (theoretical min: {theoretical_min_block_size}, data-based: {data_block_size})")
194
195    def process_reasoning(s, desc="encode"):
196        split_text = s.split('\n')
197        if nw and nw > 1 and len(split_text) > 0:
198            import multiprocessing as mp
199            chunk_size = max(1, min(10000, len(split_text) // (nw * 100) or 1))
200            chunks = _chunk_list(split_text, chunk_size)
201            ctx = mp.get_context('fork')
202            ret = []
203            with ctx.Pool(processes=nw,
204                          initializer=_prep_worker_init,
205                          initargs=(stoi, block_size)) as pool:
206                with tqdm(total=len(split_text), desc=desc) as pbar:
207                    # Use imap (ordered) to preserve original line order in output.
208                    for i, r in enumerate(pool.imap(_prep_encode_batch, chunks)):
209                        ret.extend(r)
210                        step = len(chunks[i])
211                        pbar.update(step)
212            return ret
213        # Serial
214        ret = []
215        for st in tqdm(split_text, desc=desc):
216            if st != "":
217                enc_str = encode(st) + [1]
218                ret += enc_str + [0] * (block_size + 1 - len(enc_str))
219        return ret
220
221    train_ids = process_reasoning(train_data, desc="encode train")
222    val_ids = process_reasoning(val_data, desc="encode val")
223
224    print(f"train has {len(train_ids):,} tokens")
225    print(f"val has {len(val_ids):,} tokens")
226
227    train_ids = np.array(train_ids, dtype=np.uint16)
228    val_ids = np.array(val_ids, dtype=np.uint16)
229
230    # Save bin files with appropriate tag
231    train_ids.tofile(os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/train_{tasks_tag}_{train_label}.bin'))
232    val_ids.tofile(os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/val_{tasks_tag}_{test_label}.bin'))
233
234    unreachable = 'x' in chars
235    simple_format = ':' not in chars
236
237    meta = {
238        'unreachable': unreachable,
239        'simple_format': simple_format,
240        'block_size': block_size,
241        'vocab_size': vocab_size,
242        'itos': itos,
243        'stoi': stoi,
244        'no_task_tag': no_task_tag_mode,
245    }
246
247    print(stoi)
248    print(itos)
249    with open(os.path.join(os.path.dirname(__file__), f'{args.num_nodes}/meta_{tasks_tag}.pkl'), 'wb') as f:
250        pickle.dump(meta, f)
251
252    print(f"Saved files with tag: {tasks_tag}")
253    return tasks_tag
254
255
256def find_characters(data_string):
257    pattern = r'\d+|\D'
258    matches = re.findall(pattern, data_string)
259    return set(matches)
260
261
262# ---- Parallel worker helpers (must be at module scope for pickling) ----
263_W_STOI = None
264_W_BLOCK_SIZE = None
265
266
267def _prep_worker_init(stoi_arg, block_size_arg):
268    global _W_STOI, _W_BLOCK_SIZE
269    _W_STOI = stoi_arg
270    _W_BLOCK_SIZE = block_size_arg
271
272
273def _prep_max_len_batch(lines):
274    """Return the max encoded length (tokens + EOL) over a batch of lines."""
275    bs = 0
276    for st in lines:
277        if st == "":
278            continue
279        # encoded length = number of space-separated tokens + 1 (EOL token)
280        L = st.count(" ") + 2
281        if L > bs:
282            bs = L
283    return bs
284
285
286def _prep_encode_batch(lines):
287    """Encode + pad a batch of lines; returns a flat list of token IDs."""
288    stoi = _W_STOI
289    bs1 = _W_BLOCK_SIZE + 1
290    out = []
291    for st in lines:
292        if st == "":
293            continue
294        enc = [stoi[ch] for ch in st.split(" ")]
295        enc.append(1)
296        out.extend(enc)
297        out.extend([0] * (bs1 - len(enc)))
298    return out
299
300
301def _chunk_list(lst, chunk_size):
302    return [lst[i:i + chunk_size] for i in range(0, len(lst), chunk_size)]
303
304
305# Main execution
306if args.both:
307    # Process both versions
308    print("=" * 60)
309    print("Generating both with-tag and without-tag versions")
310    print("=" * 60)
311
312    # Process with task tags
313    process_data_for_tag_mode(no_task_tag_mode=False)
314
315    # Process without task tags
316    process_data_for_tag_mode(no_task_tag_mode=True)
317
318    print("=" * 60)
319    print("Successfully generated both with-tag and without-tag datasets.")
320    print("=" * 60)
321else:
322    # Original logic: generate only one version based on no_task_tag
323    process_data_for_tag_mode(no_task_tag_mode=args.no_task_tag)