from preprocessor import PreProcessor

from MIMIC_Bert_NER import *

if __name__ == '__main__':

    #Training data

    save_path = "./cleaned_files_train"

    concept_path = "../dnc_code/medical_data/train_data/concept"

    text_path = "../dnc_code/medical_data/train_data/txt"

    if not os.path.exists(save_path):

        os.mkdir(save_path)

    preproc = PreProcessor(concept_path,text_path,save_path)

    preproc.pre_process()

    #Train BERT

    bert = MIMICBertNER(save_path)

    bert.get_inputs()

    #bert_out_address = './saved_models/bert_out_model/mimic_bert'

    model = bert.train()

    # Save a trained model, configuration and tokenizer

    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self

    # If we save using the predefined names, we can load using `from_pretrained`

    output_model_file = os.path.join(bert.bert_out_address, "pytorch_model.bin")

    output_config_file = os.path.join(bert.bert_out_address, "config.json")

    torch.save(model_to_save.state_dict(), output_model_file)

    model_to_save.config.to_json_file(output_config_file)

    bert.tokenizer.save_vocabulary(bert.bert_out_address)

    #Validation data metrics

    bert.evaluate()

    #Test data

    save_path = "./cleaned_files_test"

    concept_path = "../dnc_code/medical_data/test_data/concept"

    text_path = "../dnc_code/medical_data/test_data/txt"

    if not os.path.exists(save_path):

        os.mkdir(save_path)

    preproc = PreProcessor(concept_path,text_path,save_path)

    preproc.pre_process()

    #Test BERT

    bert_test = MIMICBertNER(saved_path)

    bert_test.get_inputs("test")

    bert_test.evaluate()