# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# cd ..

# +
import torch.optim as optim
import torch.nn as nn

from tree_to_sequence.program_datasets import TreeANCDataset
from tree_to_sequence.tree_to_sequence_attention_anc import TreeToSequenceAttentionANC
from tree_to_sequence.tree_encoder import TreeEncoder
from tree_to_sequence.sequence_encoder import SequenceEncoder
from tree_to_sequence.multilayer_lstm_cell import MultilayerLSTMCell
from tree_to_sequence.translating_trees import tree_to_list, decode_tokens, print_tree
from neural_nets_library import training
from ANC.util import printProgram

# + code_folding=[3, 19]
num_vars = 10
num_ints = 11

for_ops = {
    "Var": 0,
    "Const": 1,
    "Plus": 2,
    "Minus": 3,
    "EqualFor": 4,
    "LeFor": 5,
    "GeFor": 6,
    "Assign": 7,
    "If": 8,
    "Seq": 9,
    "For": 10
}

for_ops = {"<" + k.upper() + ">": v for k,v in for_ops.items()}

lambda_calculus_ops = {
                "<VARIABLE>": 0,
                "<ABSTRACTION>": 1,
                "<NUMBER>": 2,
                "<BOOLEAN>": 3,
                "<NIL>": 4,
                "<IF>": 5,
                "<CONS>": 6,
                "<MATCH>": 7,
                "<UNARYOPER>": 8,
                "<BINARYOPER>": 9,
                "<LET>": 10,
                "<LETREC>": 11,
                "<TRUE>": 12,
                "<FALSE>": 13,
                "<TINT>": 14,
                "<TBOOL>": 15,
                "<TINTLIST>": 16,
                "<TFUN>": 17,
                "<ARGUMENT>": 18,
                "<NEG>": 19,
                "<NOT>": 20,
                "<PLUS>": 21,
                "<MINUS>": 22,
                "<TIMES>": 23,
                "<DIVIDE>": 24,
                "<AND>": 25,
                "<OR>": 26,
                "<EQUAL>": 27,
                "<LESS>": 28,
                "<APPLICATION>": 29,
                "<HEAD>": 30,
                "<TAIL>": 31
            }

# + code_folding=[] hide_input=false
input_eos_token = False
input_as_seq = False
use_embedding = True
eos_bonus = 1 if input_eos_token and input_as_seq else 0
long_base_case = True
binarize = True

# + code_folding=[] hide_input=false
is_lambda_calculus = False
M, R = 5, 3
N = 11

for_anc_dset = TreeANCDataset("ANC/Easy-arbitraryForListWithOutput.json", 
                              is_lambda_calculus, num_ints = M,
                              binarize=binarize, input_eos_token=input_eos_token, 
                              use_embedding=use_embedding, long_base_case=long_base_case, 
                              input_as_seq=input_as_seq, cuda=True)


# + code_folding=[0]
def reset_all_parameters_uniform(model, stdev):
    for param in model.parameters():
        nn.init.uniform(param, -stdev, stdev)


# +
embedding_size = 200
hidden_size = 200
num_layers = 1
alignment_size = 50
align_type = 1
encoder_input_size = num_vars + num_ints + len(for_ops.keys()) + eos_bonus

if input_as_seq:
    encoder = SequenceEncoder(encoder_input_size, hidden_size, num_layers, attention=True, use_embedding=use_embedding)
else:
    encoder = TreeEncoder(encoder_input_size, hidden_size, num_layers, [1, 2], attention=True, use_embedding=use_embedding)

decoder = MultilayerLSTMCell(N + 3*R + hidden_size, hidden_size, num_layers)
program_model = TreeToSequenceAttentionANC(encoder, decoder, hidden_size, embedding_size, M, R, 
                                           alignment_size=alignment_size, align_type=align_type,
                                           correctness_weight=5, halting_weight=1, confidence_weight=5, 
                                           efficiency_weight=1, diversity_weight=0, mix_probabilities=True,
                                           predict_registers=True)
    
reset_all_parameters_uniform(program_model, 0.1)
encoder.initialize_forget_bias(3)
decoder.initialize_forget_bias(3)
# -

program_model = program_model.cuda()

# + code_folding=[0]
optimizer = optim.Adam(program_model.parameters(), lr=0.005)
lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, verbose=True, patience=20000, factor=0.8)
# -

training.train_model_tree_to_anc(program_model, itertools.repeat(for_anc_dset[1], 100000), 
                                 optimizer, lr_scheduler=lr_scheduler, num_epochs=1, 
                                 batch_size=1, plateau_lr=True, print_every=200)

for prog, target in for_anc_dset[0:5]:
    print(decode_tokens(tree_to_list(prog), 10, M, for_ops))
    controller = program_model.forward_prediction(prog)
    
    input_memory = target[0][0]
    correct_memory = target[1][0]
    prediction_memory, _ = controller.forward_prediction([input_memory])
    
    printProgram(controller, 0.9)
    print(correct_memory[0])
    print(prediction_memory[0])
    
    mem_diff = correct_memory[0] - prediction_memory[0].data
    
    print(torch.sum(mem_diff * mem_diff))

StarcoderdataJupyterScriptsDedupFiltered