diff --git a/python/dolma/taggers/code/code_taggers.py b/python/dolma/taggers/code/code_taggers.py
index 31b57087..62b7ecb2 100644
--- a/python/dolma/taggers/code/code_taggers.py
+++ b/python/dolma/taggers/code/code_taggers.py
@@ -156,7 +156,7 @@ def __init__(self):
         super().__init__()
 
     def _get_num_tokens(self, text: str) -> int:
-        return len(self.whitespace_regex.split(text))
+        return len(self.whitespace_regex.findall(text))
 
     def predict(self, doc: Document) -> DocResult:
         """Main runner."""