from simpletransformers.classification import (ClassificationArgs, ClassificationModel)

import pandas as pd

import pickle

# ### Read data from files

# In[ ]:

train_df = pd.read_csv('/Users/aakansha/Desktop/NCCS NLP for Histology Reports/Datasets for Trials/train.csv')

train_df['Combined Diagnosis'] = train_df['Diagnosis'] + train_df['Gross Description'] + train_df[

    'Microscgopic Description']

train_df = train_df[['Combined Diagnosis', 'Cancerous?']]

# ### Pre-process data

# In[ ]:

# Capitalize values

train_df['Cancerous?'] = train_df['Cancerous?'].str.upper()

train_df['Cancerous?'] = train_df['Cancerous?'].str.strip()

# Drop all NA rows

train_df = train_df.dropna(subset=['Cancerous?']).reset_index(drop=True)

# ## Roberta-Large

# In[ ]:

for n in [1]:

    curr_epoch = "Epoch" + str(n)

    # Configure model args

    model_args = ClassificationArgs(num_train_epochs=n)

    model_args.evaluate_during_training_steps = -1

    model_args.save_eval_checkpoints = False

    model_args.save_model_every_epoch = False

    model_args.learning_rate = 1e-5

    model_args.manual_seed = 4

    model_args.multiprocessing_chunksize = 5000

    model_args.no_cache = True

    model_args.reprocess_input_data = True

    model_args.train_batch_size = 16

    model_args.gradient_accumulation_steps = 2

    model_args.use_multiprocessing = True

    model_args.overwrite_output_dir = True

    model_args.labels_list = ['YES', 'NO']

    # model

    model_type = "roberta"

    model_name = "roberta-large"

    # Create Transformer Model

    model = ClassificationModel(model_type, model_name, num_labels=2, use_cuda=False, args=model_args)

    if __name__ == '__main__':

        # Train the model

        model.train_model(train_df)

        pickle.dump(model, open('model.pkl', 'wb'))