|
""" |
|
|
|
""" |
|
from torch.utils.data import IterableDataset |
|
import numpy as np |
|
from typing import List |
|
from ..readers import InputExample |
|
import logging |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
class SentenceLabelDataset(IterableDataset): |
|
""" |
|
This dataset can be used for some specific Triplet Losses like BATCH_HARD_TRIPLET_LOSS which requires |
|
multiple examples with the same label in a batch. |
|
|
|
It draws n consecutive, random and unique samples from one label at a time. This is repeated for each label. |
|
|
|
Labels with fewer than n unique samples are ignored. |
|
This also applied to drawing without replacement, once less than n samples remain for a label, it is skipped. |
|
|
|
This *DOES NOT* check if there are more labels than the batch is large or if the batch size is divisible |
|
by the samples drawn per label. |
|
""" |
|
def __init__(self, examples: List[InputExample], samples_per_label: int = 2, with_replacement: bool = False): |
|
""" |
|
Creates a LabelSampler for a SentenceLabelDataset. |
|
|
|
:param examples: |
|
a list with InputExamples |
|
:param samples_per_label: |
|
the number of consecutive, random and unique samples drawn per label. Batch size should be a multiple of samples_per_label |
|
:param with_replacement: |
|
if this is True, then each sample is drawn at most once (depending on the total number of samples per label). |
|
if this is False, then one sample can be drawn in multiple draws, but still not multiple times in the same |
|
drawing. |
|
""" |
|
super().__init__() |
|
|
|
self.samples_per_label = samples_per_label |
|
|
|
|
|
label2ex = {} |
|
for example in examples: |
|
if example.label not in label2ex: |
|
label2ex[example.label] = [] |
|
label2ex[example.label].append(example) |
|
|
|
|
|
self.grouped_inputs = [] |
|
self.groups_right_border = [] |
|
num_labels = 0 |
|
|
|
for label, label_examples in label2ex.items(): |
|
if len(label_examples) >= self.samples_per_label: |
|
self.grouped_inputs.extend(label_examples) |
|
self.groups_right_border.append(len(self.grouped_inputs)) |
|
num_labels += 1 |
|
|
|
self.label_range = np.arange(num_labels) |
|
self.with_replacement = with_replacement |
|
np.random.shuffle(self.label_range) |
|
|
|
logger.info("SentenceLabelDataset: {} examples, from which {} examples could be used (those labels appeared at least {} times). {} different labels found.".format(len(examples), len(self.grouped_inputs), self.samples_per_label, num_labels )) |
|
|
|
def __iter__(self): |
|
label_idx = 0 |
|
count = 0 |
|
already_seen = {} |
|
while count < len(self.grouped_inputs): |
|
label = self.label_range[label_idx] |
|
if label not in already_seen: |
|
already_seen[label] = set() |
|
|
|
left_border = 0 if label == 0 else self.groups_right_border[label-1] |
|
right_border = self.groups_right_border[label] |
|
|
|
if self.with_replacement: |
|
selection = np.arange(left_border, right_border) |
|
else: |
|
selection = [i for i in np.arange(left_border, right_border) if i not in already_seen[label]] |
|
|
|
if len(selection) >= self.samples_per_label: |
|
for element_idx in np.random.choice(selection, self.samples_per_label, replace=False): |
|
count += 1 |
|
already_seen[label].add(element_idx) |
|
yield self.grouped_inputs[element_idx] |
|
|
|
label_idx += 1 |
|
if label_idx >= len(self.label_range): |
|
label_idx = 0 |
|
already_seen = {} |
|
np.random.shuffle(self.label_range) |
|
|
|
def __len__(self): |
|
return len(self.grouped_inputs) |
