|
""" |
|
This examples trains a CrossEncoder for the NLI task. A CrossEncoder takes a sentence pair |
|
as input and outputs a label. Here, it learns to predict the labels: "contradiction": 0, "entailment": 1, "neutral": 2. |
|
|
|
It does NOT produce a sentence embedding and does NOT work for individual sentences. |
|
|
|
Usage: |
|
python training_nli.py |
|
""" |
|
from torch.utils.data import DataLoader |
|
import math |
|
from sentence_transformers import LoggingHandler, util |
|
from sentence_transformers.cross_encoder import CrossEncoder |
|
from sentence_transformers.cross_encoder.evaluation import CESoftmaxAccuracyEvaluator |
|
from sentence_transformers.readers import InputExample |
|
import logging |
|
from datetime import datetime |
|
import os |
|
import gzip |
|
import csv |
|
|
|
|
|
logging.basicConfig(format='%(asctime)s - %(message)s', |
|
datefmt='%Y-%m-%d %H:%M:%S', |
|
level=logging.INFO, |
|
handlers=[LoggingHandler()]) |
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
|
|
|
nli_dataset_path = 'datasets/AllNLI.tsv.gz' |
|
|
|
if not os.path.exists(nli_dataset_path): |
|
util.http_get('https://sbert.net/datasets/AllNLI.tsv.gz', nli_dataset_path) |
|
|
|
|
|
|
|
logger.info("Read AllNLI train dataset") |
|
|
|
label2int = {"contradiction": 0, "entailment": 1, "neutral": 2} |
|
train_samples = [] |
|
dev_samples = [] |
|
with gzip.open(nli_dataset_path, 'rt', encoding='utf8') as fIn: |
|
reader = csv.DictReader(fIn, delimiter='\t', quoting=csv.QUOTE_NONE) |
|
for row in reader: |
|
label_id = label2int[row['label']] |
|
if row['split'] == 'train': |
|
train_samples.append(InputExample(texts=[row['sentence1'], row['sentence2']], label=label_id)) |
|
else: |
|
dev_samples.append(InputExample(texts=[row['sentence1'], row['sentence2']], label=label_id)) |
|
|
|
|
|
|
|
train_batch_size = 16 |
|
num_epochs = 4 |
|
model_save_path = 'output/training_allnli-'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S") |
|
|
|
|
|
model = CrossEncoder('distilroberta-base', num_labels=len(label2int)) |
|
|
|
|
|
train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=train_batch_size) |
|
|
|
|
|
evaluator = CESoftmaxAccuracyEvaluator.from_input_examples(dev_samples, name='AllNLI-dev') |
|
|
|
|
|
warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1) |
|
logger.info("Warmup-steps: {}".format(warmup_steps)) |
|
|
|
|
|
|
|
model.fit(train_dataloader=train_dataloader, |
|
evaluator=evaluator, |
|
epochs=num_epochs, |
|
evaluation_steps=10000, |
|
warmup_steps=warmup_steps, |
|
output_path=model_save_path) |
|
|
|
|
|
|