Merge pull request #157 from ChEB-AI/fix/load-electra-checkpoints

fix checkpoint loading for electra, return attentions

Author: sfluegel05

chebai/models/electra.py

@@ -203,6 +203,7 @@ def _process_batch(self, batch: Dict[str, Any], batch_idx: int) -> Dict[str, Any
             )
             * CLS_TOKEN
         )
+        model_kwargs["output_attentions"] = True
         return dict(
             features=torch.cat((cls_tokens, batch.x), dim=1),
             labels=batch.y,

chebai/result/prediction.py

@@ -58,9 +58,18 @@ def __init__(
         self._model_hparams = ckpt_file["hyper_parameters"]
         self._model_hparams.pop("_instantiator", None)
         self._model_hparams.pop("classes_txt_file_path", None)
-        self._model = ChebaiBaseNet.load_from_checkpoint(
-            checkpoint_path, map_location=self.device
-        )
+        try:
+            self._model = ChebaiBaseNet.load_from_checkpoint(
+                checkpoint_path,
+                map_location=self.device,
+            )
+        except Exception:
+            # models trained on a pretrained checkpoint have an additional path argument that we need to set to None
+            self._model = ChebaiBaseNet.load_from_checkpoint(
+                checkpoint_path,
+                map_location=self.device,
+                pretrained_checkpoint=None,
+            )
         assert (
             isinstance(self._model, ChebaiBaseNet)
             and type(self._model) is not ChebaiBaseNet