| """ARC dataset loader.""" |
| import json |
| import os |
| import random |
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
| from pathlib import Path |
| from typing import Dict, List, Union |
|
|
|
|
| class ARCDataset(Dataset): |
| def __init__(self, data_dir, split="training", max_grid_size=64, num_colors=16, |
| augment=False, mode="pairs"): |
| self.data_dir = Path(data_dir) |
| self.split = split |
| self.max_grid_size = max_grid_size |
| self.num_colors = num_colors |
| self.augment = augment |
| self.mode = mode |
| self.samples = self._load() |
| |
| def _load(self): |
| samples = [] |
| task_dir = self.data_dir / self.split |
| if not task_dir.exists(): |
| raise FileNotFoundError(task_dir) |
| for f in sorted(task_dir.glob("*.json")): |
| with open(f) as fh: |
| task = json.load(fh) |
| tid = f.stem |
| for p in task.get("train", []): |
| samples.append({"task_id": tid, "input": np.array(p["input"], dtype=np.int64), |
| "output": np.array(p["output"], dtype=np.int64), "split": "train"}) |
| for p in task.get("test", []): |
| samples.append({"task_id": tid, "input": np.array(p["input"], dtype=np.int64), |
| "output": np.array(p["output"], dtype=np.int64), "split": "test"}) |
| return samples |
| |
| def _pad(self, g): |
| h, w = g.shape |
| if h > self.max_grid_size or w > self.max_grid_size: |
| g = g[:self.max_grid_size, :self.max_grid_size] |
| h, w = g.shape |
| ph = self.max_grid_size - h |
| pw = self.max_grid_size - w |
| if ph > 0 or pw > 0: |
| g = np.pad(g, ((0, ph), (0, pw)), mode="constant", constant_values=0) |
| return g |
| |
| def _aug(self, inp, out): |
| if not self.augment: |
| return inp, out |
| if random.random() < 0.5: |
| inp, out = np.flip(inp, 0).copy(), np.flip(out, 0).copy() |
| if random.random() < 0.5: |
| inp, out = np.flip(inp, 1).copy(), np.flip(out, 1).copy() |
| k = random.randint(0, 3) |
| if k > 0: |
| inp, out = np.rot90(inp, k=k).copy(), np.rot90(out, k=k).copy() |
| if random.random() < 0.5: |
| colors = np.unique(np.concatenate([inp.flatten(), out.flatten()])) |
| colors = colors[colors > 0] |
| if len(colors) > 1: |
| shuffled = colors.copy() |
| np.random.shuffle(shuffled) |
| cmap = {c: shuffled[i] for i, c in enumerate(colors)} |
| cmap[0] = 0 |
| inp = np.vectorize(cmap.get)(inp) |
| out = np.vectorize(cmap.get)(out) |
| return inp, out |
| |
| def __len__(self): |
| return len(self.samples) |
| |
| def __getitem__(self, idx): |
| s = self.samples[idx] |
| inp, out = self._aug(s["input"], s["output"]) |
| inp = self._pad(inp) |
| out = self._pad(out) |
| return { |
| "task_id": s["task_id"], |
| "input_grid": torch.from_numpy(np.clip(inp, 0, self.num_colors - 1)).long(), |
| "output_grid": torch.from_numpy(np.clip(out, 0, self.num_colors - 1)).long(), |
| "split": s["split"], |
| } |
|
|
|
|
| def collate_fn(batch): |
| return { |
| "task_ids": [b["task_id"] for b in batch], |
| "input_grids": torch.stack([b["input_grid"] for b in batch]), |
| "output_grids": torch.stack([b["output_grid"] for b in batch]), |
| "splits": [b["split"] for b in batch], |
| } |
|
|
|
|
| def load_arc_data(source="local", split="training", cache_dir="./data/arc_agi_1", **kwargs): |
| return ARCDataset(cache_dir, split=split, **kwargs) |
|
|
