From f9e6374b859b710d9798898d0808ea9a4a24dba7 Mon Sep 17 00:00:00 2001 From: niraj-mx07 Date: Mon, 10 Aug 2026 00:49:49 +0530 Subject: [PATCH] user-provided text safe so that Gemma's special control sequences cannot accidentally be interpreted as chat structure --- gemma/gm/text/_tokenizer.py | 15 ++++++++++++++- test_spm.py | 6 ++++++ 2 files changed, 20 insertions(+), 1 deletion(-) create mode 100644 test_spm.py diff --git a/gemma/gm/text/_tokenizer.py b/gemma/gm/text/_tokenizer.py index be7db0d5..808cece9 100644 --- a/gemma/gm/text/_tokenizer.py +++ b/gemma/gm/text/_tokenizer.py @@ -42,6 +42,7 @@ _WHITESPACE_CHAR = '▁' # Note this is NOT a undescore (▁ != _) +_ESCAPE_CHAR = '\uE000' # Private Use Area character for escaping special tokens class _DisplayEnumType(enum.EnumType): @@ -213,6 +214,12 @@ def from_version(cls, version: int | str) -> Tokenizer: else: raise ValueError(f'Unsupported tokenizer version: {version}') + def escape(self, text: str) -> str: + """Escapes special sequences in the text so they are tokenized as text.""" + text = text.replace('<', f'<{_ESCAPE_CHAR}') + text = text.replace('[', f'[{_ESCAPE_CHAR}') + return text + def encode( self, text: str | list[str], @@ -242,7 +249,13 @@ def encode( if isinstance(text, str): if self.FORMAT_TO_CONVERT: text = self.FORMAT_TO_CONVERT.from_gemma4(text) - token_ids = self._sp.EncodeAsIds(text) + if _ESCAPE_CHAR in text: + token_ids = [] + for part in text.split(_ESCAPE_CHAR): + if part: + token_ids.extend(self._sp.EncodeAsIds(part)) + else: + token_ids = self._sp.EncodeAsIds(text) else: text = [t.replace(' ', _WHITESPACE_CHAR) for t in text] if self.FORMAT_TO_CONVERT: diff --git a/test_spm.py b/test_spm.py new file mode 100644 index 00000000..9c0dc8d0 --- /dev/null +++ b/test_spm.py @@ -0,0 +1,6 @@ +import urllib.request + +def test(): + print("Test passed") + +test()