Tools for working with the Abstraction & Reasoning Corpus
from .data import Task, load_data, fmt_grid, load_single
from .vis import draw_grid, draw_task, output_drawing
arckit/__init__.py
mxbi/arckit
import sys
import argparse
from .data import load_single
from .vis import draw_task, output_drawing

def taskprint():
    if len(sys.argv) < 2:
        print("Usage: arctask <task_id>")
        return

    task = load_single(sys.argv[1])

    task.show()

def tasksave():
    parser = argparse.ArgumentParser(description='Save a task to a image file.')
    parser.add_argument('task_id', type=str, help='The task id to save. Can either be a task ID or a string e.g. `train0`')
    parser.add_argument('width', type=int, help='The width of the output image', default=20)
    parser.add_argument('height', type=int, help='The height of the output image', default=10)
    parser.add_argument('--output', type=str, help='The output file to save to. Must end in .svg/.pdf/.png. By default, pdf is used.', default=False, required=False)
    args = parser.parse_args()

    task = load_single(args.task_id)

    drawing = draw_task(task, width=args.width, height=args.height)
    out_fn = args.output or f'{task.id}_{drawing.width}x{drawing.height}.pdf'
    print(f'Drawn task {task.id} ({drawing.width}x{drawing.height}), saving to {out_fn}')

    output_drawing(drawing, out_fn)
arckit/cli.py
mxbi/arckit
from glob import glob
import json
import numpy as np
import os
import csv

from rich.traceback import install
install()#(show_locals=True)

from rich import print
from rich.panel import Panel
from rich.table import Table
from rich.text import Text

def idx2chr(idx):
    return chr(idx + 65)

def fmt_grid(grid, colour=True, spaces=True):
    grid_str = []
    if not colour:
        for row in grid:
            if spaces == 'gpt':
                grid_str.append(''.join([' ' + str(x) for x in row]))
            elif spaces:
                grid_str.append(' '.join([str(x) for x in row]))
            else:
                grid_str.append(''.join([str(x) for x in row]))

        return "\n".join(grid_str)
    else:
        if spaces:
            cmap = dict({i: (str(i) + ' ', f"color({i})") for i in range(10)}, **{str(i): (str(i) + ' ', f"color({i})") for i in range(10)})
        else:
            cmap = dict({i: (str(i), f"color({i})") for i in range(10)}, **{str(i): (str(i), f"color({i})") for i in range(10)})

        for row in grid:
            grid_str += [cmap[digit] for digit in row]
            grid_str += ["\n"]

        return Text.assemble(*grid_str[:-1])

