Upload adapt_tokenizer.py

llava/adapt_tokenizer.py

@@ -0,0 +1,41 @@
+from typing import Union
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
+Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
+NUM_SENTINEL_TOKENS: int = 100
+
+def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
+    """Adds sentinel tokens and padding token (if missing).
+
+    Expands the tokenizer vocabulary to include sentinel tokens
+    used in mixture-of-denoiser tasks as well as a padding token.
+
+    All added tokens are added as special tokens. No tokens are
+    added if sentinel tokens and padding token already exist.
+    """
+    sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]
+    tokenizer.add_tokens(sentinels_to_add, special_tokens=True)
+    if tokenizer.pad_token is None:
+        tokenizer.add_tokens('<pad>', special_tokens=True)
+        tokenizer.pad_token = '<pad>'
+        assert tokenizer.pad_token_id is not None
+    sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])
+    _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids
+    tokenizer.sentinel_token_ids = _sentinel_token_ids
+
+class AutoTokenizerForMOD(AutoTokenizer):
+    """AutoTokenizer + Adaptation for MOD.
+
+    A simple wrapper around AutoTokenizer to make instantiating
+    an MOD-adapted tokenizer a bit easier.
+
+    MOD-adapted tokenizers have sentinel tokens (e.g., <extra_id_0>),
+    a padding token, and a property to get the token ids of the
+    sentinel tokens.
+    """
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        """See `AutoTokenizer.from_pretrained` docstring."""
+        tokenizer = super().from_pretrained(*args, **kwargs)
+        adapt_tokenizer_for_denoising(tokenizer)
+        return tokenizer