# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
# ---

# %load_ext autoreload
# %autoreload 2

# +
# %matplotlib inline

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# +
import pickle

seq_length = 53

with open('train.pkl', 'rb') as f:
    gs_selected_train, ocr_selected_train = pickle.load(f)
    
with open('val.pkl', 'rb') as f:
    gs_selected_val, ocr_selected_val = pickle.load(f)
    
with open('ci.pkl', 'rb') as f:
    ci = pickle.load(f)
# -

gs_selected_train[:50]

# +
from ochre.datagen import DataGenerator

dg_val = DataGenerator(xData=ocr_selected_val[:10], yData=gs_selected_val[:10], char_to_int=ci,
                       seq_length=seq_length, padding_char='\n', oov_char='@',
                       batch_size=10, shuffle=False)
dg_train = DataGenerator(xData=ocr_selected_train[:50], yData=gs_selected_train[:50], char_to_int=ci,
                       seq_length=seq_length, padding_char='\n', oov_char='@',
                       batch_size=10, shuffle=False)

# +
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import LSTM
from keras.layers import TimeDistributed
from keras.layers import Bidirectional
from keras.layers import RepeatVector
from keras.layers import Embedding
from keras.callbacks import ModelCheckpoint

n_nodes = 1000
dropout = 0.2
n_embed = 256
n_vocab = len(ci)

loss='categorical_crossentropy'
optimizer='adam'
metrics=['accuracy']

model = Sequential()

# encoder

model.add(Embedding(n_vocab, n_embed, input_length=seq_length))
model.add(LSTM(n_nodes, input_shape=(seq_length, n_vocab)))
# For the decoder's input, we repeat the encoded input for each time step
model.add(RepeatVector(seq_length))
model.add(LSTM(n_nodes, return_sequences=True))

# For each of step of the output sequence, decide which character should be
# chosen
model.add(TimeDistributed(Dense(n_vocab, activation='softmax')))
model.compile(loss=loss, optimizer=optimizer, metrics=metrics)

# +
# initialize saving of weights
#filepath = os.path.join(weights_dir, '{loss:.4f}-{epoch:02d}.hdf5')
filepath = '{loss:.4f}-{epoch:02d}.hdf5'
checkpoint = ModelCheckpoint(filepath, monitor='loss', verbose=1,
                                 save_best_only=True, mode='min')
callbacks_list = [checkpoint]

# do training (and save weights)
model.fit_generator(dg_train, steps_per_epoch=len(dg_train), epochs=10, 
                    validation_data=dg_val, 
                    validation_steps=len(dg_val), callbacks=callbacks_list,
                    use_multiprocessing=True,
                    workers=3)

StarcoderdataJupyterScriptsDedupFiltered