class Task:
    def __init__(self, id, train, test, dataset=None):
        self.dataset = dataset
        self.id = id
        self.train = [(np.array(example['input']), np.array(example['output'])) for example in train]
        self.test = [(np.array(example['input']), np.array(example['output'])) for example in test]

        self.color_count = max

    def __lt__(self, other):
        return self.id < other.id

    def __hash__(self):
        return hash(self.id)

    @classmethod
    def from_json(cls, json_file):
        with open(json_file) as f:
            data = json.load(f)

            # train_examples = [(np.array(example['input']), np.array(example['output'])) for example in data['train']]
            # test_examples = [(np.array(example['input']), np.array(example['output'])) for example in data['test']]

            return cls(os.path.basename(json_file)[:-5], data['train'], data['test'])

    def __repr__(self):
        return f"<Task-{self.dataset} {self.id} | {len(self.train)} train | {len(self.test)} test>"

    def show(self, answer=False):
        table = Table(title=repr(self), show_lines=True)

        data = []
        for i, (input, output) in enumerate(self.train):
            data += [fmt_grid(input), fmt_grid(output)]
            ix, iy = input.shape
            ox, oy = output.shape
            table.add_column(f"{idx2chr(i)}-in {ix}x{iy}", justify="center", no_wrap=True)
            table.add_column(f"{idx2chr(i)}-out {ox}x{oy}", justify="center", no_wrap=True)

        data.append('')
        table.add_column("")
        for i, (input, output) in enumerate(self.test):
            table.add_column(f"T{idx2chr(i)}-in", justify="center", header_style="bold", no_wrap=True)
            if answer:
                table.add_column(f"T{idx2chr(i)}-out", justify="center", header_style="bold", no_wrap=True)
                data += [fmt_grid(input), fmt_grid(output)]
            else:
                data += [fmt_grid(input)]

        table.add_row(*data)
        print(table)
        return table

    def to_dict(self):
        return {
            "id": self.id,
            "train": [{"input": input.tolist(), "output": output.tolist()} for input, output in self.train],
            "test": [{"input": input.tolist(), "output": output.tolist()} for input, output in self.test]
        }
    
    def dreamcoder_format(self):
        # TODO: Separate out the train/test data
        # Currently, dreamcoder gets to search on the train examples too
        # This is included to keep compatibility with Simon's results, but should be corrected and written about
        grids = []
        for inp, out in self.train:
            grids += [inp, out]
        for inp, out in self.test:
            grids += [inp, out]
        return {"name": self.id, "grids": grids}

    def scoreA(self, output):
        output = np.array(output).astype(int)
        if output.shape != self.test[0][1].shape:
            return False
        return (output == self.test[0][1]).all()

    def gpt_prompt(self, i_test, mode="chatgpt", include_completion=False, rot90=False, transpose=False, spaces=True):
        if mode == "chatgpt":
            prompt = "We are playing a game which involves transforming an input grid of digits into an output grid of digits. In general, digits form objects in 2D and the task is to perform some spatial transformation of these objects to go from the input grid to the output grid. All the information about the transformation is contained within the input pairs themselves, and your answer will only be correct if the output grid is exactly correct, so this is what I expect from you. I will begin by giving you several examples of input-output pairs. You will then be given a new input grid, and you must provide the corresponding output grid.\n"
        elif mode == "gpt3":
            # prompt = "We are playing a game which involves transforming an input grid of digits into an output grid of digits. Every below pair of grids contains the same transformation. In general, digits form objects in 2D and the task is to perform some spatial transformation of these objects to go from the input tile to the output tile. One such example of tiles is below.\n"
            prompt = "We are playing a game which involves transformaing a 2D input grid of digits into an output grid of digits. Every below pair of grids contains the same transformation (e.g. rotation, symmetry, manipulation of objects). Each Input grid is followed by an Output grid which applies the same transformation as previous Input/Output pairs. One such example is below.\n"
        else:
            raise ValueError(f"Unknown mode: {mode}")

        i = 1
        for input, output in self.train:
            if rot90:
                input = np.rot90(input)
                output = np.rot90(output)
            if transpose:
                input = np.transpose(input)
                output = np.transpose(output)

            prompt += f"""Input {i}: 
{fmt_grid(input, colour=False, spaces=spaces)}
Output {i}: 
{fmt_grid(output, colour=False, spaces=spaces)}\n
"""
            i += 1

        test_grid = self.test[i_test][0]
        if rot90:
            test_grid = np.rot90(test_grid)
        if transpose:
            test_grid = np.transpose(test_grid)

        prompt += f"Input {i}:\n"
        prompt += f"{fmt_grid(test_grid, colour=False, spaces=spaces)}"
        if mode == "gpt3":
            prompt += f"\nOutput {i}:"
        else:
            prompt += f"\nOutput {i}: (please provide the output grid only)\n"

        if include_completion:
            if rot90:
                output = np.rot90(self.test[i_test][1])
            if transpose:
                output = np.transpose(self.test[i_test][1])
            prompt += f"\n{fmt_grid(self.test[i_test][1], colour=False, spaces=spaces)}"
        return prompt

