Merge pull request #159 from dreadnode/fix/transformers_type_check

fix: typing check for transformers

Author: moohax

docs/api/data.mdx

@@ -257,8 +257,8 @@ chats\_to\_tokens
 
 ```python
 chats_to_tokens(
-    chat: Chat | None,
-    tokenizer: AutoTokenizer,
+    chat: Chat,
+    tokenizer: PreTrainedTokenizerBase,
     *,
     apply_chat_template_kwargs: dict[str, Any]
     | None = None,
@@ -272,10 +272,10 @@ Transform a chat into a tokenized format with structured slices.
 **Parameters:**
 
 * **`chat`**
-  (`Chat | None`)
+  (`Chat`)
   –The chat object to tokenize.
 * **`tokenizer`**
-  (`AutoTokenizer`)
+  (`PreTrainedTokenizerBase`)
   –The tokenizer to use for encoding and decoding.
 
 **Returns:**
@@ -286,8 +286,8 @@ Transform a chat into a tokenized format with structured slices.
 <Accordion title="Source code in rigging/data.py" icon="code">
 ```python
 async def chats_to_tokens(
-    chat: Chat | None,
-    tokenizer: AutoTokenizer,
+    chat: Chat,
+    tokenizer: "PreTrainedTokenizerBase",
     *,
     apply_chat_template_kwargs: dict[str, t.Any] | None = None,
     encode_kwargs: dict[str, t.Any] | None = None,
@@ -323,8 +323,9 @@ async def chats_to_tokens(
         if chat.params and chat.params.tools
         else None
     )
+    # the tools above return dict[str, Any], but Transformers expects list[dict[Any, Any]]
 
-    chat_text = tokenizer.apply_chat_template(messages, tools=tools, **apply_chat_template_kwargs)
+    chat_text = tokenizer.apply_chat_template(messages, tools=tools, **apply_chat_template_kwargs)  # type: ignore[arg-type]
     chat_tokens = tokenizer.encode(chat_text, **encode_kwargs)
 
     slices: list[TokenSlice] = []
@@ -334,7 +335,13 @@ async def chats_to_tokens(
     for message in chat.all:
         # Find this message
         if not (
-            match := find_in_tokens(message.content, chat_tokens, tokenizer.decode, 0, search_start)
+            match := find_in_tokens(
+                message.content,
+                chat_tokens,
+                lambda tokens: tokenizer.decode(tokens),
+                0,
+                search_start,
+            )
         ):
             warnings.warn(
                 f"Warning: Could not find message '{message.content[:50]}...' in chat tokens",
@@ -370,7 +377,7 @@ async def chats_to_tokens(
             part_match = find_in_tokens(
                 part_text,
                 message_tokens,
-                tokenizer.decode,
+                lambda tokens: tokenizer.decode(tokens),
                 msg_start,
                 part_search_start,
             )
@@ -399,8 +406,9 @@ async def chats_to_tokens(
         # Continue searching after this message
         search_start = msg_end
 
+    # we ask for a string by default in apply_chat_template_kwargs with the tokenize=False
     return TokenizedChat(
-        text=chat_text,
+        text=chat_text,  # type: ignore[arg-type]
         tokens=chat_tokens,
         slices=slices,
         obj=chat,

docs/api/tokenize.mdx

@@ -12,7 +12,7 @@ get\_tokenizer
 ```python
 get_tokenizer(
     tokenizer_id: str, **tokenizer_kwargs: Any
-) -> AutoTokenizer | None
+) -> t.Any
 ```
 
 Get the tokenizer from transformers model identifier, or from an already loaded tokenizer.
@@ -30,15 +30,15 @@ Get the tokenizer from transformers model identifier, or from an already loaded
 
 **Returns:**
 
-* `AutoTokenizer | None`
+* `Any`
   –An instance of `AutoTokenizer`.
 
 <Accordion title="Source code in rigging/tokenize/tokenizer.py" icon="code">
 ```python
 def get_tokenizer(
     tokenizer_id: str,
     **tokenizer_kwargs: t.Any,
-) -> AutoTokenizer | None:
+) -> t.Any:
     """
     Get the tokenizer from transformers model identifier, or from an already loaded tokenizer.
 
@@ -49,18 +49,20 @@ def get_tokenizer(
     Returns:
         An instance of `AutoTokenizer`.
     """
-    tokenizer: AutoTokenizer | None = None
-
     try:
+        from transformers import AutoTokenizer
+
         tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_id,
             **tokenizer_kwargs,
         )
         logger.success(f"Loaded tokenizer for model '{tokenizer_id}'")
 
-    except Exception as e:  # noqa: BLE001
+    except Exception as e:
         # Catch all exceptions to handle any issues with loading the tokenizer
-        logger.error(f"Failed to load tokenizer for model '{tokenizer_id}': {e}")
+        raise RuntimeError(
+            f"Failed to load tokenizer for model '{tokenizer_id}': {e}",
+        ) from e
 
     return tokenizer
 ```

rigging/data.py

@@ -12,14 +12,16 @@
 import pandas as pd
 from elastic_transport import ObjectApiResponse
 from mypy_boto3_s3 import S3Client
-from transformers import AutoTokenizer
 
 from rigging.chat import Chat
 from rigging.error import TokenizeWarning
 from rigging.message import Message
 from rigging.tokenize import find_in_tokens
 from rigging.tokenize.base import TokenizedChat, TokenSlice
 
+if t.TYPE_CHECKING:
+    from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+
 
 def flatten_chats(chats: Chat | t.Sequence[Chat]) -> list[dict[t.Any, t.Any]]:
     """
@@ -294,8 +296,8 @@ async def chats_to_elastic(
 
 
 async def chats_to_tokens(
-    chat: Chat | None,
-    tokenizer: AutoTokenizer,
+    chat: Chat,
+    tokenizer: "PreTrainedTokenizerBase",
     *,
     apply_chat_template_kwargs: dict[str, t.Any] | None = None,
     encode_kwargs: dict[str, t.Any] | None = None,
@@ -331,8 +333,9 @@ async def chats_to_tokens(
         if chat.params and chat.params.tools
         else None
     )
+    # the tools above return dict[str, Any], but Transformers expects list[dict[Any, Any]]
 
-    chat_text = tokenizer.apply_chat_template(messages, tools=tools, **apply_chat_template_kwargs)
+    chat_text = tokenizer.apply_chat_template(messages, tools=tools, **apply_chat_template_kwargs)  # type: ignore[arg-type]
     chat_tokens = tokenizer.encode(chat_text, **encode_kwargs)
 
     slices: list[TokenSlice] = []
@@ -342,7 +345,13 @@ async def chats_to_tokens(
     for message in chat.all:
         # Find this message
         if not (
-            match := find_in_tokens(message.content, chat_tokens, tokenizer.decode, 0, search_start)
+            match := find_in_tokens(
+                message.content,
+                chat_tokens,
+                lambda tokens: tokenizer.decode(tokens),
+                0,
+                search_start,
+            )
         ):
             warnings.warn(
                 f"Warning: Could not find message '{message.content[:50]}...' in chat tokens",
@@ -378,7 +387,7 @@ async def chats_to_tokens(
             part_match = find_in_tokens(
                 part_text,
                 message_tokens,
-                tokenizer.decode,
+                lambda tokens: tokenizer.decode(tokens),
                 msg_start,
                 part_search_start,
             )
@@ -407,8 +416,9 @@ async def chats_to_tokens(
         # Continue searching after this message
         search_start = msg_end
 
+    # we ask for a string by default in apply_chat_template_kwargs with the tokenize=False
     return TokenizedChat(
-        text=chat_text,
+        text=chat_text,  # type: ignore[arg-type]
         tokens=chat_tokens,
         slices=slices,
         obj=chat,

rigging/tokenize/tokenizer.py

@@ -1,20 +1,14 @@
-import importlib.util
 import typing as t
 
-if importlib.util.find_spec("transformers") is None:
-    raise ModuleNotFoundError("Please install the `transformers` package to use this module.")
-
-
 from loguru import logger
-from transformers import AutoTokenizer
 
 from rigging.tokenize.base import Decoder
 
 
 def get_tokenizer(
     tokenizer_id: str,
     **tokenizer_kwargs: t.Any,
-) -> AutoTokenizer | None:
+) -> t.Any:
     """
     Get the tokenizer from transformers model identifier, or from an already loaded tokenizer.
 
@@ -25,18 +19,20 @@ def get_tokenizer(
     Returns:
         An instance of `AutoTokenizer`.
     """
-    tokenizer: AutoTokenizer | None = None
-
     try:
+        from transformers import AutoTokenizer
+
         tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_id,
             **tokenizer_kwargs,
         )
         logger.success(f"Loaded tokenizer for model '{tokenizer_id}'")
 
-    except Exception as e:  # noqa: BLE001
+    except Exception as e:
         # Catch all exceptions to handle any issues with loading the tokenizer
-        logger.error(f"Failed to load tokenizer for model '{tokenizer_id}': {e}")
+        raise RuntimeError(
+            f"Failed to load tokenizer for model '{tokenizer_id}': {e}",
+        ) from e
 
     return tokenizer