mirror of
https://github.com/wassname/Castor.git
synced 2026-09-10 11:40:44 +08:00
100 lines
4.6 KiB
Python
100 lines
4.6 KiB
Python
import os
|
|
|
|
import torch
|
|
|
|
from torchtext.data.dataset import Dataset
|
|
from torchtext.data.example import Example
|
|
from torchtext.data.field import Field
|
|
from torchtext.data.iterator import BucketIterator
|
|
from torchtext.data.iterator import Iterator
|
|
from torchtext.vocab import Vectors
|
|
from torchtext.data import Pipeline
|
|
|
|
from datasets.castor_dataset import CastorPairDataset
|
|
from datasets.idf_utils import get_pairwise_word_to_doc_freq, get_pairwise_overlap_features
|
|
|
|
|
|
class TWITTER(Dataset):
|
|
NAME = 'twitter'
|
|
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,
|
|
postprocessing=Pipeline(lambda arr, _, train: [float(y) for y in arr]))
|
|
LABEL_FIELD = Field(sequential=False, use_vocab=False, batch_first=True)
|
|
VOCAB_SIZE = 0
|
|
|
|
@staticmethod
|
|
def sort_key(ex):
|
|
return len(ex.sentence_1)
|
|
|
|
def __init__(self, pardir, subdirs):
|
|
"""
|
|
Create a Twitter dataset instance.
|
|
"""
|
|
fields = [('id', self.ID_FIELD), ('sentence_1', self.TEXT_FIELD), ('sentence_2', self.TEXT_FIELD), ('ext_feats',
|
|
self.EXT_FEATS_FIELD), ('label', self.LABEL_FIELD), ('aid', self.AID_FIELD)]
|
|
|
|
examples = []
|
|
for subdir in subdirs:
|
|
path = os.path.join(pardir, subdir)
|
|
with open(os.path.join(path, 'a.toks'), 'r') as f1, open(os.path.join(path, 'b.toks'), 'r') as f2:
|
|
sent_list_1 = [l.rstrip('.\n').split(' ') for l in f1]
|
|
sent_list_2 = [l.rstrip('.\n').split(' ') for l in f2]
|
|
|
|
word_to_doc_cnt = get_pairwise_word_to_doc_freq(sent_list_1, sent_list_2)
|
|
overlap_feats = get_pairwise_overlap_features(sent_list_1, sent_list_2, word_to_doc_cnt)
|
|
|
|
for subdir_i, subdir in enumerate(subdirs):
|
|
path = os.path.join(pardir, subdir)
|
|
with open(os.path.join(path, 'id.txt'), 'r') as id_file, open(os.path.join(path, 'sim.txt'), 'r') as label_file:
|
|
for i, (pair_id, l1, l2, ext_feats, label) in enumerate(zip(id_file, sent_list_1, sent_list_2, overlap_feats, label_file)):
|
|
pair_id = pair_id.rstrip('.\n')
|
|
label = label.rstrip('.\n')
|
|
example_list = [pair_id, l1, l2, ext_feats, label, (subdir_i) * 100000 + (i + 1)]
|
|
example = Example.fromlist(example_list, fields)
|
|
examples.append(example)
|
|
|
|
super(TWITTER, self).__init__(examples, fields)
|
|
|
|
@classmethod
|
|
def splits(cls, path, train_paths, test_paths, **kwargs):
|
|
train_data = cls(path, train_paths, **kwargs)
|
|
test_data = cls(path, test_paths, **kwargs)
|
|
return train_data, test_data
|
|
|
|
@classmethod
|
|
def set_vectors(cls, field, vector_path):
|
|
return CastorPairDataset.set_vectors(field, vector_path)
|
|
|
|
@classmethod
|
|
def iters(cls, path, train_dirs, test_dirs, vectors_name, vectors_dir, batch_size=64, shuffle=True, device=0, pt_file=False, vectors=None, unk_init=torch.Tensor.zero_):
|
|
"""
|
|
:param path: directory containing train, test, dev files
|
|
:param train_dirs: list of directory names used for training
|
|
:param test_dirs: list of directory name used for testing
|
|
: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 unk_init: function used to generate vector for OOV words
|
|
:return:
|
|
"""
|
|
|
|
train, test = cls.splits(path, train_dirs, test_dirs)
|
|
if not pt_file:
|
|
if vectors is None:
|
|
vectors = Vectors(name=vectors_name, cache=vectors_dir, unk_init=unk_init)
|
|
cls.TEXT_FIELD.build_vocab(train, test, vectors=vectors)
|
|
else:
|
|
cls.TEXT_FIELD.build_vocab(train, test)
|
|
cls.TEXT_FIELD = cls.set_vectors(cls.TEXT_FIELD, os.path.join(vectors_dir, vectors_name))
|
|
|
|
cls.LABEL_FIELD.build_vocab(train, test)
|
|
|
|
cls.VOCAB_SIZE = len(cls.TEXT_FIELD.vocab)
|
|
|
|
return BucketIterator.splits((train, test), batch_size=batch_size, repeat=False, shuffle=shuffle, device=device)
|