diff --git a/train.py b/train.py new file mode 100644 index 0000000000000000000000000000000000000000..fdf3936a593eaf3ceecc042167694b57caec06d2 --- /dev/null +++ b/train.py @@ -0,0 +1,17 @@ +import torch +from Configuration import Configuration +from Linker import * +from utils import read_csv_pgbar + +torch.cuda.empty_cache() +batch_size = int(Configuration.modelTrainingConfig['batch_size']) +nb_sentences = batch_size * 800 +epochs = int(Configuration.modelTrainingConfig['epoch']) + +file_path_axiom_links = 'Datasets/goldANDsilver_dataset_links.csv' +df_axiom_links = read_csv_pgbar(file_path_axiom_links, nb_sentences) + +print("Linker") +linker = Linker("models/flaubert_super_98_V2_50e.pt") +print("\nLinker Training\n") +linker.train_linker(df_axiom_links, validation_rate=0.1, epochs=epochs, batch_size=batch_size, checkpoint=False, tensorboard=True) \ No newline at end of file