class TaskSet:
    def __init__(self, tasks):
        tasks = sorted(tasks)
        self.tasks = tasks
        self.task_dict = {task.id: task for task in tasks}

    def __getitem__(self, item):
        if isinstance(item, slice):
            return TaskSet(self.tasks[item])
        get = self.task_dict.get(item)
        if get is None:
            try:
                return self.tasks[item]
            except (TypeError, IndexError):
                raise KeyError(f"Task {item} not found")
        return get

    def __len__(self):
        return len(self.tasks)

    def __iter__(self):
        return iter(self.tasks)

    def __repr__(self):
        return f"<TaskSet: {len(self.tasks)} tasks>"

    def score_submission(self,fn: str, topn=3, return_correct=False) -> int:
        """
        Score a submission file, in Kaggle csv format
        Two columns: output_id,output
        """
        from collections import defaultdict
        preds = defaultdict(list)

        with open(fn) as f:
            reader = csv.DictReader(f)
            for row in reader:
                task_id, test_num = row['output_id'].split('_')
                test_num = int(test_num)
                assert test_num == len(preds[task_id]), f'Predictions must be in order'
                 # Predictions must be in order

                row_preds = row['output'].strip().split(' ')
                row_preds = row_preds[:topn] # max 3 preds

                try:
                    array_preds = []
                    for p in row_preds:
                        p = p.strip('|').split('|')
                        p = [[int(x) for x in row] for row in p]
                        array_preds.append(np.array(p))
                except:
                    raise ValueError(f'Could not parse prediction: {row_preds}')

                preds[task_id].append(array_preds)

        total_score = 0
        correct = set()
        for task in self.tasks:
            test_examples = task.test
            assert len(test_examples) == len(preds[task.id])

            score = 0
            for ex, ex_preds in zip(test_examples, preds[task.id]):
                if any([np.array_equal(ex[1], pred) for pred in ex_preds]):
                    score += 1

            if score == len(test_examples):
                total_score += 1
                correct.add(task.id)

        if return_correct:
            return total_score, correct
        else:
            return total_score
        

def load_data() -> (TaskSet, TaskSet):
    data = json.load(open(f"{os.path.dirname(__file__)}/arc1.json"))
    train_tasks = []
    eval_tasks = []
    for id, task in data['train'].items():
        train_tasks.append(Task(id, task['train'], task['test'], 'train'))

    for id, task in data['eval'].items():
        eval_tasks.append(Task(id, task['train'], task['test'], 'eval'))

    return TaskSet(train_tasks), TaskSet(eval_tasks)

def load_single(id: str) -> Task:
    """
    Load a single task from disk.
    """
    data = json.load(open(f"{os.path.dirname(__file__)}/arc1.json"))
    # task = data['train'][id]
    # return Task(id, task['train'], task['test'], 'train'
    if id.startswith('train'):
        dataset_tasks = sorted(data['train'].items())
        taskid, task = dataset_tasks[int(id[5:])]
        return Task(taskid, task['train'], task['test'], 'train')
    elif id.startswith('eval'):
        dataset_tasks = sorted(data['eval'].items())
        taskid, task = dataset_tasks[int(id[4:])]
        return Task(taskid, task['train'], task['test'], 'eval')
    elif id in data['train']:
        task = data['train'][id]
        return Task(id, task['train'], task['test'], 'train')
    elif id in data['eval']:
        task = data['eval'][id]
        return Task(id, task['train'], task['test'], 'eval')
    else:
        raise ValueError(f"Unknown task id: {id}")

if __name__ == "__main__":
    train_tasks, eval_tasks = load_data()
    print(train_tasks)
    print(eval_tasks)
    print(train_tasks["007bbfb7"])
    for i in range(10):
        train_tasks[i].show()

    print(train_tasks["08ed6ac7"].gpt_prompt("gpt3"))
arckit/data.py
mxbi/arckit
import drawsvg
import numpy as np
import io

