| |
| |
| |
| |
|
|
| import numpy as np |
| import torch |
| from typing import Dict |
|
|
| from fairseq.data.monolingual_dataset import MonolingualDataset |
|
|
| from . import FairseqDataset |
|
|
|
|
| class LMContextWindowDataset(FairseqDataset): |
| """ |
| Wraps a MonolingualDataset and provides more context for evaluation. |
| |
| Each item in the new dataset will have a maximum size of |
| ``tokens_per_sample + context_window``. |
| |
| Args: |
| dataset: dataset to wrap |
| tokens_per_sample (int): the max number of tokens in each dataset item |
| context_window (int): the number of accumulated tokens to add to each |
| dataset item |
| pad_idx (int): padding symbol |
| """ |
|
|
| def __init__( |
| self, |
| dataset: MonolingualDataset, |
| tokens_per_sample: int, |
| context_window: int, |
| pad_idx: int, |
| ): |
| assert context_window > 0 |
| self.dataset = dataset |
| self.tokens_per_sample = tokens_per_sample |
| self.context_window = context_window |
| self.pad_idx = pad_idx |
| self.prev_tokens = np.empty([0]) |
|
|
| def __getitem__(self, index): |
| return self.dataset[index] |
|
|
| def __len__(self): |
| return len(self.dataset) |
|
|
| def collater(self, samples) -> Dict: |
| sample = self.dataset.collater(samples) |
|
|
| pad = self.pad_idx |
| max_sample_len = self.tokens_per_sample + self.context_window |
|
|
| bsz, tsz = sample["net_input"]["src_tokens"].shape |
| start_idxs = [0] * bsz |
| toks = sample["net_input"]["src_tokens"] |
| lengths = sample["net_input"]["src_lengths"] |
| tgt = sample["target"] |
| new_toks = np.empty([bsz, tsz + self.context_window], dtype=np.int64) |
| new_tgt = np.full([bsz, tsz + self.context_window], pad, dtype=np.int64) |
| sample_lens = toks.ne(pad).long().sum(dim=1).cpu() |
| for i in range(bsz): |
| sample_len = sample_lens[i] |
| extra = len(self.prev_tokens) + sample_len - max_sample_len |
| if extra > 0: |
| self.prev_tokens = self.prev_tokens[extra:] |
| pads = np.full(self.context_window - len(self.prev_tokens), pad) |
| new_toks[i] = np.concatenate([self.prev_tokens, toks[i].numpy(), pads]) |
| new_tgt[ |
| i, len(self.prev_tokens) : len(self.prev_tokens) + len(tgt[i]) |
| ] = tgt[i] |
| start_idxs[i] = len(self.prev_tokens) |
| lengths[i] += len(self.prev_tokens) |
| self.prev_tokens = new_toks[i][new_toks[i] != pad][-self.context_window :] |
| sample["net_input"]["src_tokens"] = torch.from_numpy(new_toks) |
| sample["target"] = torch.from_numpy(new_tgt) |
| sample["start_indices"] = start_idxs |
| return sample |
|
|
| def num_tokens(self, index): |
| return self.dataset.num_tokens(index) |
|
|
| def size(self, index): |
| return self.dataset.size(index) |
|
|
| def ordered_indices(self): |
| |
| return np.arange(len(self.dataset)) |
|
|
| @property |
| def supports_prefetch(self): |
| return getattr(self.dataset, "supports_prefetch", False) |
|
|
| def prefetch(self, indices): |
| return self.dataset.prefetch(indices) |
|
|
