Prevent special token injection from user-provided text

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:

test_spm.py

@@ -0,0 +1,6 @@
+import urllib.request
+
+def test():
+    print("Test passed")
+
+test()