cmap = [
    '#252525', # black
    '#0074D9', # blue
    '#FF4136', # red
    '#37D449', #2ECC40', # green
    '#FFDC00', # yellow
    '#E6E6E6', # grey
    '#F012BE', # pink
    '#FF871E', # orange
    '#54D2EB', #7FDBFF', # light blue
    '#8D1D2C',#870C25', # brown
    '#FFFFFF'
        ]

def draw_grid(grid, xmax=10, ymax=10, padding=.5, extra_bottom_padding=0.5, group=False, add_size=True, label='', bordercol='#111111ff'):
    """
    Draws a grid, 

    Parameters
    ----------
    grid : np.ndarray
        The grid to draw
    xmax : float, optional
        The maximum horizontal size of the drawing, by default 10
    ymax : float, optional
        The maximum vertical size of the drawing, by default 10
    padding : float, optional
        The padding around the grid, half on each side, by default .5
    extra_bottom_padding : float, optional
        Extra padding at the bottom of the drawing, by default 0.5
        This is used to draw the label
    group : bool, optional
        If enabled, return a drawsvg.Group to include within a bigger drawing
        Otherwise, return a drawsvg.Drawing (default)
    label : str, optional
        A label to draw at the bottom left of the drawing, by default ''
        This does not affect drawing of the size.
    bordercol : str, optional
        The colour of the border, by default '#111111ff'
    """
    # Size is the total size of the LARGER axis.
    # With 0.5 cell padding, we consider the grid to be 1 unit larger than the number of cells
    # padding *= size # padding is proportional
    
    gridy, gridx = grid.shape
    # Calculate cell size based on the two restrictions
    # The actual cell size is the minimum of the two
    cellsize_x = xmax / gridx
    cellsize_y = ymax / gridy
    cellsize = min(cellsize_x, cellsize_y)

    xsize = gridx * cellsize
    ysize = gridy * cellsize

    line_thickness = 0.01
    circle_radius = 0
    border_width = 0.08
    lt = line_thickness / 2

    if group:
        drawing = drawsvg.Group()
    else:
        drawing = drawsvg.Drawing(xsize+padding, ysize+padding+extra_bottom_padding, origin=(-0.5*padding, -0.5*padding))
        drawing.set_pixel_scale(40)
    # drawing = drawsvg.Group()
    for j, row in enumerate(grid):
        for i, cell in enumerate(row):
            drawing.append(drawsvg.Rectangle(i*cellsize+lt, j*cellsize+lt, cellsize-lt, cellsize-lt, fill=cmap[cell]))
    
    # white dot at each vertex
    if circle_radius > 0:
        for i in range(1, gridx):
            for j in range(1, gridy):
                drawing.append(drawsvg.Circle(i*cellsize, j*cellsize, circle_radius, fill='white'))
        
    # Add a border
    bw = border_width / 3 # slightly more than 2 to avoid white border
    drawing.append(drawsvg.Rectangle(-bw, -bw, xsize+bw*2, ysize+bw*2, fill='none', stroke=bordercol, stroke_width=border_width))

    drawing.embed_google_font('Anuphan:wght@400;600;700', text=set(f'Input Output 0123456789x Test Task ABCDEFGHIJ? abcdefghjklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ'))

    # Write size on the bottom right
    # drawing.append(drawsvg.Text(text=f'{gridx}x{gridy}', x=-0.05, y=-0.25, font_size=padding/4, fill='black', text_anchor='start'))
    fontsize = (padding/2 + extra_bottom_padding)/2
    if add_size:
        drawing.append(drawsvg.Text(text=f'{gridx}x{gridy}', x=xsize, y=ysize+fontsize*1.25+0, font_size=fontsize, fill='black', text_anchor='end', font_family='Anuphan'))
    if label:
        drawing.append(drawsvg.Text(text=label, x=-0.1*fontsize, y=ysize+fontsize*1.25+0, font_size=fontsize, fill='black', text_anchor='start', font_family='Anuphan', font_weight='600'))
    
    if group:
        # return group, origin, (xsize, ysize)
        return drawing, (-0.5*padding, -0.5*padding), (xsize+padding, ysize+padding+extra_bottom_padding)
    return drawing

