|
""" |
|
This examples trains a CrossEncoder for the Quora Duplicate Questions Detection task. A CrossEncoder takes a sentence pair |
|
as input and outputs a label. Here, it output a continious labels 0...1 to indicate the similarity between the input pair. |
|
|
|
It does NOT produce a sentence embedding and does NOT work for individual sentences. |
|
|
|
Usage: |
|
python training_quora_duplicate_questions.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 CEBinaryClassificationEvaluator |
|
from sentence_transformers.readers import InputExample |
|
import logging |
|
from datetime import datetime |
|
import os |
|
import gzip |
|
import csv |
|
from zipfile import ZipFile |
|
|
|
|
|
logging.basicConfig(format='%(asctime)s - %(message)s', |
|
datefmt='%Y-%m-%d %H:%M:%S', |
|
level=logging.INFO, |
|
handlers=[LoggingHandler()]) |
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
|
|
|
dataset_path = 'quora-dataset/' |
|
|
|
if not os.path.exists(dataset_path): |
|
logger.info("Dataset not found. Download") |
|
zip_save_path = 'quora-IR-dataset.zip' |
|
util.http_get(url='https://sbert.net/datasets/quora-IR-dataset.zip', path=zip_save_path) |
|
with ZipFile(zip_save_path, 'r') as zip: |
|
zip.extractall(dataset_path) |
|
|
|
|
|
|
|
logger.info("Read train dataset") |
|
train_samples = [] |
|
with open(os.path.join(dataset_path, 'classification', 'train_pairs.tsv'), 'r', encoding='utf8') as fIn: |
|
reader = csv.DictReader(fIn, delimiter='\t', quoting=csv.QUOTE_NONE) |
|
for row in reader: |
|
train_samples.append(InputExample(texts=[row['question1'], row['question2']], label=int(row['is_duplicate']))) |
|
train_samples.append(InputExample(texts=[row['question2'], row['question1']], label=int(row['is_duplicate']))) |
|
|
|
|
|
logger.info("Read dev dataset") |
|
dev_samples = [] |
|
with open(os.path.join(dataset_path, 'classification', 'dev_pairs.tsv'), 'r', encoding='utf8') as fIn: |
|
reader = csv.DictReader(fIn, delimiter='\t', quoting=csv.QUOTE_NONE) |
|
for row in reader: |
|
dev_samples.append(InputExample(texts=[row['question1'], row['question2']], label=int(row['is_duplicate']))) |
|
|
|
|
|
|
|
train_batch_size = 16 |
|
num_epochs = 4 |
|
model_save_path = 'output/training_quora-'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S") |
|
|
|
|
|
|
|
model = CrossEncoder('distilroberta-base', num_labels=1) |
|
|
|
|
|
train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=train_batch_size) |
|
|
|
|
|
|
|
evaluator = CEBinaryClassificationEvaluator.from_input_examples(dev_samples, name='Quora-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=5000, |
|
warmup_steps=warmup_steps, |
|
output_path=model_save_path) |
|
|
|
|
|
|