Add Support to Gemma2

requirements.txt

@@ -1,9 +1,10 @@
 torch>=2.0.0
 numpy>=1.24.2
-transformers>=4.37.2
+transformers>=4.42.0
 tokenizers>=0.15.2
 termcolor>=2.4.0
 sentencepiece>=0.1.99
 protobuf>=4.25.2
 setuptools>=69.0.3
-line_profiler
+datasets
+line_profiler

tests/test_accept_token_sequence/test_gemma2.py

@@ -0,0 +1,16 @@
+import unittest
+
+from transformers import AutoTokenizer
+
+from tests.test_accept_token_sequence._test_accept_tokens_mixin import (
+    TokenizerTesterMixin,
+)
+
+
+class Gemma2TokenizerTest(TokenizerTesterMixin, unittest.TestCase):
+
+    tokenizer_class = AutoTokenizer
+    pretrained_name = "google/gemma-2-2b-it"
+
+    def setUp(self):
+        super().setUp()

tests/test_metrics/test_metric.py

@@ -1,8 +1,11 @@
 import math
 from unittest import TestCase
+
 import torch
-from transformers_cfg.metrics import ConstrainedDecodingMetric
-from transformers_cfg.metrics.metrics import ConstrainedDecodingMetricOutput
+from transformers_cfg.metrics.metrics import (
+    ConstrainedDecodingMetric,
+    ConstrainedDecodingMetricOutput,
+)
 
 
 class TestConstrainedDecodingMetric(TestCase):

transformers_cfg/tokenization/SUPPORTED_TOKENIZERS.py

@@ -5,6 +5,7 @@
     T5TokenizerFast,
     CodeGenTokenizerFast,
     PreTrainedTokenizerFast,
+    GemmaTokenizerFast
 )
 
 SUPPORTED_TOKENIZERS = {
@@ -14,4 +15,5 @@
     T5TokenizerFast,
     CodeGenTokenizerFast,
     PreTrainedTokenizerFast,
+    GemmaTokenizerFast
 }

transformers_cfg/tokenization/middle/TokenizerMiddleMapping.py

@@ -8,6 +8,7 @@
     CodeGenTokenizerFast,
     LlamaTokenizerFast,
     PreTrainedTokenizerFast,
+    GemmaTokenizerFast
 )
 
 from transformers_cfg.tokenization.utils import get_tokenizer_charset
@@ -38,7 +39,7 @@ def from_hf_tokenizer(cls, hf_tokenizer):
             hf_tokenizer, (GPT2TokenizerFast, BartTokenizerFast, CodeGenTokenizerFast)
         ):
             return GPT2TokenizerMiddleMapping(hf_tokenizer)
-        elif isinstance(hf_tokenizer, LlamaTokenizerFast):
+        elif isinstance(hf_tokenizer, (LlamaTokenizerFast, GemmaTokenizerFast)):
             # deepseek, though inheriting from LlamaTokenizerFast, is actually a GPT2TokenizerFast
             # check https://github.com/epfl-dlab/transformers-CFG/issues/72
             if hf_tokenizer.name_or_path.startswith("deepseek-ai/deepseek-coder"):

transformers_cfg/tokenization/tokenizer.py

@@ -7,6 +7,7 @@
     T5TokenizerFast,
     CodeGenTokenizerFast,
     PreTrainedTokenizerFast,
+    GemmaTokenizerFast
 )
 
 from transformers_cfg.tokenization.SUPPORTED_TOKENIZERS import SUPPORTED_TOKENIZERS
@@ -53,7 +54,7 @@ def from_hf_tokenizer(cls, hf_tokenizer):
             (GPT2TokenizerFast, BartTokenizerFast),
         ):
             return TCFG_GPT2Tokenizer(hf_tokenizer)
-        elif isinstance(hf_tokenizer, (LlamaTokenizerFast, T5TokenizerFast)):
+        elif isinstance(hf_tokenizer, (LlamaTokenizerFast, GemmaTokenizerFast, T5TokenizerFast)):
             return TCFG_LlamaTokenizer(hf_tokenizer)
         elif isinstance(hf_tokenizer, CodeGenTokenizerFast):
             # phi reuses the codegen tokenizer