From b91d67a7233c7324f3dd374a1557e78847c1e5b6 Mon Sep 17 00:00:00 2001 From: Taksh Date: Fri, 17 Jul 2026 10:35:53 +0300 Subject: [PATCH] fix: whitespace tokenizer counts tokens with findall split() overcounts (N+1); match RedPajama tagger / token regex semantics. Co-authored-by: Cursor --- python/dolma/taggers/length.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/dolma/taggers/length.py b/python/dolma/taggers/length.py index a1420c7b..73b7745d 100644 --- a/python/dolma/taggers/length.py +++ b/python/dolma/taggers/length.py @@ -60,7 +60,7 @@ class WhitespaceLengthV1(BaseTagger): WHITESPACE_REGEX = regex.compile(r"\w+|[^\w\s]+") def predict(self, doc: Document) -> DocResult: - score = len(self.WHITESPACE_REGEX.split(doc.text)) + score = len(self.WHITESPACE_REGEX.findall(doc.text)) return DocResult(doc=doc, spans=[Span(start=0, end=len(doc.text), type="length", score=score)]) @@ -68,7 +68,7 @@ def predict(self, doc: Document) -> DocResult: class WhitespaceLengthParagraphsV1(WhitespaceLengthV1): def predict(self, doc: Document) -> DocResult: spans = [ - Span(start=p.start, end=p.end, type="paragraph", score=len(self.WHITESPACE_REGEX.split(p.text))) + Span(start=p.start, end=p.end, type="paragraph", score=len(self.WHITESPACE_REGEX.findall(p.text))) for p in split_paragraphs(doc.text) ] spans.append(Span(start=0, end=len(doc.text), type="document", score=sum(s.score for s in spans)))