| 14 | |
| 15 | |
| 16 | class PathContextDataModule(LightningDataModule): |
| 17 | _train = "train" |
| 18 | _val = "val" |
| 19 | _test = "test" |
| 20 | |
| 21 | def __init__(self, data_dir: str, config: DictConfig, is_class: bool = False): |
| 22 | super().__init__() |
| 23 | self._config = config |
| 24 | self._data_dir = data_dir |
| 25 | self._name = basename(data_dir) |
| 26 | self._is_class = is_class |
| 27 | |
| 28 | self._vocabulary = self.setup_vocabulary() |
| 29 | |
| 30 | @property |
| 31 | def vocabulary(self) -> Vocabulary: |
| 32 | if self._vocabulary is None: |
| 33 | raise RuntimeError(f"Setup data module for initializing vocabulary") |
| 34 | return self._vocabulary |
| 35 | |
| 36 | def prepare_data(self): |
| 37 | if exists(self._data_dir): |
| 38 | print(f"Dataset is already downloaded") |
| 39 | return |
| 40 | if "url" not in self._config: |
| 41 | raise ValueError(f"Config doesn't contain url for, can't download it automatically") |
| 42 | download_dataset(self._config.url, self._data_dir, self._name) |
| 43 | |
| 44 | def setup_vocabulary(self) -> Vocabulary: |
| 45 | vocabulary_path = join(self._data_dir, Vocabulary.vocab_filename) |
| 46 | if not exists(vocabulary_path): |
| 47 | print("Can't find vocabulary, collect it from train holdout") |
| 48 | build_from_scratch(join(self._data_dir, f"{self._train}.c2s"), Vocabulary) |
| 49 | return Vocabulary(vocabulary_path, self._config.labels_count, self._config.tokens_count, self._is_class) |
| 50 | |
| 51 | @staticmethod |
| 52 | def collate_wrapper(batch: List[Optional[LabeledPathContext]]) -> BatchedLabeledPathContext: |
| 53 | return BatchedLabeledPathContext(batch) |
| 54 | |
| 55 | def _create_dataset(self, holdout_file: str, random_context: bool) -> PathContextDataset: |
| 56 | if self._vocabulary is None: |
| 57 | raise RuntimeError(f"Setup vocabulary before creating data loaders") |
| 58 | return PathContextDataset(holdout_file, self._config, self._vocabulary, random_context) |
| 59 | |
| 60 | def _shared_dataloader(self, holdout: str) -> DataLoader: |
| 61 | if self._vocabulary is None: |
| 62 | raise RuntimeError(f"Setup vocabulary before creating data loaders") |
| 63 | |
| 64 | holdout_file = join(self._data_dir, f"{holdout}.c2s") |
| 65 | random_context = self._config.random_context if holdout == self._train else False |
| 66 | dataset = self._create_dataset(holdout_file, random_context) |
| 67 | |
| 68 | batch_size = self._config.batch_size if holdout == self._train else self._config.test_batch_size |
| 69 | shuffle = holdout == self._train |
| 70 | |
| 71 | return DataLoader( |
| 72 | dataset, |
| 73 | batch_size, |
no outgoing calls