run_dt_ARC.py

import csv
import json
import logging
# make deterministic

from functools import reduce 
import operator
import os
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
import math
from torch.utils.data import Dataset

from minigpt.utils import set_seed
from minigpt.model_ARC import GPT, GPTConfig
from minigpt.trainer import Trainer, TrainerConfig
from minigpt.tester import Tester, TesterConfig, Test_StateActionReturnDataset
#from mingpt.utils import sample

from collections import deque
import random
import pickle
import argparse


parser = argparse.ArgumentParser()
parser.add_argument('--seed', type=int, default=123)
parser.add_argument('--context_length', type=int, default=6)
parser.add_argument('--epochs', type=int, default=5)
parser.add_argument('--batch_size', type=int, default=1)
parser.add_argument('--learning_rate', type=float, default=3e-4)
parser.add_argument('--lr_decay', type=bool, default=False)

parser.add_argument('--grid_x', type=int, default=5)
parser.add_argument('--grid_y', type=int, default=5)
parser.add_argument('--color_num', type=int, default=10)
parser.add_argument('--max_timestep', type=int, default=200)
parser.add_argument('--n_embd', type=int, default=768)
parser.add_argument('--n_layer', type=int, default=12)
parser.add_argument('--n_head', type=int, default=12)

parser.add_argument('--data_dir_prefix', type=str, default='/home/jovyan/workspace/Jaehyun/Mini_ARC_DT_2/')
parser.add_argument('--train_data_folder', type=str, default='mini_arc_train')
parser.add_argument('--test_data_dir_prefix', type=str, default='/home/jovyan/workspace/Jaehyun/Mini_ARC_DT/')
parser.add_argument('--test_data_folder', type=str, default='data')
parser.add_argument('--ckpt_path', type=str, default='./model_save')
parser.add_argument('--save_cycle', type=int, default=200)
args = parser.parse_args()

set_seed(args.seed)

ACTION_DIC={
    "start":0,
    "undo":1,
    "edit":2,
    "copyFromInput":3,
    "rotate":4,
    "reflecty":5,
    "reflectx":6,
    "translate":7,
    "resizeOutputGrid":8,
    "resetOutputGrid":9,
    "selected_cells":10, 
    "none":11,
    "end":12,
}
    
class Train_StateActionReturnDataset(Dataset):
    
    def __init__(self, data_path, data_type=None, step_gap=5, block_size=5*5+2): 
        self.actions_type = ACTION_DIC

        self.vocab_size = len(self.actions_type)
        self.block_size = block_size
        self.step_gap = step_gap
        self.data_path = Path(data_path)
        self.training_path = self.data_path / data_type
        self.training_tasks = []


        for path, _, files in os.walk(self.training_path):
            for name in files:
                self.training_tasks.append(os.path.join(path, name))

        self.num_dataset = len(self.training_tasks)
        
    def __len__(self):
        return self.num_dataset

    def __getitem__(self, idx):
        task_file = self.training_tasks[idx]
            
        with open(task_file, 'r') as f:
            task = json.load(f)
        
        total_step = int(task["total_step"])

        # start_idx = np.random.randint(-self.step_gap+2, total_step-self.step_gap+1)
        
        states = []
        actions_str = []
        actions = []
        rtgs = []
        timesteps = []
        for i in range(len(task["action_sequence"]["action_sequence"])):
            states.append(reduce(operator.add, task["action_sequence"]["action_sequence"][i]["grid"]))
            actions_str.append(task["action_sequence"]["action_sequence"][i]["action"]["tool"])
            actions.append(self.actions_type[task["action_sequence"]["action_sequence"][i]["action"]["tool"]])
            rtgs.append(float(task["action_sequence"]["action_sequence"][i]["reward"]))
            timesteps.append(int(task["action_sequence"]["action_sequence"][i]["time_step"]))

        states = torch.LongTensor(states)
        states = states.view([-1])
        states = states.contiguous()
        actions = torch.LongTensor(actions)
        rtgs = torch.FloatTensor(rtgs)
        timesteps = torch.LongTensor(timesteps)

        return states, actions, rtgs, timesteps

    
# set up logging
logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
)

train_dataset = Train_StateActionReturnDataset(
        data_path=args.data_dir_prefix,
        data_type=args.train_data_folder, 
        step_gap=args.context_length, 
        block_size=args.context_length
    )


mconf = GPTConfig(
        train_dataset.vocab_size, 
        train_dataset.block_size,
        n_layer=args.n_layer, 
        n_head=args.n_head, 
        n_embd=args.n_embd,
        max_timestep=args.max_timestep,
        grid_x=args.grid_x,
        grid_y=args.grid_y,
        color_num=args.color_num,
        context_length=args.context_length,
        train=True
    )
model = GPT(mconf)

# initialize a trainer instance and kick off training
epochs = args.epochs
tconf = TrainerConfig(
            max_epochs=epochs, 
            batch_size=args.batch_size, 
            learning_rate=args.learning_rate,
            lr_decay=args.lr_decay, 
            warmup_tokens=512*20, 
            final_tokens=2*len(train_dataset)*args.context_length*3,
            num_workers=4, 
            seed=args.seed, 
            max_timestep=args.max_timestep,
            ckpt_path=args.ckpt_path,
            save_cycle=args.save_cycle,
            action_dic=ACTION_DIC
        )

trainer = Trainer(model, train_dataset, None, tconf)

trainer.train()