| 20 | |
| 21 | |
| 22 | class Code2Seq(LightningModule): |
| 23 | def __init__( |
| 24 | self, |
| 25 | model_config: DictConfig, |
| 26 | optimizer_config: DictConfig, |
| 27 | vocabulary: Vocabulary, |
| 28 | teacher_forcing: float = 0.0, |
| 29 | ): |
| 30 | super().__init__() |
| 31 | self.save_hyperparameters() |
| 32 | self._optim_config = optimizer_config |
| 33 | self._vocabulary = vocabulary |
| 34 | |
| 35 | if vocabulary.SOS not in vocabulary.label_to_id: |
| 36 | raise ValueError(f"Can't find SOS token in label to id vocabulary") |
| 37 | |
| 38 | self.__pad_idx = vocabulary.label_to_id[vocabulary.PAD] |
| 39 | eos_idx = vocabulary.label_to_id[vocabulary.EOS] |
| 40 | ignore_idx = [vocabulary.label_to_id[vocabulary.SOS], vocabulary.label_to_id[vocabulary.UNK]] |
| 41 | metrics: Dict[str, Metric] = { |
| 42 | f"{holdout}_f1": SequentialF1Score(pad_idx=self.__pad_idx, eos_idx=eos_idx, ignore_idx=ignore_idx) |
| 43 | for holdout in ["train", "val", "test"] |
| 44 | } |
| 45 | id2label = {v: k for k, v in vocabulary.label_to_id.items()} |
| 46 | metrics.update( |
| 47 | {f"{holdout}_chrf": ChrF(id2label, ignore_idx + [self.__pad_idx, eos_idx]) for holdout in ["val", "test"]} |
| 48 | ) |
| 49 | self.__metrics = MetricCollection(metrics) |
| 50 | |
| 51 | self._encoder = self._get_encoder(model_config) |
| 52 | decoder_step = LSTMDecoderStep(model_config, len(vocabulary.label_to_id), self.__pad_idx) |
| 53 | self._decoder = Decoder( |
| 54 | decoder_step, len(vocabulary.label_to_id), vocabulary.label_to_id[vocabulary.SOS], teacher_forcing |
| 55 | ) |
| 56 | |
| 57 | self.__loss = SequenceCrossEntropyLoss(self.__pad_idx, reduction="batch-mean") |
| 58 | |
| 59 | @property |
| 60 | def vocabulary(self) -> Vocabulary: |
| 61 | return self._vocabulary |
| 62 | |
| 63 | def _get_encoder(self, config: DictConfig) -> nn.Module: |
| 64 | return PathEncoder( |
| 65 | config, |
| 66 | len(self._vocabulary.token_to_id), |
| 67 | self._vocabulary.token_to_id[Vocabulary.PAD], |
| 68 | len(self._vocabulary.node_to_id), |
| 69 | self._vocabulary.node_to_id[Vocabulary.PAD], |
| 70 | ) |
| 71 | |
| 72 | # ========== Main PyTorch-Lightning hooks ========== |
| 73 | |
| 74 | def configure_optimizers(self) -> Tuple[List[Optimizer], List[_LRScheduler]]: |
| 75 | return configure_optimizers_alon(self._optim_config, self.parameters()) |
| 76 | |
| 77 | def forward( # type: ignore |
| 78 | self, |
| 79 | from_token: torch.Tensor, |