import { Dataset } from '../../src/dataset/dataset'
import { DatabaseStorage } from '../../src/embeddings/database/storage'
import { BotonicNer } from '../../src/tasks/ner/botonic-ner'

const LOCALE = 'en'
const MAX_SEQUENCE_LENGTH = 12

const EMBEDDINGS_DIMENSION = 50
const EMBEDDINGS_TYPE = 'glove'

const EPOCHS = 4
const BATCH_SIZE = 8

const DATASET_DIR_PATH = '../dataset'

const MODEL_DIR = './model'

const dataset = Dataset.load(DATASET_DIR_PATH)
const { trainSet, testSet } = dataset.split()

;(async () => {
  const ner = BotonicNer.with(LOCALE, MAX_SEQUENCE_LENGTH, dataset.entities)

  ner.generateVocabulary(trainSet)

  ner.compile()

  await ner.createModel(
    'biLstm',
    await DatabaseStorage.with(LOCALE, EMBEDDINGS_TYPE, EMBEDDINGS_DIMENSION)
  )

  await ner.train(trainSet, EPOCHS, BATCH_SIZE)

  const { accuracy, loss } = await ner.evaluate(testSet)
  console.log(`Test Accuracy: ${accuracy}`)
  console.log(`Test loss: ${loss}`)

  const entities = ner.recognizeEntities('I want to return this black shirt')
  console.log(entities)

  await ner.saveModel(MODEL_DIR)
})()
