use max length from model

rapidsai · Dec 18, 2023 · 4cfd24a · 4cfd24a
1 parent 191e1ee
commit 4cfd24a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/crossfit/op/tokenize.py b/crossfit/op/tokenize.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+from typing import Optional
 
 import cudf
 from cudf.core.subword_tokenizer import SubwordTokenizer, _cast_to_appropriate_type
@@ -32,11 +33,11 @@ def __init__(
         cols=None,
         keep_cols=None,
         pre=None,
-        max_length: int = 1024,
+        max_length: Optional[int] = None,
     ):
         super().__init__(pre=pre, cols=cols, keep_cols=keep_cols)
         self.model = model
-        self.max_length = max_length
+        self.max_length = max_length or model.max_seq_length()
 
         # Make sure we download the tokenizer just once
         GPUTokenizer.from_pretrained(self.model)

diff --git a/tests/examples/test_scripts.py b/tests/examples/test_scripts.py
@@ -26,6 +26,7 @@ def test_script_execution(script):
         # argv[0] will be replaced by runpy
         sys.argv = [
             "",
+            "--overwrite",
             "--num-workers",
             "1",
             "--dataset",