def draw_task(task, width=30, height=12, include_test=False, label=True, bordercols=['#111111ff', '#111111ff'], shortdesc=False):
    """
    Plot an entire task vertically, fitting the desired dimensions.
    The output is displayed below the input, with an arrow in between.

    Note that dimensions are a best effort, you should check .width and .height on the output
    
    Parameters
    ----------
    task : Task
        The task to plot
    width : float, optional
        The desired width of the drawing, by default 30
    height : float, optional
        The desired height of the drawing, by default 12
    include_test : bool, optional
        If enabled, include the test examples in the plot, by default False
        If set to 'all', ALSO include the output of the test examples
    label: bool, default True
    bordercols: list of str, default None
    """

    padding = 0.5
    bonus_padding = 0.25
    io_gap = 0.4
    ymax = (height - padding - bonus_padding - io_gap)/2

    if include_test:
        examples = task.train + task.test
    else:
        examples = task.train
    n_train = len(task.train)
    paddingless_width = width - padding * len(examples)

    max_widths = np.zeros(len(examples))
    # If any examples would exceed the height restriction, scale their width down and redistribute the space
    for i, (input_grid, output_grid) in enumerate(examples):
        input_grid_ratio = input_grid.shape[1] / input_grid.shape[0]
        output_grid_ratio = output_grid.shape[1] / output_grid.shape[0]
        max_ratio = max(input_grid_ratio, output_grid_ratio) # could be min
        xmax = ymax * max_ratio
        max_widths[i] = xmax

    # Allocate paddingless width to each example
    allocation = np.zeros_like(max_widths)
    increment = 0.01
    for i in range(int(paddingless_width//increment)):
        incr = (allocation + increment) <= max_widths
        allocation[incr] += increment / incr.sum()

    drawlist = []
    x_ptr = 0
    y_ptr = 0
    for i, (input_grid, output_grid) in enumerate(examples):
        if shortdesc:
            if i >= n_train:
                input_label = ''#f'T{i-n_train+1}'
                output_label = ''#f'T{i-n_train+1}'
            else:
                input_label = ''#f'I{i+1}'
                output_label = ''#f'O{i+1}'
        else:
            if i >= n_train:
                input_label = f'Test {i-n_train+1}'
                output_label = f'Test {i-n_train+1}'
            else:
                input_label = f'Input {i+1}'
                output_label = f'Output {i+1}'

        # input_label, output_label = '', ''

        input_grid, offset, (input_x, input_y) = draw_grid(input_grid, padding=padding, xmax=allocation[i], ymax=ymax, group=True, label=input_label, extra_bottom_padding=0.5, bordercol=bordercols[0])
        output_grid, offset, (output_x, output_y) = draw_grid(output_grid, padding=padding, xmax=allocation[i], ymax=ymax, group=True, label=output_label, extra_bottom_padding=0.5, bordercol=bordercols[1])

        drawlist.append(drawsvg.Use(input_grid, x=x_ptr + (allocation[i]+padding-input_x)/2 - offset[0], y=-offset[1]))
        # drawlist.append(drawsvg.Use(output_grid, x=x_ptr + (allocation[i]+padding-output_x)/2 - offset[0], y=ymax-offset[1]+2))
        x_ptr += max(input_x, output_x)
        y_ptr = max(y_ptr, input_y)

    x_ptr = 0
    y_ptr2 = 0
    for i, (input_grid, output_grid) in enumerate(examples):
        if shortdesc:
            if i >= n_train:
                input_label = ''#f'T{i-n_train+1}'
                output_label = ''#f'T{i-n_train+1}'
            else:
                input_label = ''#f'I{i+1}'
                output_label = ''#f'O{i+1}'
        else:
            if i >= n_train:
                input_label = f'Test {i-n_train+1}'
                output_label = f'Test {i-n_train+1}'
            else:
                input_label = f'Input {i+1}'
                output_label = f'Output {i+1}'
        # input_label, output_label = '', ''

        input_grid, offset, (input_x, input_y) = draw_grid(input_grid, padding=padding, xmax=allocation[i], ymax=ymax, group=True, label=input_label, extra_bottom_padding=0.5, bordercol=bordercols[0])
        output_grid, offset, (output_x, output_y) = draw_grid(output_grid, padding=padding, xmax=allocation[i], ymax=ymax, group=True, label=output_label, extra_bottom_padding=0.5, bordercol=bordercols[1])
                # Down arrow
        drawlist.append(drawsvg.Line(
            x_ptr + input_x/2, 
            y_ptr + padding - 0.6, 
            x_ptr + input_x/2, 
            y_ptr + padding + io_gap - 0.6, 
            stroke_width=0.05, stroke='#888888'))
        drawlist.append(drawsvg.Line(
            x_ptr + input_x/2 - 0.15,
            y_ptr + padding + io_gap - 0.8,
            x_ptr + input_x/2,
            y_ptr + padding + io_gap - 0.6,
            stroke_width=0.05, stroke='#888888'))
        drawlist.append(drawsvg.Line(
            x_ptr + input_x/2 + 0.15,
            y_ptr + padding + io_gap - 0.8,
            x_ptr + input_x/2,
            y_ptr + padding + io_gap - 0.6,
            stroke_width=0.05, stroke='#888888'))
        

        if i < n_train or include_test == 'all':
            drawlist.append(drawsvg.Use(output_grid, x=x_ptr + (allocation[i]+padding-output_x)/2 - offset[0], y=y_ptr-offset[1]+io_gap))
        else:
            # Add a question mark
            drawlist.append(drawsvg.Text(
                '?',
                x=x_ptr + (allocation[i]+padding)/2,
                y=y_ptr + output_y/2+bonus_padding,
                font_size=1,
                font_family='Anuphan',
                font_weight='700',
                fill='#333333',
                text_anchor='middle',
                alignment_baseline='middle',
            ))
        x_ptr += max(input_x, output_x)
        y_ptr2 = max(y_ptr2, y_ptr+output_y+io_gap)

    x_ptr = round(x_ptr, 1)
    y_ptr2 = round(y_ptr2, 1)
    d = drawsvg.Drawing(x_ptr, y_ptr2+0.2, origin=(0, 0))
    d.append(drawsvg.Rectangle(0, 0, '100%', '100%', fill='#eeeff6'))
    d.embed_google_font('Anuphan:wght@400;600;700', text=set(f'Input Output 0123456789x Test Task ABCDEFGHIJ? abcdefghjklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ'))
    for item in drawlist:
        d.append(item)

    fontsize=0.3
    d.append(drawsvg.Text(f"Task {task.id}", x=x_ptr-0.1, y=y_ptr2+0.1, font_size=fontsize, font_family='Anuphan', font_weight='600', fill='#666666', text_anchor='end', alignment_baseline='bottom'))

    d.set_pixel_scale(40)
    return d

def output_drawing(d: drawsvg.Drawing, filename: str, context=None):
    if filename.endswith('.svg'):
        d.save_svg(filename)
    elif filename.endswith('.png'):
        d.save_png(filename)
    elif filename.endswith('.pdf'):
        buffer = io.StringIO()
        d.as_svg(output_file=buffer, context=context)

        import cairosvg
        cairosvg.svg2pdf(bytestring=buffer.getvalue(), write_to=filename)
    else:
        raise ValueError(f'Unknown file extension for {filename}')
arckit/vis.py
mxbi/arckit