Update return types

modeling_cased.py

@@ -197,8 +197,7 @@ class CaSEDModel(PreTrainedModel):
 
         return vocabularies
 
-    @torch.no_grad()
-    def forward(self, images: dict, alpha: Optional[float] = None) -> torch.Tensor():
+    def forward(self, images: dict, alpha: Optional[float] = None) -> torch.Tensor:
         """Forward pass.
 
         Args:
@@ -248,8 +247,8 @@ class CaSEDModel(PreTrainedModel):
             vocabularies.append(vocabulary)
 
         # get the scores
-        samples_p = torch.stack(samples_p, dim=0)
-        scores = sample_p.cpu()
+        samples_p = torch.cat(samples_p, dim=0)
+        scores = samples_p.cpu()
 
         # define the results
         results = {"vocabularies": vocabularies, "scores": scores}

transforms_cased.py

@@ -28,7 +28,7 @@ class BaseTextTransform(ABC):
     """Base class for string transforms."""
 
     @abstractmethod
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         raise NotImplementedError
 
     def __repr__(self) -> str:
@@ -38,7 +38,7 @@ class BaseTextTransform(ABC):
 class DropFileExtensions(BaseTextTransform):
     """Remove file extensions from the input text."""
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove file extensions from.
@@ -51,7 +51,7 @@ class DropFileExtensions(BaseTextTransform):
 class DropNonAlpha(BaseTextTransform):
     """Remove non-alpha words from the input text."""
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove non-alpha words from.
@@ -72,7 +72,7 @@ class DropShortWords(BaseTextTransform):
         super().__init__()
         self.min_length = min_length
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove short words from.
@@ -92,7 +92,7 @@ class DropSpecialCharacters(BaseTextTransform):
     hyphen, period, apostrophe, or ampersand.
     """
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove special characters from.
@@ -108,7 +108,7 @@ class DropTokens(BaseTextTransform):
     Tokens are defined as strings enclosed in angle brackets, e.g. <token>.
     """
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove tokens from.
@@ -121,7 +121,7 @@ class DropTokens(BaseTextTransform):
 class DropURLs(BaseTextTransform):
     """Remove URLs from the input text."""
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove URLs from.
@@ -142,7 +142,7 @@ class DropWords(BaseTextTransform):
         self.words = words
         self.pattern = r"\b(?:{})\b".format("|".join(words))
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove words from.
@@ -177,7 +177,7 @@ class FilterPOS(BaseTextTransform):
         elif engine == "flair":
             self.tagger = SequenceTagger.load("flair/pos-english-fast").predict
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove words with specific POS tags from.
@@ -234,7 +234,7 @@ class FrequencyMinWordCount(BaseTextTransform):
         super().__init__()
         self.min_count = min_count
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove infrequent words from.
@@ -270,7 +270,7 @@ class FrequencyTopK(BaseTextTransform):
         super().__init__()
         self.top_k = top_k
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove infrequent words from.
@@ -297,7 +297,7 @@ class FrequencyTopK(BaseTextTransform):
 class ReplaceSeparators(BaseTextTransform):
     """Replace underscores and dashes with spaces."""
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to replace separators in.
@@ -313,7 +313,7 @@ class ReplaceSeparators(BaseTextTransform):
 class RemoveDuplicates(BaseTextTransform):
     """Remove duplicate words from the input text."""
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to remove duplicate words from.
@@ -337,7 +337,11 @@ class TextCompose:
     def __init__(self, transforms: list[BaseTextTransform]) -> None:
         self.transforms = transforms
 
-    def __call__(self, text: Union[str, list[str]]) -> Any:
+    def __call__(self, text: Union[str, list[str]]) -> list[str]:
+        """
+        Args:
+            text (Union[str, list[str]]): Text to transform.
+        """
         if isinstance(text, list):
             text = " ".join(text)
 
@@ -357,7 +361,7 @@ class TextCompose:
 class ToLowercase(BaseTextTransform):
     """Convert text to lowercase."""
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to convert to lowercase.
@@ -374,7 +378,7 @@ class ToSingular(BaseTextTransform):
         super().__init__()
         self.transform = inflect.engine().singular_noun
 
-    def __call__(self, text: str):
+    def __call__(self, text: str) -> str:
         """
         Args:
             text (str): Text to convert to singular form.