from datasets import load_dataset
from sentence_transformers import SentenceTransformer, losses, InputExample
from datetime import datetime
from torch.utils.data import DataLoader
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
import torch
import math

# Read the dataset
model_name = 'bert-base-nli-mean-tokens'
train_batch_size = 128
num_epochs = 1
model_save_path = 'output/training_paws_continue_training-'+model_name+'-'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

# Load a pre-trained sentence transformer model
model = SentenceTransformer('../models/stsb-bert-large')

# Convert the dataset to a DataLoader ready for training
train_samples = []
dev_samples = []

train_dataset = load_dataset('paws', 'labeled_final', split='train')
train_dataset.set_format(type='pandas')
train_dataset = train_dataset[:]
for index, row in train_dataset.iterrows():
    x = torch.FloatTensor([row['label']])
    train_samples.append(InputExample(texts=[row['sentence1'], row['sentence2']], label=x))

dev_dataset = load_dataset('paws', 'labeled_final', split='validation')
dev_dataset.set_format(type='pandas')
dev_dataset = dev_dataset[:]


for index, row in dev_dataset.iterrows():
    dev_samples.append(InputExample(texts=[row['sentence1'], row['sentence2']], label=row['label']))


train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=train_batch_size)
train_loss = losses.CosineSimilarityLoss(model=model)

# Development set: Measure correlation between cosine score and gold labels
evaluator = EmbeddingSimilarityEvaluator.from_input_examples(dev_samples, name='paws-dev')

# Configure the training. We skip evaluation in this example
warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1)  # 10% of train data for warm-up

# Train the model
model.fit(train_objectives=[(train_dataloader, train_loss)],
          evaluator=evaluator,
          epochs=num_epochs,
          evaluation_steps=1000,
          warmup_steps=warmup_steps,
          output_path=model_save_path)
