Add special_tokens to Tokenizer interface

dottxt-ai · Sep 16, 2023 · 76e95fe · 76e95fe
1 parent ff4ebb3
commit 76e95fe
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 1 deletion.
diff --git a/outlines/models/tokenizer.py b/outlines/models/tokenizer.py
@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Dict, List, Protocol, Tuple, Union
+from typing import Dict, List, Protocol, Set, Tuple, Union
 
 import numpy as np
 from numpy.typing import NDArray
@@ -10,6 +10,7 @@ class Tokenizer(Protocol):
     eos_token_id: int
     pad_token_id: int
     vocabulary: Dict[str, int]
+    special_tokens: Set[int]
 
     @abstractmethod
     def encode(

diff --git a/outlines/models/transformers.py b/outlines/models/transformers.py
@@ -108,6 +108,8 @@ def __init__(self, model_name: str, **kwargs):
             self.pad_token_id = self.tokenizer.pad_token_id
             self.pad_token = self.tokenizer.pad_token
 
+        self.special_tokens = set(self.tokenizer.special_tokens_map.values())
+
         self.vocabulary = self.tokenizer.get_vocab()
         self.is_llama = isinstance(self.tokenizer, get_llama_tokenizer_types())
 

diff --git a/tests/text/generate/test_regex.py b/tests/text/generate/test_regex.py
@@ -13,6 +13,7 @@ class Tokenizer:
     pad_token_id = -1
     vocabulary = {"<EOS>": 0, "-": 1, "1": 2, "0.": 3, "431": 4, "a": 5, "A": 6}
     tokens = list(vocabulary.keys())
+    special_tokens = {"<EOS>"}
 
     def decode(self, token_ids):
         decoded = []