from model import *
from preprocessing import *

import pandas as pd

from cc import ce

bs = 32

data = pd.read_csv("/mnt/tmpfs1/LJSpeech-1.1/metadata.csv",
                   sep = '|',
                   header = None,
                   quoting = 3,
                   names = ['file_name', 'i', 'normalized_transcription'])
s = int(len(data) * 0.90)
train_data = data[:s]
valid_data = data[s:]

train_ds = to_dataset(train_data, batch_size = bs)
valid_ds = to_dataset(valid_data, batch_size = bs)

m = model(input_dim = fft_length // 2 + 1,
          output_dim = char_to_num.vocabulary_size())

m.load_weights('model40-latest.keras')
ckpt1 = kc.ModelCheckpoint('model41-latest.keras',
                           monitor = 'val_loss',
                           save_best_only = False,
                           verbose = 1)

ckpt2 = kc.ModelCheckpoint('model41-best.keras',
                           monitor = 'val_loss',
                           save_best_only = True,
                           verbose = 1)

ce1 = ce(valid_ds, m)

m.fit(train_ds,
      epochs = 40,
      validation_data = valid_ds,
      callbacks = [ckpt1, ckpt2, ce1])