From de56206fe4fadc723b5751c635a4cb1e6d72479c Mon Sep 17 00:00:00 2001 From: Linqing Liu Date: Mon, 8 Oct 2018 19:23:22 -0400 Subject: [PATCH] add SNLI / STS-2014 / Quora dataset (#148) * add SNLI dataset * add STS-2014 * add trainers and evaluators for Quora * add quora in datasets/ * process the merge confict in common/dataset.py --- common/dataset.py | 20 +++++++- common/evaluation.py | 8 ++- common/evaluators/snli_evaluator.py | 32 ++++++++++++ common/train.py | 8 ++- common/trainers/quora_trainer.py | 77 +++++++++++++++++++++++++++++ common/trainers/snli_trainer.py | 77 +++++++++++++++++++++++++++++ common/trainers/sts2014_trainer.py | 77 +++++++++++++++++++++++++++++ datasets/quora.py | 66 +++++++++++++++++++++++++ datasets/sick.py | 2 +- datasets/snli.py | 66 +++++++++++++++++++++++++ datasets/sts2014.py | 72 +++++++++++++++++++++++++++ 11 files changed, 501 insertions(+), 4 deletions(-) create mode 100644 common/evaluators/snli_evaluator.py create mode 100644 common/trainers/quora_trainer.py create mode 100644 common/trainers/snli_trainer.py create mode 100644 common/trainers/sts2014_trainer.py create mode 100644 datasets/quora.py create mode 100644 datasets/snli.py create mode 100644 datasets/sts2014.py diff --git a/common/dataset.py b/common/dataset.py index e766148..52644b3 100644 --- a/common/dataset.py +++ b/common/dataset.py @@ -8,6 +8,9 @@ from datasets.msrvid import MSRVID from datasets.trecqa import TRECQA from datasets.wikiqa import WikiQA from datasets.pit2015 import PIT2015 +from datasets.snli import SNLI +from datasets.sts2014 import STS2014 +from datasets.quora import Quora from datasets.reuters import Reuters class UnknownWordVecCache(object): @@ -66,7 +69,22 @@ class DatasetFactory(object): train_loader, dev_loader, test_loader = PIT2015.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWordVecCache.unk) embedding = nn.Embedding.from_pretrained(PIT2015.TEXT_FIELD.vocab.vectors) return PIT2015, embedding, train_loader, test_loader, dev_loader - elif dataset_name == 'reuters': + elif dataset_name == 'snli': + dataset_root = os.path.join(castor_dir, os.pardir, 'Castor-data', 'datasets', 'snli_1.0/') + train_loader, dev_loader, test_loader = SNLI.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWordVecCache.unk) + embedding = nn.Embedding.from_pretrained(SNLI.TEXT_FIELD.vocab.vectors) + return SNLI, embedding, train_loader, test_loader, dev_loader + elif dataset_name == 'sts2014': + dataset_root = os.path.join(castor_dir, os.pardir, 'Castor-data', 'datasets', 'STS-2014') + train_loader, dev_loader, test_loader = STS2014.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWordVecCache.unk) + embedding = nn.Embedding.from_pretrained(STS2014.TEXT_FIELD.vocab.vectors) + return STS2014, embedding, train_loader, test_loader, dev_loader + elif dataset_name == "quora": + dataset_root = os.path.join(castor_dir, os.pardir, 'Castor-data', 'datasets', 'quora/') + train_loader, dev_loader, test_loader = Quora.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWordVecCache.unk) + embedding = nn.Embedding.from_pretrained(Quora.TEXT_FIELD.vocab.vectors) + return Quora, embedding, train_loader, test_loader, dev_loader + elif dataset_name == 'reuters': dataset_root = os.path.join(castor_dir, os.pardir, 'Castor-data', 'datasets', 'Reuters-21578/') train_loader, dev_loader, test_loader = Reuters.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWordVecCache.unk) embedding = nn.Embedding.from_pretrained(Reuters.TEXT_FIELD.vocab.vectors) diff --git a/common/evaluation.py b/common/evaluation.py index 44064ee..d28c9d7 100644 --- a/common/evaluation.py +++ b/common/evaluation.py @@ -5,6 +5,9 @@ from .evaluators.trecqa_evaluator import TRECQAEvaluator from .evaluators.wikiqa_evaluator import WikiQAEvaluator from .evaluators.pit2015_evaluator import PIT2015Evaluator from .evaluators.reuters_evaluator import ReutersEvaluator +from .evaluators.snli_evaluator import SNLIEvaluator +from .evaluators.sts2014_evaluator import STS2014Evaluator +from .evaluators.quora_evaluator import QuoraEvaluator from nce.nce_pairwise_mp.evaluators.trecqa_evaluator import TRECQAEvaluatorNCE from nce.nce_pairwise_mp.evaluators.wikiqa_evaluator import WikiQAEvaluatorNCE @@ -22,7 +25,10 @@ class EvaluatorFactory(object): 'wikiqa': WikiQAEvaluator, 'pit2015': PIT2015Evaluator, 'twitterurl': PIT2015Evaluator, - 'Reuters': ReutersEvaluator + 'Reuters': ReutersEvaluator, + 'SNLI': SNLIEvaluator, + 'sts2014': STS2014Evaluator, + 'Quora': QuoraEvaluator } evaluator_map_nce = { diff --git a/common/evaluators/snli_evaluator.py b/common/evaluators/snli_evaluator.py new file mode 100644 index 0000000..6357b1d --- /dev/null +++ b/common/evaluators/snli_evaluator.py @@ -0,0 +1,32 @@ +import torch +import torch.nn.functional as F + +from .evaluator import Evaluator + + +class SNLIEvaluator(Evaluator): + + def get_scores(self): + self.model.eval() + test_kl_div_loss = 0 + acc_total = 0 + + for batch in self.data_loader: + # Select embedding + sent1, sent2 = self.get_sentence_embeddings(batch) + + output = self.model(sent1, sent2, batch.ext_feats, batch.dataset.word_to_doc_cnt, batch.sentence_1_raw, batch.sentence_2_raw) + test_kl_div_loss += F.kl_div(output, batch.label, size_average=False).item() + + true_label = torch.max(batch.label.data, 1)[1] + prediction = torch.max(output, 1)[1] + acc_total += ((true_label == prediction)).sum().item() + + del output + + test_kl_div_loss /= len(batch.dataset.examples) + + accuracy = acc_total / len(self.data_loader.dataset.examples) + + return [accuracy, test_kl_div_loss], ['accuracy', 'KL-divergence loss'] + diff --git a/common/train.py b/common/train.py index bfd29df..abe1123 100644 --- a/common/train.py +++ b/common/train.py @@ -5,6 +5,9 @@ from .trainers.wikiqa_trainer import WikiQATrainer from .trainers.pit2015_trainer import PIT2015Trainer from .trainers.sst_trainer import SSTTrainer from .trainers.reuters_trainer import ReutersTrainer +from .trainers.snli_trainer import SNLITrainer +from .trainers.sts2014_trainer import STS2014Trainer +from .trainers.quora_trainer import QuoraTrainer from nce.nce_pairwise_mp.trainers.trecqa_trainer import TRECQATrainerNCE from nce.nce_pairwise_mp.trainers.wikiqa_trainer import WikiQATrainerNCE @@ -22,7 +25,10 @@ class TrainerFactory(object): 'wikiqa': WikiQATrainer, 'pit2015': PIT2015Trainer, 'twitterurl': PIT2015Trainer, - 'Reuters': ReutersTrainer + 'Reuters': ReutersTrainer, + 'snli': SNLITrainer, + 'sts2014': STS2014Trainer, + 'quora': QuoraTrainer } trainer_map_nce = { diff --git a/common/trainers/quora_trainer.py b/common/trainers/quora_trainer.py new file mode 100644 index 0000000..9b3071b --- /dev/null +++ b/common/trainers/quora_trainer.py @@ -0,0 +1,77 @@ +import time + +import torch.nn as nn +import torch.nn.functional as F +from torch.optim.lr_scheduler import ReduceLROnPlateau + +from .trainer import Trainer +from utils.serialization import save_checkpoint + +class QuoraTrainer(Trainer): + def train_epoch(self, epoch): + self.model.train() + total_loss = 0 + for batch_idx, batch in enumerate(self.train_loader): + self.optimizer.zero_grad() + + # Select embedding + sent1, sent2 = self.get_sentence_embeddings(batch) + + output = self.model(sent1, sent2, batch.ext_feats, batch.dataset.word_to_doc_cnt, batch.sentence_1_raw, batch.sentence_2_raw) + loss = F.kl_div(output, batch.label, size_average=False) + total_loss += loss.item() + loss.backward() + if self.clip_norm: + nn.utils.clip_grad_norm(self.model.parameters(), self.clip_norm) + self.optimizer.step() + if batch_idx % self.log_interval == 0: + self.logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + epoch, min(batch_idx * self.batch_size, len(batch.dataset.examples)), + len(batch.dataset.examples), + 100. * batch_idx / (len(self.train_loader)), loss.item() / len(batch)) + ) + + if self.use_tensorboard: + self.writer.add_scalar('quora/train/kl_div_loss', total_loss / len(self.train_loader.dataset.examples), epoch) + + return total_loss + + def train(self, epochs): + scheduler = None + if self.lr_reduce_factor != 1 and self.lr_reduce_factor != None: + scheduler = ReduceLROnPlateau(self.optimizer, mode='max', factor=self.lr_reduce_factor, + patience=self.patience) + epoch_times = [] + prev_loss = -1 + best_dev_score = -1 + for epoch in range(1, epochs + 1): + start = time.time() + self.logger.info('Epoch {} started...'.format(epoch)) + self.train_epoch(epoch) + + accuracy, new_loss = self.evaluate(self.dev_evaluator, 'dev') + + if self.use_tensorboard: + self.writer.add_scalar('quora/lr', self.optimizer.param_groups[0]['lr'], epoch) + self.writer.add_scalar('quora/dev/accuracy', accuracy, epoch) + self.writer.add_scalar('quora/dev/kl_div_loss', new_loss, epoch) + + end = time.time() + duration = end - start + self.logger.info('Epoch {} finished in {:.2f} minutes'.format(epoch, duration / 60)) + epoch_times.append(duration) + + if accuracy > best_dev_score: + best_dev_score = accuracy + save_checkpoint(epoch, self.model.arch, self.model.state_dict(), self.optimizer.state_dict(), + best_dev_score, self.model_outfile) + + if abs(prev_loss - new_loss) <= 0.0002: + self.logger.info('Early stopping. Loss changed by less than 0.0002.') + break + + prev_loss = new_loss + if scheduler is not None: + scheduler.step(accuracy) + + self.logger.info('Training took {:.2f} minutes overall...'.format(sum(epoch_times) / 60)) diff --git a/common/trainers/snli_trainer.py b/common/trainers/snli_trainer.py new file mode 100644 index 0000000..2ee18de --- /dev/null +++ b/common/trainers/snli_trainer.py @@ -0,0 +1,77 @@ +import time + +import torch.nn as nn +import torch.nn.functional as F +from torch.optim.lr_scheduler import ReduceLROnPlateau + +from .trainer import Trainer +from utils.serialization import save_checkpoint + +class SNLITrainer(Trainer): + def train_epoch(self, epoch): + self.model.train() + total_loss = 0 + for batch_idx, batch in enumerate(self.train_loader): + self.optimizer.zero_grad() + + # Select embedding + sent1, sent2 = self.get_sentence_embeddings(batch) + + output = self.model(sent1, sent2, batch.ext_feats, batch.dataset.word_to_doc_cnt, batch.sentence_1_raw, batch.sentence_2_raw) + loss = F.kl_div(output, batch.label, size_average=False) + total_loss += loss.item() + loss.backward() + if self.clip_norm: + nn.utils.clip_grad_norm(self.model.parameters(), self.clip_norm) + self.optimizer.step() + if batch_idx % self.log_interval == 0: + self.logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + epoch, min(batch_idx * self.batch_size, len(batch.dataset.examples)), + len(batch.dataset.examples), + 100. * batch_idx / (len(self.train_loader)), loss.item() / len(batch)) + ) + + if self.use_tensorboard: + self.writer.add_scalar('snli/train/kl_div_loss', total_loss / len(self.train_loader.dataset.examples), epoch) + + return total_loss + + def train(self, epochs): + scheduler = None + if self.lr_reduce_factor != 1 and self.lr_reduce_factor != None: + scheduler = ReduceLROnPlateau(self.optimizer, mode='max', factor=self.lr_reduce_factor, + patience=self.patience) + epoch_times = [] + prev_loss = -1 + best_dev_score = -1 + for epoch in range(1, epochs + 1): + start = time.time() + self.logger.info('Epoch {} started...'.format(epoch)) + self.train_epoch(epoch) + + accuracy, new_loss = self.evaluate(self.dev_evaluator, 'dev') + + if self.use_tensorboard: + self.writer.add_scalar('snli/lr', self.optimizer.param_groups[0]['lr'], epoch) + self.writer.add_scalar('snli/dev/accuracy', accuracy, epoch) + self.writer.add_scalar('snli/dev/kl_div_loss', new_loss, epoch) + + end = time.time() + duration = end - start + self.logger.info('Epoch {} finished in {:.2f} minutes'.format(epoch, duration / 60)) + epoch_times.append(duration) + + if accuracy > best_dev_score: + best_dev_score = accuracy + save_checkpoint(epoch, self.model.arch, self.model.state_dict(), self.optimizer.state_dict(), + best_dev_score, self.model_outfile) + + if abs(prev_loss - new_loss) <= 0.0002: + self.logger.info('Early stopping. Loss changed by less than 0.0002.') + break + + prev_loss = new_loss + if scheduler is not None: + scheduler.step(accuracy) + + self.logger.info('Training took {:.2f} minutes overall...'.format(sum(epoch_times) / 60)) diff --git a/common/trainers/sts2014_trainer.py b/common/trainers/sts2014_trainer.py new file mode 100644 index 0000000..dbb81db --- /dev/null +++ b/common/trainers/sts2014_trainer.py @@ -0,0 +1,77 @@ +import time + +import torch.nn as nn +import torch.nn.functional as F +from torch.optim.lr_scheduler import ReduceLROnPlateau + +from .trainer import Trainer +from utils.serialization import save_checkpoint + + +class STS2014Trainer(Trainer): + + def train_epoch(self, epoch): + self.model.train() + total_loss = 0 + for batch_idx, batch in enumerate(self.train_loader): + self.optimizer.zero_grad() + + # Select embedding + sent1, sent2 = self.get_sentence_embeddings(batch) + + output = self.model(sent1, sent2, batch.ext_feats, batch.dataset.word_to_doc_cnt, batch.sentence_1_raw, batch.sentence_2_raw) + loss = F.kl_div(output, batch.label, size_average=False) + total_loss += loss.item() + loss.backward() + if self.clip_norm: + nn.utils.clip_grad_norm(self.model.parameters(), self.clip_norm) + self.optimizer.step() + if batch_idx % self.log_interval == 0: + self.logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + epoch, min(batch_idx * self.batch_size, len(batch.dataset.examples)), + len(batch.dataset.examples), + 100. * batch_idx / (len(self.train_loader)), loss.item() / len(batch)) + ) + + if self.use_tensorboard: + self.writer.add_scalar('sts2014/train/kl_div_loss', total_loss / len(self.train_loader.dataset.examples), epoch) + + return total_loss + + def train(self, epochs): + scheduler = None + if self.lr_reduce_factor != 1 and self.lr_reduce_factor != None: + scheduler = ReduceLROnPlateau(self.optimizer, mode='max', factor=self.lr_reduce_factor, patience=self.patience) + epoch_times = [] + prev_loss = -1 + best_dev_score = -1 + for epoch in range(1, epochs + 1): + start = time.time() + self.logger.info('Epoch {} started...'.format(epoch)) + self.train_epoch(epoch) + + pearson, spearman, mse, new_loss = self.evaluate(self.dev_evaluator, 'dev') + + if self.use_tensorboard: + self.writer.add_scalar('sts2014/lr', self.optimizer.param_groups[0]['lr'], epoch) + self.writer.add_scalar('sts2014/dev/pearson_r', pearson, epoch) + self.writer.add_scalar('sts2014/dev/kl_div_loss', new_loss, epoch) + + end = time.time() + duration = end - start + self.logger.info('Epoch {} finished in {:.2f} minutes'.format(epoch, duration / 60)) + epoch_times.append(duration) + + if pearson > best_dev_score: + best_dev_score = pearson + save_checkpoint(epoch, self.model.arch, self.model.state_dict(), self.optimizer.state_dict(), best_dev_score, self.model_outfile) + + if abs(prev_loss - new_loss) <= 0.0002: + self.logger.info('Early stopping. Loss changed by less than 0.0002.') + break + + prev_loss = new_loss + if scheduler is not None: + scheduler.step(pearson) + + self.logger.info('Training took {:.2f} minutes overall...'.format(sum(epoch_times) / 60)) diff --git a/datasets/quora.py b/datasets/quora.py new file mode 100644 index 0000000..82ad988 --- /dev/null +++ b/datasets/quora.py @@ -0,0 +1,66 @@ +import numpy as np +import torch +from torchtext.data.field import Field, RawField +from torchtext.data.iterator import BucketIterator +from torchtext.vocab import Vectors +from torchtext.data.pipeline import Pipeline + +from datasets.castor_dataset import CastorPairDataset + +def get_class_probs(sim, *args): + """ + Convert a single label into class probabilities. + """ + class_probs = np.zeros(Quora.NUM_CLASSES) + class_probs[int(sim)] = 1 + return class_probs + + +class Quora(CastorPairDataset): + NAME = 'Quora' + NUM_CLASSES = 2 + ID_FIELD = Field(sequential=False, tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True) + AID_FIELD = Field(sequential=False, use_vocab=False, batch_first=True) + TEXT_FIELD = Field(batch_first=True, tokenize=lambda x: x) # tokenizer is identity since we already tokenized it to compute external features + EXT_FEATS_FIELD = Field(tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, tokenize=lambda x: x) + LABEL_FIELD = Field(sequential=False, tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, postprocessing=Pipeline(get_class_probs)) + RAW_TEXT_FIELD = RawField() + + @staticmethod + def sort_key(ex): + return len(ex.sentence_1) + + def __init__(self, path): + """ + Create a Quora dataset instance + """ + super(Quora, self).__init__(path) + + @classmethod + def splits(cls, path, train='train', validation='dev', test='test', **kwargs): + return super(Quora, cls).splits(path, train=train, validation=validation, test=test, **kwargs) + + @classmethod + def iters(cls, path, vectors_name, vectors_cache, batch_size=64, shuffle=True, device=0, vectors=None, + unk_init=torch.Tensor.zero_): + """ + :param path: directory containing train, test, dev files + :param vectors_name: name of word vectors file + :param vectors_dir: directory containing word vectors file + :param batch_size: batch size + :param device: GPU device + :param vectors: custom vectors - either predefined torchtext vectors or your own custom Vector classes + :param pt_file: load cached embedding file from disk if it is true + :param unk_init: function used to generate vector for OOV words + :return: + """ + + if vectors is None: + vectors = Vectors(name=vectors_name, cache=vectors_cache, unk_init=unk_init) + + train, validation, test = cls.splits(path) + + cls.LABEL_FIELD.build_vocab(train, validation, test) + cls.TEXT_FIELD.build_vocab(train, validation, test, vectors=vectors) + return BucketIterator.splits((train, validation, test), batch_size=batch_size, repeat=False, shuffle=shuffle, + sort_within_batch=True, device=device) \ No newline at end of file diff --git a/datasets/sick.py b/datasets/sick.py index a78f686..5daf8b9 100644 --- a/datasets/sick.py +++ b/datasets/sick.py @@ -68,4 +68,4 @@ class SICK(CastorPairDataset): cls.TEXT_FIELD.build_vocab(train, val, test, vectors=vectors) return BucketIterator.splits((train, val, test), batch_size=batch_size, repeat=False, shuffle=shuffle, - sort_within_batch=True, device=device) + sort_within_batch=True, device=device) \ No newline at end of file diff --git a/datasets/snli.py b/datasets/snli.py new file mode 100644 index 0000000..8ba5c1e --- /dev/null +++ b/datasets/snli.py @@ -0,0 +1,66 @@ +import numpy as np +import torch +from torchtext.data.field import Field, RawField +from torchtext.data.iterator import BucketIterator +from torchtext.vocab import Vectors +from torchtext.data.pipeline import Pipeline + +from datasets.castor_dataset import CastorPairDataset + +def get_class_probs(sim, *args): + """ + Convert a single label into class probabilities. + """ + class_probs = np.zeros(SNLI.NUM_CLASSES) + class_probs[int(sim)] = 1 + return class_probs + + +class SNLI(CastorPairDataset): + NAME = 'SNLI' + NUM_CLASSES = 3 + ID_FIELD = Field(sequential=False, tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True) + AID_FIELD = Field(sequential=False, use_vocab=False, batch_first=True) + TEXT_FIELD = Field(batch_first=True, tokenize=lambda x: x) # tokenizer is identity since we already tokenized it to compute external features + EXT_FEATS_FIELD = Field(tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, tokenize=lambda x: x) + LABEL_FIELD = Field(sequential=False, tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, postprocessing=Pipeline(get_class_probs)) + RAW_TEXT_FIELD = RawField() + + @staticmethod + def sort_key(ex): + return len(ex.sentence_1) + + def __init__(self, path): + """ + Create a SNLI dataset instance + """ + super(SNLI, self).__init__(path) + + @classmethod + def splits(cls, path, train='train', validation='dev', test='test', **kwargs): + return super(SNLI, cls).splits(path, train=train, validation=validation, test=test, **kwargs) + + @classmethod + def iters(cls, path, vectors_name, vectors_cache, batch_size=64, shuffle=True, device=0, vectors=None, + unk_init=torch.Tensor.zero_): + """ + :param path: directory containing train, test, dev files + :param vectors_name: name of word vectors file + :param vectors_dir: directory containing word vectors file + :param batch_size: batch size + :param device: GPU device + :param vectors: custom vectors - either predefined torchtext vectors or your own custom Vector classes + :param pt_file: load cached embedding file from disk if it is true + :param unk_init: function used to generate vector for OOV words + :return: + """ + + if vectors is None: + vectors = Vectors(name=vectors_name, cache=vectors_cache, unk_init=unk_init) + + train, validation, test = cls.splits(path) + + cls.LABEL_FIELD.build_vocab(train, validation, test) + cls.TEXT_FIELD.build_vocab(train, validation, test, vectors=vectors) + return BucketIterator.splits((train, validation, test), batch_size=batch_size, repeat=False, shuffle=shuffle, + sort_within_batch=True, device=device) \ No newline at end of file diff --git a/datasets/sts2014.py b/datasets/sts2014.py new file mode 100644 index 0000000..68edf3c --- /dev/null +++ b/datasets/sts2014.py @@ -0,0 +1,72 @@ +import math + +import numpy as np +import torch +from torchtext.data.field import Field, RawField +from torchtext.data.iterator import BucketIterator +from torchtext.data.pipeline import Pipeline +from torchtext.vocab import Vectors + +from datasets.castor_dataset import CastorPairDataset + + +def get_class_probs(sim, *args): + """ + Convert a single label into class probabilities. + """ + class_probs = np.zeros(STS2014.NUM_CLASSES) + ceil, floor = math.ceil(sim), math.floor(sim) + + if ceil == floor: + class_probs[ceil] = 1 + else: + class_probs[floor] = ceil - sim + class_probs[ceil] = sim - floor + + return class_probs + + +class STS2014(CastorPairDataset): + NAME = 'sts2014' + NUM_CLASSES = 6 + ID_FIELD = Field(sequential=False, use_vocab=False, batch_first=True) + TEXT_FIELD = Field(batch_first=True, tokenize=lambda x: x) # tokenizer is identity since we already tokenized it to compute external features + EXT_FEATS_FIELD = Field(tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, tokenize=lambda x: x) + LABEL_FIELD = Field(sequential=False, tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, postprocessing=Pipeline(get_class_probs)) + RAW_TEXT_FIELD = RawField() + + @staticmethod + def sort_key(ex): + return len(ex.sentence_1) + + def __init__(self, path): + """ + Create a STS2014 dataset instance + """ + super(STS2014, self).__init__(path) + + @classmethod + def splits(cls, path, train='train', validation='dev', test='test', **kwargs): + return super(STS2014, cls).splits(path, train=train, validation=validation, test=test, **kwargs) + + @classmethod + def iters(cls, path, vectors_name, vectors_cache, batch_size=64, shuffle=True, device=0, vectors=None, unk_init=torch.Tensor.zero_): + """ + :param path: directory containing train, test, dev files + :param vectors_name: name of word vectors file + :param vectors_cache: path to word vectors file + :param batch_size: batch size + :param device: GPU device + :param vectors: custom vectors - either predefined torchtext vectors or your own custom Vector classes + :param unk_init: function used to generate vector for OOV words + :return: + """ + if vectors is None: + vectors = Vectors(name=vectors_name, cache=vectors_cache, unk_init=unk_init) + + train, val, test = cls.splits(path) + + cls.TEXT_FIELD.build_vocab(train, val, test, vectors=vectors) + + return BucketIterator.splits((train, val, test), batch_size=batch_size, repeat=False, shuffle=shuffle, + sort_within_batch=True, device=device) \ No newline at end of file