Merge pull request #588 from Butanium/fix/prepare-input-typing

fix(typing): correct _prepare_input type annotations in LanguageModel

Author: JadenFiotto-Kaufman

src/nnsight/modeling/language.py

@@ -207,26 +207,24 @@ def _tokenize(
 
     def _prepare_input(
         self,
-        *inputs: Tuple[
-            Union[
-                str,
-                List[str],
-                List[List[str]],
-                List[int],
-                List[List[int]],
-                torch.Tensor,
-                List[torch.Tensor],
-                Dict[str, Any],
-                BatchEncoding,
-            ]
+        *inputs: Union[
+            str,
+            List[str],
+            List[List[str]],
+            List[int],
+            List[List[int]],
+            torch.Tensor,
+            List[torch.Tensor],
+            Dict[str, Any],
+            BatchEncoding,
         ],
         input_ids: Union[
             List[int], List[List[int]], torch.Tensor, List[torch.Tensor]
         ] = None,
         labels: Any = None,
         attention_mask: Any = None,
         **kwargs,
-    ) -> Tuple[BatchEncoding, int]:
+    ) -> Tuple[Tuple[()], Dict[str, Any]]:
 
         if input_ids is not None: