add missing is_train kwarg in tokenize_id

eole-nlp · Jan 15, 2025 · 928c6fe · 928c6fe
1 parent 7947a6b
commit 928c6fe
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/eole/transforms/tokenize_id.py b/eole/transforms/tokenize_id.py
@@ -96,10 +96,10 @@ def tokenize_string(self, string, side="src", is_train=False):
         return tokens
 
     def apply(self, example, is_train=False, stats=None, **kwargs):
-        src_tokens = self.tokenize_string(" ".join(example["src"]), side="src")
+        src_tokens = self.tokenize_string(" ".join(example["src"]), side="src", is_train=is_train)
         example["src_ids"] = src_tokens
         if example.get("tgt", None) is not None:
-            tgt_tokens = self.tokenize_string(" ".join(example["tgt"]), side="tgt")
+            tgt_tokens = self.tokenize_string(" ".join(example["tgt"]), side="tgt", is_train=is_train)
             example["tgt_ids"] = tgt_tokens
         return example