🪡 Fix typing issues.

This should silence warnings.

perbert/constants.py

@@ -2,8 +2,7 @@
 
 import os
 from enum import Enum
-
-from typing_extensions import Self
+from typing import Any
 
 
 class StrEnum(str, Enum):
@@ -13,7 +12,7 @@ def __str__(self) -> str:
     def __hash__(self) -> int:
         return hash(str(self))
 
-    def __eq__(self, other: Self | str) -> bool:
+    def __eq__(self, other: Any) -> bool:
         return str(self) == str(other)
 
 

perbert/models/models.py

@@ -1,3 +1,4 @@
+# pyright: reportIncompatibleMethodOverride=false
 from __future__ import annotations
 
 import typing
@@ -58,8 +59,8 @@ def __init__(self, cfg: DictConfig) -> None:
     def bert_config(self) -> BertConfig:
         return self.lm.bert.config
 
-    def forward(self, **kwargs: Any) -> BertOutput:
-        return self.lm(**kwargs)
+    def forward(self, *args: Any, **kwargs: Any) -> BertOutput:
+        return self.lm(*args, **kwargs)
 
     def _step(self, batch: BatchEncoding, batch_idx: int, name: str) -> Tensor:
         loguru.logger.trace("{} step batch: {}", name, batch_idx)
@@ -91,6 +92,7 @@ def training_step(self, batch: BatchEncoding, batch_idx: int) -> Tensor:
 
     @torch.no_grad()
     def test_step(self, batch: BatchEncoding, batch_idx: int) -> Tensor:
+        super().test_step
         return self._step(batch, batch_idx=batch_idx, name="test")
 
     @torch.no_grad()