class PDataset(Dataset):
NUM_LABELS = NUM_LABELS
LABEL2INDEX = LABEL2INDEX
INDEX2LABEL = INDEX2LABEL
def load_dataset(self, path):
df = pd.read_csv(path)
df['lesson_id'] = df['lesson_id'].apply(lambda lesson_id: self.LABEL2INDEX[lesson_id])
return df
def __init__(self, dataset_path, tokenizer, no_special_token=False, *args, **kwargs):
self.data = self.load_dataset(dataset_path)
self.tokenizer = tokenizer
self.no_special_token = no_special_token
def __getitem__(self, index):
data = self.data.loc[index,:]
content, lesson_id = data['content'], data['lesson_id']
subwords = self.tokenizer.encode(content, add_special_tokens=not self.no_special_token)
return np.array(subwords), np.array(lesson_id), data['content']
def __len__(self):
return len(self.data)
class PContentDataLoader(DataLoader):
def __init__(self, max_seq_len=512, *args, **kwargs):
super(PContentDataLoader, self).__init__(*args, **kwargs)
self.collate_fn = self._collate_fn
self.max_seq_len = max_seq_len
def _collate_fn(self, batch):
batch_size = len(batch)
max_seq_len = max(map(lambda x: len(x[0]), batch))
max_seq_len = min(self.max_seq_len, max_seq_len)
subword_batch = np.zeros((batch_size, max_seq_len), dtype=np.int64)
mask_batch = np.zeros((batch_size, max_seq_len), dtype=np.float32)
sentiment_batch = np.zeros((batch_size, 1), dtype=np.int64)
seq_list = []
for i, (subwords, sentiment, raw_seq) in enumerate(batch):
subwords = subwords[:max_seq_len]
subword_batch[i,:len(subwords)] = subwords
mask_batch[i,:len(subwords)] = 1
sentiment_batch[i,0] = sentiment
seq_list.append(raw_seq)
return subword_batch, mask_batch, sentiment_batch, seq_list