Fix tokenizer type

ProbablyFaiz · Jun 4, 2024 · 80b8e0b · 80b8e0b
1 parent f907366
commit 80b8e0b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rl/llm/token_utils.py b/rl/llm/token_utils.py
@@ -1,13 +1,13 @@
 import more_itertools as mit
-from transformers import PreTrainedTokenizer
+from transformers import PreTrainedTokenizerFast
 
 
 def token_chunk(
     text: str,
     *,
     chunk_size: int,
     stride: int | None = None,
-    tokenizer: PreTrainedTokenizer,
+    tokenizer: PreTrainedTokenizerFast,
 ) -> list[str]:
     """Tokenize a long text into chunks of a specified number of tokens."""
     if stride is None: