Replace GPT2TokenizerFast with tiktoken-based tokenizer (#44)

autopr/repos/completions_repo.py

@@ -25,7 +25,7 @@ def __init__(
         self.context_limit = context_limit
         self.temperature = temperature
 
-        self.tokenizer = tokenizer.get_tokenizer(max_tokens)
+        self.tokenizer = tokenizer.get_tiktoken_tokenizer()
         self.log = structlog.get_logger(repo=self.__class__.__name__)
 
     def complete(

autopr/utils/repo.py

@@ -121,21 +121,18 @@ def repo_to_file_descriptors(repo: Repo, context_window: int, file_chunk_size: i
             log.debug(f"Error decoding file: {blob.path}")
             continue
 
-        tokenizer = get_tokenizer(context_window)
+        tokenizer = get_tiktoken_tokenizer(context_window)
 
-        tokens = tokenizer.encode(content)
-        # Split into chunks up to the last newline
+        tokens = tokenizer(content)
         chunks: list[list[tuple[int, str]]] = []
         line_buffer = []
         for i, line in enumerate(content.splitlines()):
             line_buffer.append((i, line))
-            # FIXME speed this up
-            token_length = len(tokenizer.encode(
-                '\n'.join([l[1] for l in line_buffer])
-            ))
+            token_length = len(tokenizer('\n'.join([l[1] for l in line_buffer])))
             if token_length >= file_chunk_size:
                 chunks.append(line_buffer)
                 line_buffer = []
+                line_buffer = []
         if line_buffer:
             chunks.append(line_buffer)
 

autopr/utils/tokenizer.py

@@ -1,11 +1,14 @@
-import transformers
+import tiktoken
+from tiktoken import Tokenizer
 
-_tokenizer_cache: dict[int, transformers.GPT2TokenizerFast] = {}
+_tokenizer_cache: dict[int, Tokenizer] = {}
 
-
-def get_tokenizer(model_max_length: int):
+def get_tiktoken_tokenizer(model_max_length: int):
     global _tokenizer_cache
 
     if model_max_length not in _tokenizer_cache:
-        _tokenizer_cache[model_max_length] = transformers.GPT2TokenizerFast.from_pretrained('gpt2', model_max_length=model_max_length)
+        _tokenizer_cache[model_max_length] = Tokenizer()
     return _tokenizer_cache[model_max_length]
+
+def get_tokenizer(model_max_length: int):
+    return get_tiktoken_tokenizer(model_max_length)

tests/test_tokenizer.py

@@ -0,0 +1,27 @@
+import pytest
+from autopr.utils.tokenizer import get_tiktoken_tokenizer
+
+def test_tiktoken_tokenizer_with_simple_strings():
+    tokenizer = get_tiktoken_tokenizer()
+
+    text = "This is a simple text."
+    expected_token_count = 6
+
+    token_count = tokenizer(text)
+    assert token_count == expected_token_count, f"Expected {expected_token_count} tokens, got {token_count}"
+
+def test_tiktoken_tokenizer_with_chat_messages():
+    tokenizer = get_tiktoken_tokenizer()
+
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "name": "Alice", "content": "What's the capital of France?"},
+        {"role": "assistant", "content": "The capital of France is Paris."},
+    ]
+    expected_token_count = 38
+
+    token_count = tokenizer(messages)
+    assert token_count == expected_token_count, f"Expected {expected_token_count} tokens, got {token_count}"
+
+if __name__ == "__main__":
+    pytest.main()