| 102 | raise NotImplementedError(f'Unknown dataset {self.dataset}. Only WebNLG dataset has been implemented') |
| 103 | |
| 104 | def setup(self, stage=None): |
| 105 | |
| 106 | edge_classes_path = os.path.join(self.output_path, 'edge.classes') |
| 107 | |
| 108 | if stage == 'fit': |
| 109 | text_path = os.path.join(self.output_path, 'train.text') |
| 110 | graph_path = os.path.join(self.output_path, 'train.graph') |
| 111 | self.dataset_train = GraphDataset(tokenizer=self.tokenizer, |
| 112 | text_data_path=text_path, |
| 113 | graph_data_path=graph_path, |
| 114 | edge_classes_path=edge_classes_path, |
| 115 | max_nodes=self.max_nodes, |
| 116 | max_edges=self.max_edges, |
| 117 | edges_as_classes=self.edges_as_classes) |
| 118 | elif stage == 'validate': |
| 119 | text_path = os.path.join(self.output_path, 'dev.text') |
| 120 | graph_path = os.path.join(self.output_path, 'dev.graph') |
| 121 | self.dataset_dev = GraphDataset(tokenizer=self.tokenizer, |
| 122 | text_data_path=text_path, |
| 123 | graph_data_path=graph_path, |
| 124 | edge_classes_path=edge_classes_path, |
| 125 | max_nodes=self.max_nodes, |
| 126 | max_edges=self.max_edges, |
| 127 | edges_as_classes=self.edges_as_classes) |
| 128 | |
| 129 | elif stage == 'test': |
| 130 | text_path = os.path.join(self.output_path, 'test.text') |
| 131 | graph_path = os.path.join(self.output_path, 'test.graph') |
| 132 | self.dataset_test = GraphDataset(tokenizer=self.tokenizer, |
| 133 | text_data_path=text_path, |
| 134 | graph_data_path=graph_path, |
| 135 | edge_classes_path=edge_classes_path, |
| 136 | max_nodes=self.max_nodes, |
| 137 | max_edges=self.max_edges, |
| 138 | edges_as_classes=self.edges_as_classes) |
| 139 | else: |
| 140 | return |
| 141 | |
| 142 | def train_dataloader(self): |
| 143 | return DataLoader(self.dataset_train, |