"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sequence token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the end of sequence.
- The token used is the `sep_token`.
-
-
-
- prefix_token (`str`, *optional*, defaults to `"▁"`):
- Prefix token used for infilling.
- middle_token (`str`, *optional*, defaults to `"▁"`):
- Middle token used for infilling.
- suffix_token (`str`, *optional*, defaults to `"▁"`):
- Suffix token used for infilling.
- eot_token (`str`, *optional*, defaults to `"▁"`):
- End of text token used for infilling.
- fill_token (`str`, *optional*, defaults to `""`):
- The token used to split the input between the prefix and suffix.
- suffix_first (`bool`, *optional*, defaults to `False`):
- Whether the input prompt and suffix should be formatted with the suffix first.
- sp_model_kwargs (`dict`, *optional*):
- Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
- SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
- to set:
-
- - `enable_sampling`: Enable subword regularization.
- - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
-
- - `nbest_size = {0,1}`: No sampling is performed.
- - `nbest_size > 1`: samples from the nbest_size results.
- - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
- using forward-filtering-and-backward-sampling algorithm.
-
- - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
- BPE-dropout.
- add_bos_token (`bool`, *optional*, defaults to `True`):
- Whether to add a beginning of sequence token at the start of sequences.
- add_eos_token (`bool`, *optional*, defaults to `False`):
- Whether to add an end of sequence token at the end of sequences.
- clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
- Whether or not to clean up the tokenization spaces.
- additional_special_tokens (`List[str]`, *optional*):
- Additional special tokens used by the tokenizer.
- use_default_system_prompt (`bool`, *optional*, defaults to `False`):
- Whether or not the default system prompt for Llama should be used.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- unk_token="",
- bos_token="",
- eos_token="",
- prefix_token="▁",
- middle_token="▁",
- suffix_token="▁",
- eot_token="▁",
- fill_token="",
- suffix_first=False,
- sp_model_kwargs: Optional[Dict[str, Any]] = None,
- add_bos_token=True,
- add_eos_token=False,
- clean_up_tokenization_spaces=False,
- additional_special_tokens=None,
- use_default_system_prompt=False,
- **kwargs,
- ):
- requires_backends(self, "protobuf")
- self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
- bos_token = AddedToken(bos_token, normalized=False, special=True) if isinstance(bos_token, str) else bos_token
- eos_token = AddedToken(eos_token, normalized=False, special=True) if isinstance(eos_token, str) else eos_token
- unk_token = AddedToken(unk_token, normalized=False, special=True) if isinstance(unk_token, str) else unk_token
-
- self.use_default_system_prompt = use_default_system_prompt
- # mark tokens special to skip them
- additional_special_tokens = additional_special_tokens or []
- for token in [prefix_token, middle_token, suffix_token, eot_token]:
- additional_special_tokens += [token] if token is not None else []
-
- self.vocab_file = vocab_file
- self.add_bos_token = add_bos_token
- self.add_eos_token = add_eos_token
- self._prefix_token = prefix_token
- self._middle_token = middle_token
- self._suffix_token = suffix_token
- self._eot_token = eot_token
- self.fill_token = fill_token
- self.suffix_first = suffix_first
- self.sp_model = self.get_spm_processor()
-
- super().__init__(
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- add_bos_token=add_bos_token,
- add_eos_token=add_eos_token,
- prefix_token=prefix_token,
- middle_token=middle_token,
- suffix_token=suffix_token,
- eot_token=eot_token,
- fill_token=fill_token,
- sp_model_kwargs=self.sp_model_kwargs,
- suffix_first=suffix_first,
- clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- additional_special_tokens=additional_special_tokens,
- use_default_system_prompt=use_default_system_prompt,
- **kwargs,
- )
-
- @property
- def unk_token_length(self):
- return len(self.sp_model.encode(str(self.unk_token)))
-
- def get_spm_processor(self):
- tokenizer = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- with open(self.vocab_file, "rb") as f:
- sp_model = f.read()
- model_pb2 = import_protobuf()
- model = model_pb2.ModelProto.FromString(sp_model)
- normalizer_spec = model_pb2.NormalizerSpec()
- normalizer_spec.add_dummy_prefix = False
- model.normalizer_spec.MergeFrom(normalizer_spec)
- sp_model = model.SerializeToString()
- tokenizer.LoadFromSerializedProto(sp_model)
- return tokenizer
-
- @property
- def prefix_token(self):
- return self._prefix_token
-
- @property
- def prefix_id(self):
- if self._prefix_token is None:
- return None
- return self.convert_tokens_to_ids(self.prefix_token)
-
- @property
- def middle_token(self):
- return self._middle_token
-
- @property
- def middle_id(self):
- if self._middle_token is None:
- return None
- return self.convert_tokens_to_ids(self.middle_token)
-
- @property
- def suffix_token(self):
- return self._suffix_token
-
- @property
- def suffix_id(self):
- if self._suffix_token is None:
- return None
- return self.convert_tokens_to_ids(self.suffix_token)
-
- @property
- def eot_token(self):
- return self._eot_token
-
- @property
- def eot_id(self):
- if self._eot_token is None:
- return None
- return self.convert_tokens_to_ids(self.eot_token)
-
- @property
- def vocab_size(self):
- """Returns vocab size"""
- return self.sp_model.get_piece_size()
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.get_vocab
- def get_vocab(self):
- """Returns vocab as a dict"""
- vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
- vocab.update(self.added_tokens_encoder)
- return vocab
-
- def tokenize(self, prefix, suffix=None, suffix_first=False, **kwargs) -> List[int]:
- # add a prefix space to `prefix`
- if self.fill_token is not None and self.fill_token in prefix and suffix is None:
- prefix, suffix = prefix.split(self.fill_token)
-
- if len(prefix) > 0:
- prefix = SPIECE_UNDERLINE + prefix.replace(SPIECE_UNDERLINE, " ")
-
- if suffix is None or len(suffix) < 1:
- tokens = super().tokenize(prefix, **kwargs)
- if len(tokens) > 1 and tokens[0] == SPIECE_UNDERLINE and tokens[1] in self.all_special_tokens:
- tokens = tokens[1:]
- return tokens
-
- prefix_tokens = self._tokenize(prefix) # prefix has an extra `SPIECE_UNDERLINE`
-
- if None in (self.prefix_id, self.middle_id, self.suffix_id):
- raise ValueError(
- "The input either includes a `prefix` and a `suffix` used for the infilling task,"
- f" or can be split on the {self.fill_token} token, creating a suffix and prefix,"
- " but the model does not support `infilling`."
- )
- suffix_tokens = self._tokenize(suffix) # make sure CodeLlama sp model does not mess up
-
- suffix_first = suffix_first if suffix_first is not None else self.suffix_first
- if suffix_first:
- # format as " {suf} {pre}"
- return [self.prefix_token, self.suffix_token] + suffix_tokens + [self.middle_token] + prefix_tokens
- else:
- # format as " {pre} {suf} "
- return [self.prefix_token] + prefix_tokens + [self.suffix_token] + suffix_tokens + [self.middle_token]
-
- def _tokenize(self, text, **kwargs):
- """
- Returns a tokenized string.
-
- We de-activated the `add_dummy_prefix` option, thus the sentencepiece internals will always strip any
- SPIECE_UNDERLINE. For example: `self.sp_model.encode(f"{SPIECE_UNDERLINE}Hey", out_type = str)` will give
- `['H', 'e', 'y']` instead of `['▁He', 'y']`. Thus we always encode `f"{unk_token}text"` and strip the
- `unk_token`. Here is an example with `unk_token = ""` and `unk_token_length = 4`.
- `self.tokenizer.sp_model.encode(" Hey", out_type = str)[4:]`.
- """
- tokens = self.sp_model.encode(text, out_type=str)
- if not text.startswith((SPIECE_UNDERLINE, " ")):
- return tokens
- # 1. Encode string + prefix ex: " Hey"
- tokens = self.sp_model.encode(self.unk_token + text, out_type=str)
- # 2. Remove self.unk_token from ['<','unk','>', '▁Hey']
- return tokens[self.unk_token_length :] if len(tokens) >= self.unk_token_length else tokens
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.sp_model.piece_to_id(token)
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- token = self.sp_model.IdToPiece(index)
- return token
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- # since we manually add the prefix space, we have to remove it when decoding
- if tokens[0].startswith(SPIECE_UNDERLINE):
- tokens[0] = tokens[0][1:]
-
- current_sub_tokens = []
- out_string = ""
- for _, token in enumerate(tokens):
- # make sure that special tokens are not decoded using sentencepiece model
- if token in self.all_special_tokens:
- out_string += self.sp_model.decode(current_sub_tokens) + token
- current_sub_tokens = []
- else:
- current_sub_tokens.append(token)
- out_string += self.sp_model.decode(current_sub_tokens)
- return out_string
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
- """
- Save the vocabulary and special tokens file to a directory.
-
- Args:
- save_directory (`str`):
- The directory in which to save the vocabulary.
-
- Returns:
- `Tuple(str)`: Paths to the files saved.
- """
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
- elif not os.path.isfile(self.vocab_file):
- with open(out_vocab_file, "wb") as fi:
- content_spiece_model = self.sp_model.serialized_model_proto()
- fi.write(content_spiece_model)
-
- return (out_vocab_file,)
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- bos_token_id = [self.bos_token_id] if self.add_bos_token else []
- eos_token_id = [self.eos_token_id] if self.add_eos_token else []
-
- output = bos_token_id + token_ids_0 + eos_token_id
-
- if token_ids_1 is not None:
- output = output + bos_token_id + token_ids_1 + eos_token_id
-
- return output
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.get_special_tokens_mask
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- bos_token_id = [1] if self.add_bos_token else []
- eos_token_id = [1] if self.add_eos_token else []
-
- if token_ids_1 is None:
- return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
- return (
- bos_token_id
- + ([0] * len(token_ids_0))
- + eos_token_id
- + bos_token_id
- + ([0] * len(token_ids_1))
- + eos_token_id
- )
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- if token_ids_1 is None, only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of ids.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- bos_token_id = [self.bos_token_id] if self.add_bos_token else []
- eos_token_id = [self.eos_token_id] if self.add_eos_token else []
-
- output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
-
- if token_ids_1 is not None:
- output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
-
- return output
-
- @property
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.default_chat_template
- def default_chat_template(self):
- """
- LLaMA uses [INST] and [/INST] to indicate user messages, and <> and <> to indicate system messages.
- Assistant messages do not have special tokens, because LLaMA chat models are generally trained with strict
- user/assistant/user/assistant message ordering, and so assistant messages can be identified from the ordering
- rather than needing special tokens. The system message is partly 'embedded' in the first user message, which
- results in an unusual token ordering when it is present. This template should definitely be changed if you wish
- to fine-tune a model with more flexible role ordering!
-
- The output should look something like:
-
- [INST] B_SYS SystemPrompt E_SYS Prompt [/INST] Answer [INST] Prompt [/INST] Answer
- [INST] Prompt [/INST]
-
- The reference for this chat template is [this code
- snippet](https://github.com/facebookresearch/llama/blob/556949fdfb72da27c2f4a40b7f0e4cf0b8153a28/llama/generation.py#L320-L362)
- in the original repository.
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- template = (
- "{% if messages[0]['role'] == 'system' %}"
- "{% set loop_messages = messages[1:] %}" # Extract system message if it's present
- "{% set system_message = messages[0]['content'] %}"
- "{% elif USE_DEFAULT_PROMPT == true and not '<>' in messages[0]['content'] %}"
- "{% set loop_messages = messages %}" # Or use the default system message if the flag is set
- "{% set system_message = 'DEFAULT_SYSTEM_MESSAGE' %}"
- "{% else %}"
- "{% set loop_messages = messages %}"
- "{% set system_message = false %}"
- "{% endif %}"
- "{% for message in loop_messages %}" # Loop over all non-system messages
- "{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}"
- "{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}"
- "{% endif %}"
- "{% if loop.index0 == 0 and system_message != false %}" # Embed system message in first message
- "{% set content = '<>\\n' + system_message + '\\n<>\\n\\n' + message['content'] %}"
- "{% else %}"
- "{% set content = message['content'] %}"
- "{% endif %}"
- "{% if message['role'] == 'user' %}" # After all of that, handle messages/roles in a fairly normal way
- "{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}"
- "{% elif message['role'] == 'system' %}"
- "{{ '<>\\n' + content.strip() + '\\n<>\\n\\n' }}"
- "{% elif message['role'] == 'assistant' %}"
- "{{ ' ' + content.strip() + ' ' + eos_token }}"
- "{% endif %}"
- "{% endfor %}"
- )
- template = template.replace("USE_DEFAULT_PROMPT", "true" if self.use_default_system_prompt else "false")
- default_message = DEFAULT_SYSTEM_PROMPT.replace("\n", "\\n").replace("'", "\\'")
- template = template.replace("DEFAULT_SYSTEM_MESSAGE", default_message)
-
- return template
-
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sp_model"] = None
- state["sp_model_proto"] = self.sp_model.serialized_model_proto()
- return state
-
- def __setstate__(self, d):
- self.__dict__ = d
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.LoadFromSerializedProto(self.sp_model_proto)
diff --git a/transformers/models/code_llama/tokenization_code_llama_fast.py b/transformers/models/code_llama/tokenization_code_llama_fast.py
deleted file mode 100644
index e2429aaec5d187963cf64785c91f8ce2a5cd5c56..0000000000000000000000000000000000000000
--- a/transformers/models/code_llama/tokenization_code_llama_fast.py
+++ /dev/null
@@ -1,439 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import os
-from shutil import copyfile
-from typing import List, Optional, Tuple
-
-from tokenizers import normalizers, processors
-
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import is_sentencepiece_available, logging
-from ...utils.versions import require_version
-
-
-require_version("tokenizers>=0.13.3")
-
-if is_sentencepiece_available():
- from .tokenization_code_llama import CodeLlamaTokenizer
-else:
- CodeLlamaTokenizer = None
-
-logger = logging.get_logger(__name__)
-VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model", "tokenizer_file": "tokenizer.json"}
-
-SPIECE_UNDERLINE = "▁"
-
-
-B_INST, E_INST = "[INST]", "[/INST]"
-B_SYS, E_SYS = "<>\n", "\n<>\n\n"
-
-# fmt: off
-DEFAULT_SYSTEM_PROMPT = """You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your \
-answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure\
- that your responses are socially unbiased and positive in nature.
-
-If a question does not make any sense, or is not factually coherent, explain why instead of answering something not \
-correct. If you don't know the answer to a question, please don't share false information."""
-# fmt: on
-
-
-class CodeLlamaTokenizerFast(PreTrainedTokenizerFast):
- """
- Construct a Llama tokenizer. Based on byte-level Byte-Pair-Encoding.
-
- This uses notably ByteFallback and no normalization.
-
- ```python
- >>> from transformers import CodeLlamaTokenizerFast
-
- >>> tokenizer = CodeLlamaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer")
- >>> tokenizer.encode("Hello this is a test")
- [1, 15043, 445, 338, 263, 1243]
- ```
-
- If you want to change the `bos_token` or the `eos_token`, make sure to specify them when initializing the model, or
- call `tokenizer.update_post_processor()` to make sure that the post-processing is correctly done (otherwise the
- values of the first token and final token of an encoded sequence will not be correct). For more details, checkout
- [post-processors] (https://huggingface.co/docs/tokenizers/api/post-processors) documentation.
-
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods. The default configuration match that of
- [codellama/CodeLlama-7b-Instruct-hf](https://huggingface.co/codellama/CodeLlama-7b-Instruct-hf/blob/main/tokenizer_config.json)
- which supports prompt infilling.
-
- Args:
- vocab_file (`str`, *optional*):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .model extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- tokenizer_file (`str`, *optional*):
- [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that
- contains everything needed to load the tokenizer.
- clean_up_tokenization_spaces (`str`, *optional*, defaults to `False`):
- Wether to cleanup spaces after decoding, cleanup consists in removing potential artifacts like extra
- spaces.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sequence token.
- prefix_token (`str`, *optional*, defaults to `"▁"`):
- Prefix token used for infilling.
- middle_token (`str`, *optional*, defaults to `"▁"`):
- Middle token used for infilling.
- suffix_token (`str`, *optional*, defaults to `"▁"`):
- Suffix token used for infilling.
- eot_token (`str`, *optional*, defaults to `"▁"`):
- End of text token used for infilling.
- fill_token (`str`, *optional*, defaults to `""`):
- The token used to split the input between the prefix and suffix.
- additional_special_tokens (`List[str]`, *optional*):
- Additional special tokens used by the tokenizer.
- add_bos_token (`bool`, *optional*, defaults to `True`):
- Whether to add a beginning of sequence token at the start of sequences.
- add_eos_token (`bool`, *optional*, defaults to `False`):
- Whether to add an end of sequence token at the end of sequences.
- use_default_system_prompt (`bool`, *optional*, defaults to `False`):
- Whether or not the default system prompt for Llama should be used.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = CodeLlamaTokenizer
- padding_side = "left"
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- clean_up_tokenization_spaces=False,
- unk_token="",
- bos_token="",
- eos_token="",
- prefix_token="▁",
- middle_token="▁",
- suffix_token="▁",
- eot_token="▁",
- fill_token="",
- additional_special_tokens=None,
- add_bos_token=True,
- add_eos_token=False,
- use_default_system_prompt=False,
- **kwargs,
- ):
- # mark tokens special to skip them
- additional_special_tokens = additional_special_tokens or []
- for token in [prefix_token, middle_token, suffix_token, eot_token]:
- additional_special_tokens += [token] if token is not None else []
- self.use_default_system_prompt = use_default_system_prompt
-
- super().__init__(
- vocab_file=vocab_file,
- tokenizer_file=tokenizer_file,
- clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- additional_special_tokens=additional_special_tokens,
- unk_token=unk_token,
- bos_token=bos_token,
- eos_token=eos_token,
- add_bos_token=add_bos_token,
- add_eos_token=add_eos_token,
- prefix_token=prefix_token,
- middle_token=middle_token,
- suffix_token=suffix_token,
- eot_token=eot_token,
- fill_token=fill_token,
- use_default_system_prompt=use_default_system_prompt,
- **kwargs,
- )
- self._add_bos_token = add_bos_token
- self._add_eos_token = add_eos_token
- self.update_post_processor()
-
- self.vocab_file = vocab_file
-
- self._prefix_token = prefix_token
- self._middle_token = middle_token
- self._suffix_token = suffix_token
- self._eot_token = eot_token
- self.fill_token = fill_token
-
- @property
- def can_save_slow_tokenizer(self) -> bool:
- return os.path.isfile(self.vocab_file) if self.vocab_file else False
-
- # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.update_post_processor
- def update_post_processor(self):
- """
- Updates the underlying post processor with the current `bos_token` and `eos_token`.
- """
- bos = self.bos_token
- bos_token_id = self.bos_token_id
- if bos is None and self.add_bos_token:
- raise ValueError("add_bos_token = True but bos_token = None")
-
- eos = self.eos_token
- eos_token_id = self.eos_token_id
- if eos is None and self.add_eos_token:
- raise ValueError("add_eos_token = True but eos_token = None")
-
- single = f"{(bos+':0 ') if self.add_bos_token else ''}$A:0{(' '+eos+':0') if self.add_eos_token else ''}"
- pair = f"{single}{(' '+bos+':1') if self.add_bos_token else ''} $B:1{(' '+eos+':1') if self.add_eos_token else ''}"
-
- special_tokens = []
- if self.add_bos_token:
- special_tokens.append((bos, bos_token_id))
- if self.add_eos_token:
- special_tokens.append((eos, eos_token_id))
- self._tokenizer.post_processor = processors.TemplateProcessing(
- single=single, pair=pair, special_tokens=special_tokens
- )
-
- @property
- def prefix_token(self):
- return self._prefix_token
-
- @property
- def prefix_id(self):
- if self._prefix_token is None:
- return None
- return self.convert_tokens_to_ids(self.prefix_token)
-
- @property
- def middle_token(self):
- return self._middle_token
-
- @property
- def middle_id(self):
- if self._middle_token is None:
- return None
- return self.convert_tokens_to_ids(self.middle_token)
-
- @property
- def suffix_token(self):
- return self._suffix_token
-
- @property
- def suffix_id(self):
- if self._suffix_token is None:
- return None
- return self.convert_tokens_to_ids(self.suffix_token)
-
- @property
- def eot_id(self):
- if self._eot_token is None:
- return None
- return self.convert_tokens_to_ids(self.eot_token)
-
- @property
- def eot_token(self):
- return self._eot_token
-
- @property
- def add_eos_token(self):
- return self._add_eos_token
-
- @property
- def add_bos_token(self):
- return self._add_bos_token
-
- @add_eos_token.setter
- def add_eos_token(self, value):
- self._add_eos_token = value
- self.update_post_processor()
-
- @add_bos_token.setter
- def add_bos_token(self, value):
- self._add_bos_token = value
- self.update_post_processor()
-
- def set_infilling_processor(self, reset, suffix_first=False, add_special_tokens=True):
- """
- Updates the normalizer to make sure the prompt format for `infilling` is respected. The infilling format is the
- following: if suffix_first
- " {suf} {pre}"
- else:
- " {pre} {suf} "
-
- If `reset` is set to `True`, the `normalizer` and `post_processor` are reset to their "normal" behaviour, which
- is to add a prefix space for the normalizer, and add a `bos_token` to the input text for the `post_processor`.
- """
- if reset:
- self._tokenizer.normalizer = normalizers.Sequence(
- [
- normalizers.Prepend(prepend="▁"),
- normalizers.Replace(pattern=" ", content="▁"),
- ]
- )
- self.update_post_processor()
- return
-
- self._tokenizer.normalizer = normalizers.Replace(pattern=" ", content="▁")
- pair = [self.bos_token] if self.add_bos_token and add_special_tokens else []
- special_tokens = [(self.bos_token, self.bos_token_id)] if self.add_bos_token and add_special_tokens else []
- if suffix_first:
- # format as " {suf} {pre}"
- pair += [self.prefix_token, self.suffix_token, "$B", self.middle_token, "$A"]
- special_tokens += [
- (self.prefix_token, self.prefix_id),
- (self.suffix_token, self.suffix_id),
- (self.middle_token, self.middle_id),
- ]
- else:
- # format as " {pre} {suf} "
- pair += [self.prefix_token, "$A", self.suffix_token, "$B", self.middle_token]
- special_tokens += [
- (self.prefix_token, self.prefix_id),
- (self.suffix_token, self.suffix_id),
- (self.middle_token, self.middle_id),
- ]
-
- if self.add_eos_token and add_special_tokens:
- pair += [self.eos_token]
- special_tokens += [(self.eos_token, self.eos_token_id)]
- self._tokenizer.post_processor = processors.TemplateProcessing(
- single="$A", pair=pair, special_tokens=special_tokens
- )
-
- def encode_plus(self, text, text_pair=None, suffix_first=False, add_special_tokens=True, **kwargs):
- # hack to make sure the input is pre-process but outside rust
- text_pair = kwargs.pop("suffix", text_pair)
- if self.fill_token is not None and self.fill_token in text and text_pair is None:
- text, text_pair = text.split(self.fill_token)
-
- if text_pair is None or len(text_pair) < 1:
- return super().encode_plus(text, text_pair, add_special_tokens=add_special_tokens, **kwargs)
-
- if None in (self.prefix_id, self.middle_id, self.suffix_id):
- raise ValueError(
- "Then input includes a `prefix` and a `suffix` used for the infilling task,"
- " the `prefix_id, middle_id, suffix_id` must all be initialized. Current"
- f" values : {self.prefix_id, self.middle_id, self.suffix_id}"
- )
-
- self.set_infilling_processor(False, suffix_first=suffix_first, add_special_tokens=add_special_tokens)
- tokens = super().encode_plus(" " + text, text_pair=text_pair, add_special_tokens=True, **kwargs)
- self.set_infilling_processor(True)
- return tokens
-
- # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not self.can_save_slow_tokenizer:
- raise ValueError(
- "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
- "tokenizer."
- )
-
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
-
- return (out_vocab_file,)
-
- @property
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.default_chat_template
- def default_chat_template(self):
- """
- LLaMA uses [INST] and [/INST] to indicate user messages, and <> and <> to indicate system messages.
- Assistant messages do not have special tokens, because LLaMA chat models are generally trained with strict
- user/assistant/user/assistant message ordering, and so assistant messages can be identified from the ordering
- rather than needing special tokens. The system message is partly 'embedded' in the first user message, which
- results in an unusual token ordering when it is present. This template should definitely be changed if you wish
- to fine-tune a model with more flexible role ordering!
-
- The output should look something like:
-
- [INST] B_SYS SystemPrompt E_SYS Prompt [/INST] Answer [INST] Prompt [/INST] Answer
- [INST] Prompt [/INST]
-
- The reference for this chat template is [this code
- snippet](https://github.com/facebookresearch/llama/blob/556949fdfb72da27c2f4a40b7f0e4cf0b8153a28/llama/generation.py#L320-L362)
- in the original repository.
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- template = (
- "{% if messages[0]['role'] == 'system' %}"
- "{% set loop_messages = messages[1:] %}" # Extract system message if it's present
- "{% set system_message = messages[0]['content'] %}"
- "{% elif USE_DEFAULT_PROMPT == true and not '<>' in messages[0]['content'] %}"
- "{% set loop_messages = messages %}" # Or use the default system message if the flag is set
- "{% set system_message = 'DEFAULT_SYSTEM_MESSAGE' %}"
- "{% else %}"
- "{% set loop_messages = messages %}"
- "{% set system_message = false %}"
- "{% endif %}"
- "{% for message in loop_messages %}" # Loop over all non-system messages
- "{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}"
- "{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}"
- "{% endif %}"
- "{% if loop.index0 == 0 and system_message != false %}" # Embed system message in first message
- "{% set content = '<>\\n' + system_message + '\\n<>\\n\\n' + message['content'] %}"
- "{% else %}"
- "{% set content = message['content'] %}"
- "{% endif %}"
- "{% if message['role'] == 'user' %}" # After all of that, handle messages/roles in a fairly normal way
- "{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}"
- "{% elif message['role'] == 'system' %}"
- "{{ '<>\\n' + content.strip() + '\\n<>\\n\\n' }}"
- "{% elif message['role'] == 'assistant' %}"
- "{{ ' ' + content.strip() + ' ' + eos_token }}"
- "{% endif %}"
- "{% endfor %}"
- )
- template = template.replace("USE_DEFAULT_PROMPT", "true" if self.use_default_system_prompt else "false")
- default_message = DEFAULT_SYSTEM_PROMPT.replace("\n", "\\n").replace("'", "\\'")
- template = template.replace("DEFAULT_SYSTEM_MESSAGE", default_message)
-
- return template
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. The special tokens depend on calling set_lang.
-
- An NLLB sequence has the following format, where `X` represents the sequence:
-
- - `input_ids` (for encoder) `X [eos, src_lang_code]`
- - `decoder_input_ids`: (for decoder) `X [eos, tgt_lang_code]`
-
- BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a
- separator.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: list of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return self.bos_token_id + token_ids_0 + self.eos_token_id
- return self.bos_token_id + token_ids_0 + token_ids_1 + self.eos_token_id
diff --git a/transformers/models/codegen/__init__.py b/transformers/models/codegen/__init__.py
deleted file mode 100644
index a1ce89620035d50db1c4e1878763cddec62f94f2..0000000000000000000000000000000000000000
--- a/transformers/models/codegen/__init__.py
+++ /dev/null
@@ -1,73 +0,0 @@
-# Copyright 2022 Salesforce authors, The EleutherAI, and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available
-
-
-_import_structure = {
- "configuration_codegen": ["CODEGEN_PRETRAINED_CONFIG_ARCHIVE_MAP", "CodeGenConfig", "CodeGenOnnxConfig"],
- "tokenization_codegen": ["CodeGenTokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_codegen_fast"] = ["CodeGenTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_codegen"] = [
- "CODEGEN_PRETRAINED_MODEL_ARCHIVE_LIST",
- "CodeGenForCausalLM",
- "CodeGenModel",
- "CodeGenPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_codegen import CODEGEN_PRETRAINED_CONFIG_ARCHIVE_MAP, CodeGenConfig, CodeGenOnnxConfig
- from .tokenization_codegen import CodeGenTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_codegen_fast import CodeGenTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_codegen import (
- CODEGEN_PRETRAINED_MODEL_ARCHIVE_LIST,
- CodeGenForCausalLM,
- CodeGenModel,
- CodeGenPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/codegen/__pycache__/__init__.cpython-310.pyc b/transformers/models/codegen/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 562dc61e4802e21d2ee892e1e160125e88de3ca7..0000000000000000000000000000000000000000
Binary files a/transformers/models/codegen/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/codegen/__pycache__/configuration_codegen.cpython-310.pyc b/transformers/models/codegen/__pycache__/configuration_codegen.cpython-310.pyc
deleted file mode 100644
index c39a3360d82517983a5289888cedbc4d9ca8933e..0000000000000000000000000000000000000000
Binary files a/transformers/models/codegen/__pycache__/configuration_codegen.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/codegen/__pycache__/modeling_codegen.cpython-310.pyc b/transformers/models/codegen/__pycache__/modeling_codegen.cpython-310.pyc
deleted file mode 100644
index 408bdea360a89ed5150ef136a7f76bc2a69855fe..0000000000000000000000000000000000000000
Binary files a/transformers/models/codegen/__pycache__/modeling_codegen.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/codegen/__pycache__/tokenization_codegen.cpython-310.pyc b/transformers/models/codegen/__pycache__/tokenization_codegen.cpython-310.pyc
deleted file mode 100644
index b1f9bf89521a53a2ede7ce589797c8669f60db96..0000000000000000000000000000000000000000
Binary files a/transformers/models/codegen/__pycache__/tokenization_codegen.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/codegen/__pycache__/tokenization_codegen_fast.cpython-310.pyc b/transformers/models/codegen/__pycache__/tokenization_codegen_fast.cpython-310.pyc
deleted file mode 100644
index 97baab51a937de8de442dde5c209652f179803fd..0000000000000000000000000000000000000000
Binary files a/transformers/models/codegen/__pycache__/tokenization_codegen_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/codegen/configuration_codegen.py b/transformers/models/codegen/configuration_codegen.py
deleted file mode 100644
index e16dd1fadcf74aedbc9728240a6b944c5a298553..0000000000000000000000000000000000000000
--- a/transformers/models/codegen/configuration_codegen.py
+++ /dev/null
@@ -1,229 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Salesforce authors, The EleutherAI, and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" CodeGen model configuration"""
-from collections import OrderedDict
-from typing import Any, List, Mapping, Optional
-
-from ... import PreTrainedTokenizer, TensorType, is_torch_available
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfigWithPast, PatchingSpec
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import CODEGEN_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class CodeGenConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`CodeGenModel`]. It is used to instantiate a
- CodeGen model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the CodeGen
- [Salesforce/codegen-2B-mono](https://huggingface.co/Salesforce/codegen-2B-mono) architecture. Configuration objects
- inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from
- [`PretrainedConfig`] for more information.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 50400):
- Vocabulary size of the CodeGen model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`CodeGenModel`].
- n_positions (`int`, *optional*, defaults to 2048):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- n_ctx (`int`, *optional*, defaults to 2048):
- This attribute is used in `CodeGenModel.__init__` without any real effect.
- n_embd (`int`, *optional*, defaults to 4096):
- Dimensionality of the embeddings and hidden states.
- n_layer (`int`, *optional*, defaults to 28):
- Number of hidden layers in the Transformer encoder.
- n_head (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- rotary_dim (`int`, *optional*, defaults to 64):
- Number of dimensions in the embedding that Rotary Position Embedding is applied to.
- n_inner (`int`, *optional*):
- Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
- activation_function (`str`, *optional*, defaults to `"gelu_new"`):
- Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new"]`.
- resid_pdrop (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- embd_pdrop (`int`, *optional*, defaults to 0.0):
- The dropout ratio for the embeddings.
- attn_pdrop (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention.
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
- The epsilon to use in the layer normalization layers.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
- bos_token_id (`int`, *optional*, defaults to 50256):
- Beginning of stream token id.
- eos_token_id (`int`, *optional*, defaults to 50256):
- End of stream token id.
- tie_word_embeddings (`bool`, *optional*, defaults to `False`):
- Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the
- model has a output word embedding layer.
-
- Example:
-
- ```python
- >>> from transformers import CodeGenConfig, CodeGenModel
-
- >>> # Initializing a CodeGen 6B configuration
- >>> configuration = CodeGenConfig()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = CodeGenModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "codegen"
- attribute_map = {
- "max_position_embeddings": "n_positions",
- "hidden_size": "n_embd",
- "num_attention_heads": "n_head",
- "num_hidden_layers": "n_layer",
- }
-
- def __init__(
- self,
- vocab_size=50400,
- n_positions=2048,
- n_ctx=2048,
- n_embd=4096,
- n_layer=28,
- n_head=16,
- rotary_dim=64,
- n_inner=None,
- activation_function="gelu_new",
- resid_pdrop=0.0,
- embd_pdrop=0.0,
- attn_pdrop=0.0,
- layer_norm_epsilon=1e-5,
- initializer_range=0.02,
- use_cache=True,
- bos_token_id=50256,
- eos_token_id=50256,
- tie_word_embeddings=False,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.n_ctx = n_ctx
- self.n_positions = n_positions
- self.n_embd = n_embd
- self.n_layer = n_layer
- self.n_head = n_head
- self.n_inner = n_inner
- self.rotary_dim = rotary_dim
- self.activation_function = activation_function
- self.resid_pdrop = resid_pdrop
- self.embd_pdrop = embd_pdrop
- self.attn_pdrop = attn_pdrop
- self.layer_norm_epsilon = layer_norm_epsilon
- self.initializer_range = initializer_range
- self.use_cache = use_cache
-
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
-
- super().__init__(
- bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
- )
-
-
-# Copied from transformers.models.gpt2.configuration_gpt2.GPT2OnnxConfig
-class CodeGenOnnxConfig(OnnxConfigWithPast):
- def __init__(
- self,
- config: PretrainedConfig,
- task: str = "default",
- patching_specs: List[PatchingSpec] = None,
- use_past: bool = False,
- ):
- super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past)
- if not getattr(self._config, "pad_token_id", None):
- # TODO: how to do that better?
- self._config.pad_token_id = 0
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
- if self.use_past:
- self.fill_with_past_key_values_(common_inputs, direction="inputs")
- common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"}
- else:
- common_inputs["attention_mask"] = {0: "batch", 1: "sequence"}
-
- return common_inputs
-
- @property
- def num_layers(self) -> int:
- return self._config.n_layer
-
- @property
- def num_attention_heads(self) -> int:
- return self._config.n_head
-
- def generate_dummy_inputs(
- self,
- tokenizer: PreTrainedTokenizer,
- batch_size: int = -1,
- seq_length: int = -1,
- is_pair: bool = False,
- framework: Optional[TensorType] = None,
- ) -> Mapping[str, Any]:
- common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
- tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
- )
-
- # We need to order the input in the way they appears in the forward()
- ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]})
-
- # Need to add the past_keys
- if self.use_past:
- if not is_torch_available():
- raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
- else:
- import torch
-
- batch, seqlen = common_inputs["input_ids"].shape
- # Not using the same length for past_key_values
- past_key_values_length = seqlen + 2
- past_shape = (
- batch,
- self.num_attention_heads,
- past_key_values_length,
- self._config.hidden_size // self.num_attention_heads,
- )
- ordered_inputs["past_key_values"] = [
- (torch.zeros(past_shape), torch.zeros(past_shape)) for _ in range(self.num_layers)
- ]
-
- ordered_inputs["attention_mask"] = common_inputs["attention_mask"]
- if self.use_past:
- mask_dtype = ordered_inputs["attention_mask"].dtype
- ordered_inputs["attention_mask"] = torch.cat(
- [ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1
- )
-
- return ordered_inputs
-
- @property
- def default_onnx_opset(self) -> int:
- return 13
diff --git a/transformers/models/codegen/modeling_codegen.py b/transformers/models/codegen/modeling_codegen.py
deleted file mode 100644
index 41f23900c29a2cdf3e5f19d1448f79fe1e9f1c23..0000000000000000000000000000000000000000
--- a/transformers/models/codegen/modeling_codegen.py
+++ /dev/null
@@ -1,719 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Salesforce authors, The EleutherAI, and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch CodeGen model."""
-
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
-from ...modeling_utils import PreTrainedModel
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_codegen import CodeGenConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "Salesforce/codegen-2B-mono"
-_CONFIG_FOR_DOC = "CodeGenConfig"
-
-
-from ..deprecated._archive_maps import CODEGEN_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.gptj.modeling_gptj.create_sinusoidal_positions
-def create_sinusoidal_positions(num_pos: int, dim: int) -> torch.Tensor:
- inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64) / dim))
- sinusoid_inp = torch.einsum("i , j -> i j", torch.arange(num_pos, dtype=torch.int64).float(), inv_freq).float()
- return torch.cat((torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)), dim=1)
-
-
-# Copied from transformers.models.gptj.modeling_gptj.rotate_every_two
-def rotate_every_two(x: torch.Tensor) -> torch.Tensor:
- x1 = x[:, :, :, ::2]
- x2 = x[:, :, :, 1::2]
- x = torch.stack((-x2, x1), dim=-1)
- return x.flatten(-2) # in einsum notation: rearrange(x, '... d j -> ... (d j)')
-
-
-# Copied from transformers.models.gptj.modeling_gptj.apply_rotary_pos_emb
-def apply_rotary_pos_emb(tensor: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
- sin = torch.repeat_interleave(sin[:, :, None, :], 2, 3)
- cos = torch.repeat_interleave(cos[:, :, None, :], 2, 3)
- return (tensor * cos) + (rotate_every_two(tensor) * sin)
-
-
-class CodeGenAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
-
- max_positions = config.max_position_embeddings
- self.register_buffer(
- "causal_mask",
- torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
- 1, 1, max_positions, max_positions
- ),
- persistent=False,
- )
-
- self.attn_dropout = nn.Dropout(config.attn_pdrop)
- self.resid_dropout = nn.Dropout(config.resid_pdrop)
-
- self.embed_dim = config.hidden_size
- self.num_attention_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_attention_heads
- if self.head_dim * self.num_attention_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and"
- f" `num_attention_heads`: {self.num_attention_heads})."
- )
- self.scale_attn = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype())
- self.qkv_proj = nn.Linear(self.embed_dim, self.embed_dim * 3, bias=False)
-
- self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
- self.rotary_dim = config.rotary_dim
- pos_embd_dim = self.rotary_dim or self.embed_dim
- self.embed_positions = create_sinusoidal_positions(max_positions, pos_embd_dim)
-
- def _split_heads(self, x, n_head, dim_head, mp_num):
- reshaped = x.reshape(x.shape[:-1] + (n_head // mp_num, dim_head))
- reshaped = reshaped.reshape(x.shape[:-2] + (-1,) + reshaped.shape[-1:])
- return reshaped
-
- def _merge_heads(self, tensor, num_attention_heads, attn_head_size):
- """
- Merges attn_head_size dim and num_attn_heads dim into n_ctx
- """
- if len(tensor.shape) == 5:
- tensor = tensor.permute(0, 1, 3, 2, 4).contiguous()
- elif len(tensor.shape) == 4:
- tensor = tensor.permute(0, 2, 1, 3).contiguous()
- else:
- raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}")
- new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,)
- return tensor.view(new_shape)
-
- def _attn(
- self,
- query,
- key,
- value,
- attention_mask=None,
- head_mask=None,
- ):
- # compute causal mask from causal mask buffer
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.causal_mask[:, :, key_length - query_length : key_length, :key_length]
-
- # Keep the attention weights computation in fp32 to avoid overflow issues
- query = query.to(torch.float32)
- key = key.to(torch.float32)
-
- attn_weights = torch.matmul(query, key.transpose(-1, -2))
-
- attn_weights = attn_weights / self.scale_attn
- mask_value = torch.finfo(attn_weights.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
- attn_weights = torch.where(causal_mask, attn_weights, mask_value)
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = nn.Softmax(dim=-1)(attn_weights)
- attn_weights = attn_weights.to(value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def forward(
- self,
- hidden_states: Optional[torch.FloatTensor],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[
- Tuple[torch.Tensor, Tuple[torch.Tensor]],
- Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
- ]:
- qkv = self.qkv_proj(hidden_states)
- # TODO(enijkamp): factor out number of logical TPU-v4 cores or make forward pass agnostic
- mp_num = 4
- qkv_split = qkv.reshape(qkv.shape[:-1] + (mp_num, -1))
-
- local_dim = self.head_dim * self.num_attention_heads // mp_num
- query, value, key = torch.split(qkv_split, local_dim, dim=-1)
- query = self._split_heads(query, self.num_attention_heads, self.head_dim, mp_num=mp_num)
- key = self._split_heads(key, self.num_attention_heads, self.head_dim, mp_num=mp_num)
-
- value = self._split_heads(value, self.num_attention_heads, self.head_dim, mp_num=mp_num)
- value = value.permute(0, 2, 1, 3)
-
- embed_positions = self.embed_positions
- if embed_positions.device != position_ids.device:
- embed_positions = embed_positions.to(position_ids.device)
- self.embed_positions = embed_positions
-
- sincos = embed_positions[position_ids]
- sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1)
-
- if self.rotary_dim is not None:
- k_rot = key[:, :, :, : self.rotary_dim]
- k_pass = key[:, :, :, self.rotary_dim :]
-
- q_rot = query[:, :, :, : self.rotary_dim]
- q_pass = query[:, :, :, self.rotary_dim :]
-
- k_rot = apply_rotary_pos_emb(k_rot, sin, cos)
- q_rot = apply_rotary_pos_emb(q_rot, sin, cos)
-
- key = torch.cat([k_rot, k_pass], dim=-1)
- query = torch.cat([q_rot, q_pass], dim=-1)
- else:
- key = apply_rotary_pos_emb(key, sin, cos)
- query = apply_rotary_pos_emb(query, sin, cos)
-
- key = key.permute(0, 2, 1, 3)
- query = query.permute(0, 2, 1, 3)
-
- if layer_past is not None:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- if use_cache is True:
- # Note that this cast is quite ugly, but is not implemented before ROPE as k_rot in the original codebase is always in fp32.
- # Reference: https://github.com/salesforce/CodeGen/blob/f210c3bb1216c975ad858cd4132c0fdeabf4bfc2/codegen1/jaxformer/hf/codegen/modeling_codegen.py#L38
- present = (key.to(hidden_states.dtype), value)
- else:
- present = None
-
- # compute self-attention: V x Softmax(QK^T)
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
- attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim)
- attn_output = self.out_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs # a, present, (attentions)
-
-
-# Copied from transformers.models.gptj.modeling_gptj.GPTJMLP with GPTJ->CodeGen
-class CodeGenMLP(nn.Module):
- def __init__(self, intermediate_size, config): # in MLP: intermediate_size= 4 * embed_dim
- super().__init__()
- embed_dim = config.n_embd
-
- self.fc_in = nn.Linear(embed_dim, intermediate_size)
- self.fc_out = nn.Linear(intermediate_size, embed_dim)
-
- self.act = ACT2FN[config.activation_function]
- self.dropout = nn.Dropout(config.resid_pdrop)
-
- def forward(self, hidden_states: Optional[torch.FloatTensor]) -> torch.FloatTensor:
- hidden_states = self.fc_in(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.fc_out(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.gptj.modeling_gptj.GPTJBlock with GPTJ->CodeGen
-class CodeGenBlock(nn.Module):
- # Ignore copy
- def __init__(self, config):
- super().__init__()
- inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd
- self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
- self.attn = CodeGenAttention(config)
- self.mlp = CodeGenMLP(inner_dim, config)
-
- def forward(
- self,
- hidden_states: Optional[torch.FloatTensor],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states=hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
- outputs = attn_outputs[1:]
-
- feed_forward_hidden_states = self.mlp(hidden_states)
- hidden_states = attn_output + feed_forward_hidden_states + residual
-
- if use_cache:
- outputs = (hidden_states,) + outputs
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- return outputs # hidden_states, present, (attentions)
-
-
-class CodeGenPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = CodeGenConfig
- base_model_prefix = "transformer"
- supports_gradient_checkpointing = True
- _no_split_modules = ["CodeGenBlock"]
- _skip_keys_device_placement = "past_key_values"
-
- def __init__(self, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, (nn.Linear,)):
- # Slightly different from Mesh Transformer JAX which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-CODEGEN_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`CodeGenConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CODEGEN_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoProcenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_attention_heads,)` or `(n_layer, num_attention_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_dim)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare CodeGen Model transformer outputting raw hidden-states without any specific head on top.",
- CODEGEN_START_DOCSTRING,
-)
-class CodeGenModel(CodeGenPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.embed_dim = config.n_embd
- self.vocab_size = config.vocab_size
- self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
- self.drop = nn.Dropout(config.embd_pdrop)
- self.h = nn.ModuleList([CodeGenBlock(config) for _ in range(config.n_layer)])
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
- self.rotary_dim = min(config.rotary_dim, config.n_ctx // config.num_attention_heads)
-
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.wte
-
- def set_input_embeddings(self, new_embeddings):
- self.wte = new_embeddings
-
- @add_start_docstrings_to_model_forward(CODEGEN_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPast]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- input_ids = input_ids.view(-1, input_shape[-1])
- batch_size = input_ids.shape[0]
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- batch_size = inputs_embeds.shape[0]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if token_type_ids is not None:
- token_type_ids = token_type_ids.view(-1, input_shape[-1])
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * len(self.h))
- else:
- past_length = past_key_values[0][0].size(-2)
-
- if position_ids is None:
- position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0)
-
- # Attention mask.
- if attention_mask is not None:
- if batch_size <= 0:
- raise ValueError("batch_size has to be defined and > 0")
- attention_mask = attention_mask.view(batch_size, -1)
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask = attention_mask[:, None, None, :]
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and the dtype's smallest value for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
- attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x num_attention_heads x N x N
- # head_mask has shape n_layer x batch x num_attention_heads x N x N
- head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-
- if inputs_embeds is None:
- inputs_embeds = self.wte(input_ids)
-
- hidden_states = inputs_embeds
-
- if token_type_ids is not None:
- token_type_embeds = self.wte(token_type_ids)
- hidden_states = hidden_states + token_type_embeds
-
- hidden_states = self.drop(hidden_states)
-
- output_shape = input_shape + (hidden_states.size(-1),)
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
- "`use_cache=False`..."
- )
- use_cache = False
-
- presents = () if use_cache else None
- all_self_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- None,
- attention_mask,
- position_ids,
- head_mask[i],
- use_cache,
- output_attentions,
- )
- else:
- outputs = block(
- hidden_states=hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask[i],
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
-
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
-
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = hidden_states.view(output_shape)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
-
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The CodeGen Model transformer with a language modeling head on top.
- """,
- CODEGEN_START_DOCSTRING,
-)
-class CodeGenForCausalLM(CodeGenPreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.transformer = CodeGenModel(config)
- self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # Omit tokens covered by past_key_values
- if past_key_values:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
- if token_type_ids is not None:
- token_type_ids = token_type_ids[:, -input_ids.shape[1] :]
-
- attention_mask = kwargs.get("attention_mask", None)
- position_ids = kwargs.get("position_ids", None)
-
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- return {
- "input_ids": input_ids,
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "position_ids": position_ids,
- "attention_mask": attention_mask,
- "token_type_ids": token_type_ids,
- }
-
- @add_start_docstrings_to_model_forward(CODEGEN_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
-
- # make sure sampling in fp16 works correctly and
- # compute loss in fp32 to match with mesh-tf version
- # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179
- lm_logits = self.lm_head(hidden_states).to(torch.float32)
-
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
- loss = loss.to(hidden_states.dtype)
-
- if not return_dict:
- output = (lm_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=lm_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- @staticmethod
- def _reorder_cache(
- past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
- ) -> Tuple[Tuple[torch.Tensor]]:
- """
- This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
- [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
- beam_idx at every generation step.
- """
- return tuple(
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
- for layer_past in past_key_values
- )
diff --git a/transformers/models/codegen/tokenization_codegen.py b/transformers/models/codegen/tokenization_codegen.py
deleted file mode 100644
index 1b03af7008465dd76c7138fae3a6cac5e5a4ef53..0000000000000000000000000000000000000000
--- a/transformers/models/codegen/tokenization_codegen.py
+++ /dev/null
@@ -1,417 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The Salesforce authors, The Open AI Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for CodeGen"""
-
-
-import json
-import os
-from functools import lru_cache
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
-
-import numpy as np
-import regex as re
-
-from ...utils import is_tf_available, is_torch_available, logging, to_py_obj
-
-
-if TYPE_CHECKING:
- if is_torch_available():
- import torch
- if is_tf_available():
- import tensorflow as tf
-
-from ...tokenization_utils import AddedToken, PreTrainedTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {
- "vocab_file": "vocab.json",
- "merges_file": "merges.txt",
-}
-
-
-@lru_cache()
-def bytes_to_unicode():
- """
- Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
- characters the bpe code barfs on.
-
- The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
- if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
- decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
- tables between utf-8 bytes and unicode strings.
- """
- bs = (
- list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
- )
- cs = bs[:]
- n = 0
- for b in range(2**8):
- if b not in bs:
- bs.append(b)
- cs.append(2**8 + n)
- n += 1
- cs = [chr(n) for n in cs]
- return dict(zip(bs, cs))
-
-
-def get_pairs(word):
- """
- Return set of symbol pairs in a word.
-
- Word is represented as tuple of symbols (symbols being variable-length strings).
- """
- pairs = set()
- prev_char = word[0]
- for char in word[1:]:
- pairs.add((prev_char, char))
- prev_char = char
- return pairs
-
-
-class CodeGenTokenizer(PreTrainedTokenizer):
- """
- Construct a CodeGen tokenizer. Based on byte-level Byte-Pair-Encoding.
-
- This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
- be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
- ```python
- >>> from transformers import CodeGenTokenizer
-
- >>> tokenizer = CodeGenTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
- >>> tokenizer("Hello world")["input_ids"]
- [15496, 995]
-
- >>> tokenizer(" Hello world")["input_ids"]
- [18435, 995]
- ```
-
- You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer or when you
- call it on some text, but since the model was not pretrained this way, it might yield a decrease in performance.
-
-
-
- When used with `is_split_into_words=True`, this tokenizer will add a space before each word (even the first one).
-
-
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- merges_file (`str`):
- Path to the merges file.
- errors (`str`, *optional*, defaults to `"replace"`):
- Paradigm to follow when decoding bytes to UTF-8. See
- [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
- unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The end of sequence token.
- pad_token (`str`, *optional*):
- The token used for padding, for example when batching sequences of different lengths.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial space to the input. This allows to treat the leading word just as any
- other word. (CodeGen tokenizer detect beginning of words by the preceding space).
- add_bos_token (`bool`, *optional*, defaults to `False`):
- Whether to add a beginning of sequence token at the start of sequences.
- return_token_type_ids (`bool`, *optional*, defaults to `False`):
- Whether to return token type IDs.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- merges_file,
- errors="replace",
- unk_token="<|endoftext|>",
- bos_token="<|endoftext|>",
- eos_token="<|endoftext|>",
- pad_token=None,
- add_prefix_space=False,
- add_bos_token=False,
- return_token_type_ids=False,
- **kwargs,
- ):
- bos_token = AddedToken(bos_token, special=True) if isinstance(bos_token, str) else bos_token
- eos_token = AddedToken(eos_token, special=True) if isinstance(eos_token, str) else eos_token
- unk_token = AddedToken(unk_token, special=True) if isinstance(unk_token, str) else unk_token
- pad_token = AddedToken(pad_token, special=True) if isinstance(pad_token, str) else pad_token
- self.add_bos_token = add_bos_token
- self.return_token_type_ids = return_token_type_ids
- if self.return_token_type_ids:
- self.model_input_names.append("token_type_ids")
-
- with open(vocab_file, encoding="utf-8") as vocab_handle:
- self.encoder = json.load(vocab_handle)
- self.decoder = {v: k for k, v in self.encoder.items()}
- self.errors = errors # how to handle errors in decoding
- self.byte_encoder = bytes_to_unicode()
- self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
- with open(merges_file, encoding="utf-8") as merges_handle:
- bpe_merges = merges_handle.read().split("\n")[1:-1]
- bpe_merges = [tuple(merge.split()) for merge in bpe_merges]
- self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
- self.cache = {}
- self.add_prefix_space = add_prefix_space
-
- # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
- self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
- super().__init__(
- errors=errors,
- unk_token=unk_token,
- bos_token=bos_token,
- eos_token=eos_token,
- pad_token=pad_token,
- add_prefix_space=add_prefix_space,
- add_bos_token=add_bos_token,
- return_token_type_ids=return_token_type_ids,
- **kwargs,
- )
-
- @property
- def vocab_size(self):
- return len(self.encoder)
-
- def get_vocab(self):
- return dict(self.encoder, **self.added_tokens_encoder)
-
- def bpe(self, token):
- if token in self.cache:
- return self.cache[token]
- word = tuple(token)
- pairs = get_pairs(word)
-
- if not pairs:
- return token
-
- while True:
- bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
- if bigram not in self.bpe_ranks:
- break
- first, second = bigram
- new_word = []
- i = 0
- while i < len(word):
- try:
- j = word.index(first, i)
- except ValueError:
- new_word.extend(word[i:])
- break
- else:
- new_word.extend(word[i:j])
- i = j
-
- if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
- new_word.append(first + second)
- i += 2
- else:
- new_word.append(word[i])
- i += 1
- new_word = tuple(new_word)
- word = new_word
- if len(word) == 1:
- break
- else:
- pairs = get_pairs(word)
- word = " ".join(word)
- self.cache[token] = word
- return word
-
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- if self.add_bos_token:
- bos_token_ids = [self.bos_token_id]
- else:
- bos_token_ids = []
-
- output = bos_token_ids + token_ids_0
-
- if token_ids_1 is None:
- return output
-
- return output + bos_token_ids + token_ids_1
-
- def _tokenize(self, text):
- """Tokenize a string."""
- bpe_tokens = []
- for token in re.findall(self.pat, text):
- token = "".join(
- self.byte_encoder[b] for b in token.encode("utf-8")
- ) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
- bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
- return bpe_tokens
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- text = "".join(tokens)
- text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
- return text
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id] if self.sep_token_id is not None else []
- cls = [self.cls_token_id] if self.sep_token_id is not None else []
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- merge_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
- )
-
- with open(vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
- index = 0
- with open(merge_file, "w", encoding="utf-8") as writer:
- writer.write("#version: 0.2\n")
- for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
- " Please check that the tokenizer is not corrupted!"
- )
- index = token_index
- writer.write(" ".join(bpe_tokens) + "\n")
- index += 1
-
- return vocab_file, merge_file
-
- def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
- add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space)
- if is_split_into_words or add_prefix_space:
- text = " " + text
- return (text, kwargs)
-
- def decode(
- self,
- token_ids: Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"],
- skip_special_tokens: bool = False,
- clean_up_tokenization_spaces: bool = None,
- truncate_before_pattern: Optional[List[str]] = None,
- **kwargs,
- ) -> str:
- """
- Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special
- tokens and clean up tokenization spaces.
-
- Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`.
-
- Args:
- token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`):
- List of tokenized input ids. Can be obtained using the `__call__` method.
- skip_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not to remove special tokens in the decoding.
- clean_up_tokenization_spaces (`bool`, *optional*):
- Whether or not to clean up the tokenization spaces. If `None`, will default to
- `self.clean_up_tokenization_spaces` (available in the `tokenizer_config`).
- truncate_before_pattern (`List[str]`, *optional*, defaults to `None`):
- A list of regular expression strings that will be used to truncate the returned string. This can be
- used to remove extra pieces of code (e.g. truncate if observing a comment symbol "#" at the beginning
- of a new line). An example pattern could be `["^#", re.escape("<|endoftext|>"), "^'''", "\n\n\n"]`.
- kwargs (additional keyword arguments, *optional*):
- Will be passed to the underlying model specific decode method.
-
- Returns:
- `str`: The decoded sentence.
- """
-
- token_ids = to_py_obj(token_ids)
-
- decoded_text = super()._decode(
- token_ids=token_ids,
- skip_special_tokens=skip_special_tokens,
- clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- **kwargs,
- )
-
- if truncate_before_pattern is not None and len(truncate_before_pattern) > 0:
- decoded_text = self.truncate(decoded_text, truncate_before_pattern)
-
- return decoded_text
-
- def truncate(self, completion, truncate_before_pattern):
- def find_re(string, pattern, start_pos):
- m = pattern.search(string, start_pos)
- return m.start() if m else -1
-
- terminals = [re.compile(pattern, re.MULTILINE) for pattern in truncate_before_pattern]
-
- prints = list(re.finditer("^print", completion, re.MULTILINE))
-
- if len(prints) > 1:
- completion = completion[: prints[1].start()]
-
- defs = list(re.finditer("^def", completion, re.MULTILINE))
-
- if len(defs) > 1:
- completion = completion[: defs[1].start()]
-
- start_pos = 0
-
- terminals_pos = [
- pos for pos in [find_re(completion, terminal, start_pos) for terminal in terminals] if pos != -1
- ]
-
- if len(terminals_pos) > 0:
- return completion[: min(terminals_pos)]
- else:
- return completion
diff --git a/transformers/models/codegen/tokenization_codegen_fast.py b/transformers/models/codegen/tokenization_codegen_fast.py
deleted file mode 100644
index b086fb84a65af9b632f219e5911318fa518c3637..0000000000000000000000000000000000000000
--- a/transformers/models/codegen/tokenization_codegen_fast.py
+++ /dev/null
@@ -1,273 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The Salesforce authors, The Open AI Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for OpenAI GPT."""
-
-
-import json
-import re
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
-
-import numpy as np
-
-from ...utils import is_tf_available, is_torch_available, logging
-
-
-if TYPE_CHECKING:
- if is_torch_available():
- import torch
- if is_tf_available():
- import tensorflow as tf
-
-from tokenizers import pre_tokenizers
-
-from ...tokenization_utils_base import BatchEncoding
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from .tokenization_codegen import CodeGenTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}
-
-
-class CodeGenTokenizerFast(PreTrainedTokenizerFast):
- """
- Construct a "fast" CodeGen tokenizer (backed by HuggingFace's *tokenizers* library). Based on byte-level
- Byte-Pair-Encoding.
-
- This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
- be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
- ```python
- >>> from transformers import CodeGenTokenizerFast
-
- >>> tokenizer = CodeGenTokenizerFast.from_pretrained("Salesforce/codegen-350M-mono")
- >>> tokenizer("Hello world")["input_ids"]
- [15496, 995]
-
- >>> tokenizer(" Hello world")["input_ids"]
- [18435, 995]
- ```
-
- You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer, but since
- the model was not pretrained this way, it might yield a decrease in performance.
-
-
-
- When used with `is_split_into_words=True`, this tokenizer needs to be instantiated with `add_prefix_space=True`.
-
-
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`, *optional*):
- Path to the vocabulary file.
- merges_file (`str`, *optional*):
- Path to the merges file.
- tokenizer_file (`str`, *optional*):
- Path to [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that
- contains everything needed to load the tokenizer.
- unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The end of sequence token.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial space to the input. This allows to treat the leading word just as any
- other word. (CodeGen tokenizer detect beginning of words by the preceding space).
- return_token_type_ids (`bool`, *optional*, defaults to `False`):
- Whether to return token type IDs.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
- slow_tokenizer_class = CodeGenTokenizer
-
- def __init__(
- self,
- vocab_file=None,
- merges_file=None,
- tokenizer_file=None,
- unk_token="<|endoftext|>",
- bos_token="<|endoftext|>",
- eos_token="<|endoftext|>",
- add_prefix_space=False,
- return_token_type_ids=False,
- **kwargs,
- ):
- self.return_token_type_ids = return_token_type_ids
- if self.return_token_type_ids:
- self.model_input_names.append("token_type_ids")
-
- super().__init__(
- vocab_file,
- merges_file,
- tokenizer_file=tokenizer_file,
- unk_token=unk_token,
- bos_token=bos_token,
- eos_token=eos_token,
- add_prefix_space=add_prefix_space,
- return_token_type_ids=return_token_type_ids,
- **kwargs,
- )
-
- if kwargs.pop("add_bos_token", False):
- model_id = kwargs.pop("name_or_path", "")
- raise ValueError(
- "Currenty GPT2's fast tokenizer does NOT support adding a BOS token. "
- "Instead you should use GPT2's slow tokenizer class `CodeGenTokenizer` as follows: \n"
- f"`CodeGenTokenizer.from_pretrained('{model_id}')`\nor\n"
- f"`AutoTokenizer.from_pretrained('{model_id}', use_fast=False)`\n"
- "This issue will be fixed soon, see: https://github.com/huggingface/tokenizers/pull/1005."
- " so that the fast tokenizer works correctly."
- )
-
- pre_tok_state = json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__())
- if pre_tok_state.get("add_prefix_space", add_prefix_space) != add_prefix_space:
- pre_tok_class = getattr(pre_tokenizers, pre_tok_state.pop("type"))
- pre_tok_state["add_prefix_space"] = add_prefix_space
- self.backend_tokenizer.pre_tokenizer = pre_tok_class(**pre_tok_state)
-
- self.add_prefix_space = add_prefix_space
-
- def _batch_encode_plus(self, *args, **kwargs) -> BatchEncoding:
- is_split_into_words = kwargs.get("is_split_into_words", False)
- assert self.add_prefix_space or not is_split_into_words, (
- f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
- "to use it with pretokenized inputs."
- )
-
- return super()._batch_encode_plus(*args, **kwargs)
-
- def _encode_plus(self, *args, **kwargs) -> BatchEncoding:
- is_split_into_words = kwargs.get("is_split_into_words", False)
-
- assert self.add_prefix_space or not is_split_into_words, (
- f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
- "to use it with pretokenized inputs."
- )
-
- return super()._encode_plus(*args, **kwargs)
-
- # Copied from transformers.models.codegen.tokenization_codegen.CodeGenTokenizer.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id] if self.sep_token_id is not None else []
- cls = [self.cls_token_id] if self.sep_token_id is not None else []
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
-
- def decode(
- self,
- token_ids: Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"],
- skip_special_tokens: bool = False,
- clean_up_tokenization_spaces: bool = None,
- truncate_before_pattern: Optional[List[str]] = None,
- **kwargs,
- ) -> str:
- """
- Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special
- tokens and clean up tokenization spaces.
-
- Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`.
-
- Args:
- token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`):
- List of tokenized input ids. Can be obtained using the `__call__` method.
- skip_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not to remove special tokens in the decoding.
- clean_up_tokenization_spaces (`bool`, *optional*):
- Whether or not to clean up the tokenization spaces. If `None`, will default to
- `self.clean_up_tokenization_spaces` (available in the `tokenizer_config`).
- truncate_before_pattern (`List[str]`, *optional*, defaults to `None`):
- A list of regular expression strings that will be used to truncate the returned string. This can be
- used to remove extra pieces of code (e.g. truncate if observing a comment symbol "#" at the beginning
- of a new line). An example pattern could be `["^#", re.escape("<|endoftext|>"), "^'''", "\n\n\n"]`.
- kwargs (additional keyword arguments, *optional*):
- Will be passed to the underlying model specific decode method.
-
- Returns:
- `str`: The decoded sentence.
- """
-
- decoded_text = super().decode(
- token_ids=token_ids,
- skip_special_tokens=skip_special_tokens,
- clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- **kwargs,
- )
-
- if truncate_before_pattern is not None and len(truncate_before_pattern) > 0:
- decoded_text = self.truncate(decoded_text, truncate_before_pattern)
-
- return decoded_text
-
- def truncate(self, completion, truncate_before_pattern):
- def find_re(string, pattern, start_pos):
- m = pattern.search(string, start_pos)
- return m.start() if m else -1
-
- terminals = [re.compile(pattern, re.MULTILINE) for pattern in truncate_before_pattern]
-
- prints = list(re.finditer("^print", completion, re.MULTILINE))
-
- if len(prints) > 1:
- completion = completion[: prints[1].start()]
-
- defs = list(re.finditer("^def", completion, re.MULTILINE))
-
- if len(defs) > 1:
- completion = completion[: defs[1].start()]
-
- start_pos = 0
-
- terminals_pos = [
- pos for pos in [find_re(completion, terminal, start_pos) for terminal in terminals] if pos != -1
- ]
-
- if len(terminals_pos) > 0:
- return completion[: min(terminals_pos)]
- else:
- return completion
diff --git a/transformers/models/cohere/__init__.py b/transformers/models/cohere/__init__.py
deleted file mode 100644
index d6f69d1e496d0e616b185c819fd847fcc019fc39..0000000000000000000000000000000000000000
--- a/transformers/models/cohere/__init__.py
+++ /dev/null
@@ -1,77 +0,0 @@
-# Copyright 2024 Cohere and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_sentencepiece_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_cohere": ["COHERE_PRETRAINED_CONFIG_ARCHIVE_MAP", "CohereConfig"],
-}
-
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_cohere_fast"] = ["CohereTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_cohere"] = [
- "CohereForCausalLM",
- "CohereModel",
- "CoherePreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_cohere import COHERE_PRETRAINED_CONFIG_ARCHIVE_MAP, CohereConfig
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_cohere_fast import CohereTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_cohere import (
- CohereForCausalLM,
- CohereModel,
- CoherePreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/cohere/__pycache__/__init__.cpython-310.pyc b/transformers/models/cohere/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index cbe28f568c8cd5c9e1beec3becc699ebe893a989..0000000000000000000000000000000000000000
Binary files a/transformers/models/cohere/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cohere/__pycache__/configuration_cohere.cpython-310.pyc b/transformers/models/cohere/__pycache__/configuration_cohere.cpython-310.pyc
deleted file mode 100644
index ce12d1006503e4d2aaf74b09b7cc26fb9113ce1b..0000000000000000000000000000000000000000
Binary files a/transformers/models/cohere/__pycache__/configuration_cohere.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cohere/__pycache__/modeling_cohere.cpython-310.pyc b/transformers/models/cohere/__pycache__/modeling_cohere.cpython-310.pyc
deleted file mode 100644
index 0be9ff844d7855a983b9863cbb75a42e893cd513..0000000000000000000000000000000000000000
Binary files a/transformers/models/cohere/__pycache__/modeling_cohere.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cohere/__pycache__/tokenization_cohere_fast.cpython-310.pyc b/transformers/models/cohere/__pycache__/tokenization_cohere_fast.cpython-310.pyc
deleted file mode 100644
index a716f542ccee84616e01aa41455bd08f47cab311..0000000000000000000000000000000000000000
Binary files a/transformers/models/cohere/__pycache__/tokenization_cohere_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cohere/configuration_cohere.py b/transformers/models/cohere/configuration_cohere.py
deleted file mode 100644
index 7ceca2b887af7d6529bb9a8ee3a8e65b25a8fe1f..0000000000000000000000000000000000000000
--- a/transformers/models/cohere/configuration_cohere.py
+++ /dev/null
@@ -1,159 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Cohere team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Cohere model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-COHERE_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
-
-
-class CohereConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`CohereModel`]. It is used to instantiate an Cohere
- model according to the specified arguments, defining the model architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the [CohereForAI/c4ai-command-r-v01](https://huggingface.co/CohereForAI/c4ai-command-r-v01) model.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 256000):
- Vocabulary size of the Cohere model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`CohereModel`]
- hidden_size (`int`, *optional*, defaults to 8192):
- Dimension of the hidden representations.
- intermediate_size (`int`, *optional*, defaults to 22528):
- Dimension of the MLP representations.
- logit_scale (`float`, *optional*, defaults to 0.0625):
- The scaling factor for the output logits.
- num_hidden_layers (`int`, *optional*, defaults to 40):
- Number of hidden layers in the Transformer decoder.
- num_attention_heads (`int`, *optional*, defaults to 64):
- Number of attention heads for each attention layer in the Transformer decoder.
- num_key_value_heads (`int`, *optional*):
- This is the number of key_value heads that should be used to implement Grouped Query Attention. If
- `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
- `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
- converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
- by meanpooling all the original heads within that group. For more details checkout [this
- paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
- `num_attention_heads`.
- hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
- The non-linear activation function (function or string) in the decoder.
- max_position_embeddings (`int`, *optional*, defaults to 8192):
- The maximum sequence length that this model might ever be used with.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the layer normalization.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- pad_token_id (`int`, *optional*, defaults to 0):
- Padding token id.
- bos_token_id (`int`, *optional*, defaults to 5):
- Beginning of stream token id.
- eos_token_id (`int`, *optional*, defaults to 255001):
- End of stream token id.
- tie_word_embeddings (`bool`, *optional*, defaults to `True`):
- Whether to tie weight embeddings
- rope_theta (`float`, *optional*, defaults to 10000.0):
- The base period of the RoPE embeddings.
- attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
- Whether to use a bias in the query, key, value and output projection layers during self-attention.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- use_qk_norm (`bool`, *optional*, defaults to `False`):
- Whether to use query-key normalization in the attention
-
- ```python
- >>> from transformers import CohereModel, CohereConfig
-
- >>> # Initializing a Cohere model configuration
- >>> configuration = CohereConfig()
-
- >>> # Initializing a model from the Cohere configuration
- >>> model = CohereModel(configuration) # doctest: +SKIP
-
- >>> # Accessing the model configuration
- >>> configuration = model.config # doctest: +SKIP
- ```"""
-
- model_type = "cohere"
- keys_to_ignore_at_inference = ["past_key_values"]
-
- def __init__(
- self,
- vocab_size=256000,
- hidden_size=8192,
- intermediate_size=22528,
- logit_scale=0.0625,
- num_hidden_layers=40,
- num_attention_heads=64,
- num_key_value_heads=None,
- hidden_act="silu",
- max_position_embeddings=8192,
- initializer_range=0.02,
- layer_norm_eps=1e-5,
- use_cache=True,
- pad_token_id=0,
- bos_token_id=5,
- eos_token_id=255001,
- tie_word_embeddings=True,
- rope_theta=10000.0,
- attention_bias=False,
- attention_dropout=0.0,
- use_qk_norm=False,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.hidden_size = hidden_size
- self.logit_scale = logit_scale
- self.intermediate_size = intermediate_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
-
- # for backward compatibility
- if num_key_value_heads is None:
- num_key_value_heads = num_attention_heads
-
- self.num_key_value_heads = num_key_value_heads
- self.hidden_act = hidden_act
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.attention_bias = attention_bias
- self.attention_dropout = attention_dropout
- self.use_qk_norm = use_qk_norm
-
- super().__init__(
- pad_token_id=pad_token_id,
- bos_token_id=bos_token_id,
- eos_token_id=eos_token_id,
- tie_word_embeddings=tie_word_embeddings,
- **kwargs,
- )
diff --git a/transformers/models/cohere/modeling_cohere.py b/transformers/models/cohere/modeling_cohere.py
deleted file mode 100644
index 950d45ea867a3080bb804565edff1b2bb0e09e97..0000000000000000000000000000000000000000
--- a/transformers/models/cohere/modeling_cohere.py
+++ /dev/null
@@ -1,1266 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Cohere team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# This file is based on the LLama model definition file in transformers
-
-"""PyTorch Cohere model."""
-
-import math
-import warnings
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from ...activations import ACT2FN
-from ...cache_utils import Cache, DynamicCache, StaticCache
-from ...modeling_attn_mask_utils import AttentionMaskConverter
-from ...modeling_outputs import (
- BaseModelOutputWithPast,
- CausalLMOutputWithPast,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import ALL_LAYERNORM_LAYERS
-from ...utils import (
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- logging,
- replace_return_docstrings,
-)
-from .configuration_cohere import CohereConfig
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "CohereConfig"
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-class CohereLayerNorm(nn.Module):
- def __init__(self, hidden_size=None, eps=1e-5, bias=False):
- """The hidden size can be a tuple or an int. The tuple is used for QKNorm to normalize across head_dim"""
- super().__init__()
- self.weight = nn.Parameter(torch.ones(hidden_size))
- self.variance_epsilon = eps
-
- def forward(self, hidden_states):
- input_dtype = hidden_states.dtype
- hidden_states = hidden_states.to(torch.float32)
- mean = hidden_states.mean(-1, keepdim=True)
- variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True)
- hidden_states = (hidden_states - mean) * torch.rsqrt(variance + self.variance_epsilon)
- hidden_states = self.weight.to(torch.float32) * hidden_states
- return hidden_states.to(input_dtype)
-
-
-ALL_LAYERNORM_LAYERS.append(CohereLayerNorm)
-
-
-class CohereRotaryEmbedding(nn.Module):
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
- super().__init__()
- self.scaling_factor = scaling_factor
- self.dim = dim
- self.max_position_embeddings = max_position_embeddings
- self.base = base
- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
-
- @torch.no_grad()
- def forward(self, x, position_ids):
- # x: [bs, num_attention_heads, seq_len, head_size]
- inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
- position_ids_expanded = position_ids[:, None, :].float()
-
- # Force float32 since bfloat16 loses precision on long contexts
- # See https://github.com/huggingface/transformers/pull/29285
- device_type = x.device.type
- device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
- with torch.autocast(device_type=device_type, enabled=False):
- freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
- emb = torch.repeat_interleave(freqs, 2, dim=-1)
- cos = emb.cos()
- sin = emb.sin()
- return cos, sin
-
-
-def rotate_half(x):
- # Split and rotate
- x1 = x[..., ::2]
- x2 = x[..., 1::2]
- rot_x = torch.stack([-x2, x1], dim=-1).flatten(-2)
- return rot_x
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
- """Applies Rotary Position Embedding to the query and key tensors.
-
- Args:
- q (`torch.Tensor`): The query tensor.
- k (`torch.Tensor`): The key tensor.
- cos (`torch.Tensor`): The cosine part of the rotary embedding.
- sin (`torch.Tensor`): The sine part of the rotary embedding.
- position_ids (`torch.Tensor`, *optional*):
- Deprecated and unused.
- unsqueeze_dim (`int`, *optional*, defaults to 1):
- The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
- sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
- that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
- k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
- cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
- the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
- Returns:
- `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
- """
- dtype = q.dtype
- q = q.float()
- k = k.float()
- cos = cos.unsqueeze(unsqueeze_dim)
- sin = sin.unsqueeze(unsqueeze_dim)
- q_embed = (q * cos) + (rotate_half(q) * sin)
- k_embed = (k * cos) + (rotate_half(k) * sin)
- return q_embed.to(dtype=dtype), k_embed.to(dtype=dtype)
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaMLP Llama->Cohere
-class CohereMLP(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.hidden_size = config.hidden_size
- self.intermediate_size = config.intermediate_size
- self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
- self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
- self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
- self.act_fn = ACT2FN[config.hidden_act]
-
- # Ignore copy
- def forward(self, x):
- down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
- return down_proj
-
-
-# Copied from transformers.models.llama.modeling_llama.repeat_kv
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
- """
- This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
- num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
- """
- batch, num_key_value_heads, slen, head_dim = hidden_states.shape
- if n_rep == 1:
- return hidden_states
- hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
- return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-class CohereAttention(nn.Module):
- """Multi-headed attention from 'Attention Is All You Need' paper"""
-
- def __init__(self, config: CohereConfig, layer_idx: Optional[int] = None):
- super().__init__()
- self.config = config
- self.layer_idx = layer_idx
- if layer_idx is None:
- logger.warning_once(
- f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
- "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
- "when creating this class."
- )
-
- self.attention_dropout = config.attention_dropout
- self.hidden_size = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.hidden_size // self.num_heads
- self.num_key_value_heads = config.num_key_value_heads
- self.num_key_value_groups = self.num_heads // self.num_key_value_heads
- self.max_position_embeddings = config.max_position_embeddings
- self.rope_theta = config.rope_theta
- self.is_causal = True
- self.use_qk_norm = config.use_qk_norm
-
- if (self.head_dim * self.num_heads) != self.hidden_size:
- raise ValueError(
- f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
- f" and `num_heads`: {self.num_heads})."
- )
-
- if self.use_qk_norm:
- # When sharding the model using Tensor Parallelism, need to be careful to use n_local_heads
- self.q_norm = CohereLayerNorm(hidden_size=(self.num_heads, self.head_dim), eps=config.layer_norm_eps)
- self.k_norm = CohereLayerNorm(
- hidden_size=(self.num_key_value_heads, self.head_dim), eps=config.layer_norm_eps
- )
-
- self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
- self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
- self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
- self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.attention_bias)
- self._init_rope()
-
- # Ignore copy
- def _init_rope(self):
- self.rotary_emb = CohereRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- base=self.rope_theta,
- )
-
- # Ignore copy
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- bsz, q_len, _ = hidden_states.size()
-
- query_states = self.q_proj(hidden_states)
- key_states = self.k_proj(hidden_states)
- value_states = self.v_proj(hidden_states)
-
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
- if self.use_qk_norm:
- query_states = self.q_norm(query_states)
- key_states = self.k_norm(key_states)
-
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- past_key_value = getattr(self, "past_key_value", past_key_value)
- cos, sin = self.rotary_emb(value_states, position_ids)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; position_ids needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
- key_states = repeat_kv(key_states, self.num_key_value_groups)
- value_states = repeat_kv(value_states, self.num_key_value_groups)
-
- attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-
- if attention_mask is not None: # no matter the length, we just slice it
- causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
- attn_weights = attn_weights + causal_mask
-
- # upcast attention to fp32
- attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
- attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
- attn_output = torch.matmul(attn_weights, value_states)
-
- if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.transpose(1, 2).contiguous()
-
- attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-
- attn_output = self.o_proj(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights, past_key_value
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 Llama->Cohere
-class CohereFlashAttention2(CohereAttention):
- """
- Cohere flash attention module. This module inherits from `CohereAttention` as the weights of the module stays
- untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
- flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- output_attentions = False
-
- bsz, q_len, _ = hidden_states.size()
-
- query_states = self.q_proj(hidden_states)
- key_states = self.k_proj(hidden_states)
- value_states = self.v_proj(hidden_states)
-
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
- if self.use_qk_norm:
- query_states = self.q_norm(query_states)
- key_states = self.k_norm(key_states)
-
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- cos, sin = self.rotary_emb(value_states, position_ids)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
- past_key_value = getattr(self, "past_key_value", past_key_value)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; position_ids needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
- # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
- # to be able to avoid many of these transpose/reshape/view.
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.transpose(1, 2)
-
- dropout_rate = self.attention_dropout if self.training else 0.0
-
- # Ignore copy
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in the correct dtype just to be sure everything works as expected.
- # This might slowdown training & inference so it is recommended to not cast the LayerNorms
- # in fp32. (CohereLayerNorm handles it correctly)
-
- input_dtype = query_states.dtype
- if input_dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.q_proj.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
-
- query_states = query_states.to(target_dtype)
- key_states = key_states.to(target_dtype)
- value_states = value_states.to(target_dtype)
-
- attn_output = self._flash_attention_forward(
- query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
- )
-
- attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
- attn_output = self.o_proj(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights, past_key_value
-
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in CohereFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention Llama->Cohere
-class CohereSdpaAttention(CohereAttention):
- """
- Cohere attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
- `CohereAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
- SDPA API.
- """
-
- # Ignore copy
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- if output_attentions:
- # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
- logger.warning_once(
- "CohereModel is using CohereSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
- 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
- )
- return super().forward(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- )
-
- bsz, q_len, _ = hidden_states.size()
-
- query_states = self.q_proj(hidden_states)
- key_states = self.k_proj(hidden_states)
- value_states = self.v_proj(hidden_states)
-
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
- if self.use_qk_norm:
- query_states = self.q_norm(query_states)
- key_states = self.k_norm(key_states)
-
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- cos, sin = self.rotary_emb(value_states, position_ids)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
- # In case static cache is used, it is an instance attribute.
- past_key_value = getattr(self, "past_key_value", past_key_value)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; cache_position needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
- key_states = repeat_kv(key_states, self.num_key_value_groups)
- value_states = repeat_kv(value_states, self.num_key_value_groups)
-
- causal_mask = attention_mask
- # if attention_mask is not None and cache_position is not None:
- if attention_mask is not None:
- causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
-
- # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
- # Reference: https://github.com/pytorch/pytorch/issues/112577.
- if query_states.device.type == "cuda" and causal_mask is not None:
- query_states = query_states.contiguous()
- key_states = key_states.contiguous()
- value_states = value_states.contiguous()
-
- # In case we are not compiling, we may set `causal_mask` to None, which is required to dispatch to SDPA's Flash Attention 2 backend, rather
- # relying on the `is_causal` argument.
- attn_output = torch.nn.functional.scaled_dot_product_attention(
- query_states,
- key_states,
- value_states,
- attn_mask=causal_mask,
- dropout_p=self.attention_dropout if self.training else 0.0,
- is_causal=causal_mask is None and q_len > 1,
- )
-
- attn_output = attn_output.transpose(1, 2).contiguous()
- attn_output = attn_output.view(bsz, q_len, self.hidden_size)
-
- attn_output = self.o_proj(attn_output)
-
- return attn_output, None, past_key_value
-
-
-COHERE_ATTENTION_CLASSES = {
- "eager": CohereAttention,
- "flash_attention_2": CohereFlashAttention2,
- "sdpa": CohereSdpaAttention,
-}
-
-
-class CohereDecoderLayer(nn.Module):
- def __init__(self, config: CohereConfig, layer_idx: int):
- super().__init__()
- self.hidden_size = config.hidden_size
-
- self.self_attn = COHERE_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
-
- self.mlp = CohereMLP(config)
- self.input_layernorm = CohereLayerNorm(hidden_size=(config.hidden_size), eps=config.layer_norm_eps)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: Optional[bool] = False,
- use_cache: Optional[bool] = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs,
- ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`, *optional*):
- attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
- query_sequence_length, key_sequence_length)` if default attention is used.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
- """
- if "padding_mask" in kwargs:
- warnings.warn(
- "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
- )
-
- residual = hidden_states
-
- hidden_states = self.input_layernorm(hidden_states)
-
- # Self Attention
- hidden_states_attention, self_attn_weights, present_key_value = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- **kwargs,
- )
-
- # Fully Connected
- hidden_states_mlp = self.mlp(hidden_states)
-
- # Add everything together
- hidden_states = residual + hidden_states_attention + hidden_states_mlp
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights,)
-
- if use_cache:
- outputs += (present_key_value,)
-
- return outputs
-
-
-COHERE_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`CohereConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-@add_start_docstrings(
- "The bare Cohere Model outputting raw hidden-states without any specific head on top.",
- COHERE_START_DOCSTRING,
-)
-# Copied from transformers.models.llama.modeling_llama.LlamaPreTrainedModel with Llama->Cohere
-class CoherePreTrainedModel(PreTrainedModel):
- config_class = CohereConfig
- base_model_prefix = "model"
- supports_gradient_checkpointing = True
- _no_split_modules = ["CohereDecoderLayer"]
- _skip_keys_device_placement = ["past_key_values"]
- _supports_flash_attn_2 = True
- _supports_sdpa = True
- _supports_cache_class = True
-
- def _init_weights(self, module):
- std = self.config.initializer_range
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
- def _setup_cache(self, cache_cls, max_batch_size, max_cache_len: Optional[int] = None):
- if self.config._attn_implementation == "flash_attention_2" and cache_cls == StaticCache:
- raise ValueError(
- "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "
- "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers"
- )
-
- for layer in self.model.layers:
- device = layer.input_layernorm.weight.device
- if hasattr(self.config, "_pre_quantization_dtype"):
- dtype = self.config._pre_quantization_dtype
- else:
- dtype = layer.self_attn.o_proj.weight.dtype
- layer.self_attn.past_key_value = cache_cls(
- self.config, max_batch_size, max_cache_len, device=device, dtype=dtype
- )
-
- def _reset_cache(self):
- for layer in self.model.layers:
- layer.self_attn.past_key_value = None
-
-
-COHERE_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
- it.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
- `past_key_values`).
-
- If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
- and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
- information on the default strategy.
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
- Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
- blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
- returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
-
- Two formats are allowed:
- - a [`~cache_utils.Cache`] instance;
- - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
- shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
- cache format.
-
- The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
- legacy cache format will be returned.
-
- If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
- have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
- of shape `(batch_size, sequence_length)`.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Cohere Model outputting raw hidden-states without any specific head on top.",
- COHERE_START_DOCSTRING,
-)
-# Copied from transformers.models.llama.modeling_llama.LlamaModel with Llama->Cohere
-class CohereModel(CoherePreTrainedModel):
- """
- Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`CohereDecoderLayer`]
-
- Args:
- config: CohereConfig
- """
-
- # Ignore copy
- def __init__(self, config: CohereConfig):
- super().__init__(config)
- self.padding_idx = config.pad_token_id
- self.vocab_size = config.vocab_size
-
- self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
- self.layers = nn.ModuleList(
- [CohereDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
- )
- self.norm = CohereLayerNorm(hidden_size=(config.hidden_size), eps=config.layer_norm_eps)
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embed_tokens
-
- def set_input_embeddings(self, value):
- self.embed_tokens = value
-
- # Ignore copy
- @add_start_docstrings_to_model_forward(COHERE_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Union[Tuple, BaseModelOutputWithPast]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if (input_ids is None) ^ (inputs_embeds is not None):
- raise ValueError(
- "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
- )
-
- if self.gradient_checkpointing and self.training and use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
- )
- use_cache = False
-
- if inputs_embeds is None:
- inputs_embeds = self.embed_tokens(input_ids)
-
- past_seen_tokens = 0
- if use_cache: # kept for BC (cache positions)
- if not isinstance(past_key_values, StaticCache):
- past_key_values = DynamicCache.from_legacy_cache(past_key_values)
- past_seen_tokens = past_key_values.get_seq_length()
-
- if cache_position is None:
- if isinstance(past_key_values, StaticCache):
- raise ValueError("cache_position is a required argument when using StaticCache.")
- cache_position = torch.arange(
- past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
- )
-
- if position_ids is None:
- position_ids = cache_position.unsqueeze(0)
-
- causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_seen_tokens)
-
- # embed positions
- hidden_states = inputs_embeds
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- next_decoder_cache = None
-
- for decoder_layer in self.layers:
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- decoder_layer.__call__,
- hidden_states,
- causal_mask,
- position_ids,
- past_key_values,
- output_attentions,
- use_cache,
- cache_position,
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- attention_mask=causal_mask,
- position_ids=position_ids,
- past_key_value=past_key_values,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- )
-
- hidden_states = layer_outputs[0]
-
- if use_cache:
- next_decoder_cache = layer_outputs[2 if output_attentions else 1]
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- hidden_states = self.norm(hidden_states)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- next_cache = None
- if use_cache:
- next_cache = (
- next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache
- )
- if not return_dict:
- return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=next_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- )
-
- def _update_causal_mask(
- self,
- attention_mask: torch.Tensor,
- input_tensor: torch.Tensor,
- cache_position: torch.Tensor,
- past_seen_tokens: int,
- ):
- # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static
- # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.
- # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using
- # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114
-
- if self.config._attn_implementation == "flash_attention_2":
- if attention_mask is not None and 0.0 in attention_mask:
- return attention_mask
- return None
-
- if self.config._attn_implementation == "sdpa":
- # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument,
- # in order to dispatch on Flash Attention 2.
- if AttentionMaskConverter._ignore_causal_mask_sdpa(
- attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens
- ):
- return None
-
- dtype, device = input_tensor.dtype, input_tensor.device
- min_dtype = torch.finfo(dtype).min
- sequence_length = input_tensor.shape[1]
- if hasattr(getattr(self.layers[0], "self_attn", {}), "past_key_value"): # static cache
- target_length = self.config.max_position_embeddings
- else: # dynamic cache
- target_length = (
- attention_mask.shape[-1]
- if isinstance(attention_mask, torch.Tensor)
- else past_seen_tokens + sequence_length + 1
- )
-
- causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
- if sequence_length != 1:
- causal_mask = torch.triu(causal_mask, diagonal=1)
- causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
- causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
- if attention_mask is not None:
- causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
- if attention_mask.dim() == 2:
- mask_length = attention_mask.shape[-1]
- padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
- causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype)
- elif attention_mask.dim() == 4:
- # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with
- # cache. In that case, the 4D attention mask attends to the newest tokens only.
- if attention_mask.shape[-2] < cache_position[0] + sequence_length:
- offset = cache_position[0]
- else:
- offset = 0
- mask_shape = attention_mask.shape
- mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype
- causal_mask[
- : mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]
- ] = mask_slice
-
- if (
- self.config._attn_implementation == "sdpa"
- and attention_mask is not None
- and attention_mask.device.type == "cuda"
- ):
- # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
- # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
- # Details: https://github.com/pytorch/pytorch/issues/110213
- causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
-
- return causal_mask
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with Llama->Cohere
-class CohereForCausalLM(CoherePreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- # Ignore copy
- def __init__(self, config):
- super().__init__(config)
- self.model = CohereModel(config)
- self.vocab_size = config.vocab_size
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
- self.logit_scale = config.logit_scale
- self.tie_word_embeddings = config.tie_word_embeddings
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.model.embed_tokens
-
- def set_input_embeddings(self, value):
- self.model.embed_tokens = value
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def set_decoder(self, decoder):
- self.model = decoder
-
- def get_decoder(self):
- return self.model
-
- # Ignore copy
- @add_start_docstrings_to_model_forward(COHERE_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- Args:
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- Returns:
-
- Example:
-
- ```python
- >> from transformers import AutoTokenizer, CohereForCausalLM
-
- >> model = CohereForCausalLM.from_pretrained("CohereForAI/c4ai-command-r-v01")
- >> tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")
-
- >> prompt = "Hey, are you conscious? Can you talk to me?"
- >> inputs = tokenizer(prompt, return_tensors="pt")
-
- >> # Generate
- >> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- )
-
- hidden_states = outputs[0]
- logits = self.lm_head(hidden_states)
- logits = logits * self.logit_scale
- logits = logits.float()
-
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- loss = loss_fct(shift_logits, shift_labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def prepare_inputs_for_generation(
- self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs
- ):
- # With static cache, the `past_key_values` is None
- # TODO joao: standardize interface for the different Cache classes and remove of this if
- has_static_cache = False
- if past_key_values is None:
- past_key_values = getattr(getattr(self.model.layers[0], "self_attn", {}), "past_key_value", None)
- has_static_cache = past_key_values is not None
-
- past_length = 0
- if past_key_values is not None:
- if isinstance(past_key_values, Cache):
- past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
- max_cache_length = (
- torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
- if past_key_values.get_max_length() is not None
- else None
- )
- cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
- # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
- else:
- cache_length = past_length = past_key_values[0][0].shape[2]
- max_cache_length = None
-
- # Keep only the unprocessed tokens:
- # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
- # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
- # input)
- if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
- input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
- # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
- # input_ids based on the past_length.
- elif past_length < input_ids.shape[1]:
- input_ids = input_ids[:, past_length:]
- # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
-
- # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
- if (
- max_cache_length is not None
- and attention_mask is not None
- and cache_length + input_ids.shape[1] > max_cache_length
- ):
- attention_mask = attention_mask[:, -max_cache_length:]
-
- position_ids = kwargs.get("position_ids", None)
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
- # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
- # TODO: use `next_tokens` directly instead.
- model_inputs = {"input_ids": input_ids.contiguous()}
-
- input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
- if cache_position is None:
- cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
- else:
- cache_position = cache_position[-input_length:]
-
- if has_static_cache:
- past_key_values = None
-
- model_inputs.update(
- {
- "position_ids": position_ids,
- "cache_position": cache_position,
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "attention_mask": attention_mask,
- }
- )
- return model_inputs
-
- @staticmethod
- def _reorder_cache(past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
- )
- return reordered_past
diff --git a/transformers/models/cohere/tokenization_cohere_fast.py b/transformers/models/cohere/tokenization_cohere_fast.py
deleted file mode 100644
index e733a6dfd09541c9e536e377fd0846e3e8da63fb..0000000000000000000000000000000000000000
--- a/transformers/models/cohere/tokenization_cohere_fast.py
+++ /dev/null
@@ -1,701 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Cohere team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# This file is based on the tokenization_llama_fast.py file in transformers
-
-import pickle
-from typing import Dict, List, Literal, Union
-
-from tokenizers import processors
-
-from ...pipelines.conversational import Conversation
-from ...tokenization_utils_base import BatchEncoding
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import logging
-from ...utils.versions import require_version
-
-
-require_version("tokenizers>=0.13.3")
-
-logger = logging.get_logger(__name__)
-VOCAB_FILES_NAMES = {"tokenizer_file": "tokenizer.json"}
-
-PRETRAINED_VOCAB_FILES_MAP = {
- "tokenizer_file": {
- "Cohere/Command-nightly": "https://huggingface.co/Cohere/Command-nightly/blob/main/tokenizer.json",
- },
-}
-
-# fmt: off
-DEFAULT_SYSTEM_PROMPT = "You are Command-R, a brilliant, sophisticated, AI-assistant trained to assist human users by providing thorough responses. You are trained by Cohere."
-DEFAULT_RAG_PREAMBLE = """## Task and Context
-You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.
-
-## Style Guide
-Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling."""
-# fmt: on
-
-
-class CohereTokenizerFast(PreTrainedTokenizerFast):
- """
- Construct a Cohere tokenizer. Based on byte-level Byte-Pair-Encoding.
-
- This uses notably ByteFallback and NFC normalization.
-
- ```python
- >>> from transformers import AutoTokenizer
-
- >>> tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")
- >>> tokenizer.encode("Hello this is a test")
- [5, 28339, 2075, 1801, 1671, 3282]
- ```
-
- If you want to change the `bos_token` or the `eos_token`, make sure to specify them when initializing the model, or
- call `tokenizer.update_post_processor()` to make sure that the post-processing is correctly done (otherwise the
- values of the first token and final token of an encoded sequence will not be correct). For more details, checkout
- [post-processors] (https://huggingface.co/docs/tokenizers/api/post-processors) documentation.
-
- You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer, but since
- the model was not pretrained this way, it might yield a decrease in performance.
-
-
-
- When used with `is_split_into_words=True`, this tokenizer needs to be instantiated with `add_prefix_space=True`.
-
-
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`, *optional*):
- Path to the vocabulary file.
- merges_file (`str`, *optional*):
- Path to the merges file.
- tokenizer_file (`str`, *optional*):
- [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that
- contains everything needed to load the tokenizer.
- clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
- Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
- extra spaces.
- unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
- eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|END_OF_TURN_TOKEN|>"`):
- The end of sequence token.
- add_bos_token (`bool`, *optional*, defaults to `True`):
- Whether or not to add an `bos_token` at the start of sequences.
- add_eos_token (`bool`, *optional*, defaults to `False`):
- Whether or not to add an `eos_token` at the end of sequences.
- use_default_system_prompt (`bool`, *optional*, defaults to `False`):
- Whether or not the default system prompt for Cohere tokenizer should be used.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not the tokenizer should automatically add a prefix space
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
- padding_side = "left"
- model_input_names = ["input_ids", "attention_mask"]
- slow_tokenizer_class = None
- # No `max_model_input_sizes`
-
- def __init__(
- self,
- vocab_file=None,
- merges_file=None,
- tokenizer_file=None,
- clean_up_tokenization_spaces=False,
- unk_token="",
- bos_token="",
- eos_token="<|END_OF_TURN_TOKEN|>",
- add_bos_token=True,
- add_eos_token=False,
- use_default_system_prompt=False,
- add_prefix_space=False,
- **kwargs,
- ):
- super().__init__(
- vocab_file=vocab_file,
- merges_file=merges_file,
- tokenizer_file=tokenizer_file,
- clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- unk_token=unk_token,
- bos_token=bos_token,
- eos_token=eos_token,
- add_bos_token=add_bos_token,
- add_eos_token=add_eos_token,
- use_default_system_prompt=use_default_system_prompt,
- add_prefix_space=add_prefix_space,
- **kwargs,
- )
- self._add_bos_token = add_bos_token
- self._add_eos_token = add_eos_token
- self.update_post_processor()
- self.use_default_system_prompt = use_default_system_prompt
- self.vocab_file = vocab_file
- self.grounded_generation_template = kwargs.pop("grounded_generation_template", None)
- self.tool_use_template = kwargs.pop("tool_use_template", None)
-
- # TODO @ArthurZucker this can only work one way for now, to update later-on. Tests should also properly
- # check this as they were green before.
- pre_tok_state = pickle.dumps(self.backend_tokenizer.pre_tokenizer)
- decoder_state = pickle.dumps(self.backend_tokenizer.decoder)
-
- if add_prefix_space:
- pre_tok_state = pre_tok_state.replace(b'"add_prefix_space":false', b'"add_prefix_space": true')
- decoder_state = decoder_state.replace(b'"add_prefix_space":false', b'"add_prefix_space": true')
- self.backend_tokenizer.pre_tokenizer = pickle.loads(pre_tok_state)
- self.backend_tokenizer.decoder = pickle.loads(decoder_state)
-
- self.add_prefix_space = add_prefix_space
-
- def _batch_encode_plus(self, *args, **kwargs) -> BatchEncoding:
- is_split_into_words = kwargs.get("is_split_into_words", False)
- if not (self.add_prefix_space or not is_split_into_words):
- raise Exception(
- f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"
- " pretokenized inputs."
- )
-
- return super()._batch_encode_plus(*args, **kwargs)
-
- def _encode_plus(self, *args, **kwargs) -> BatchEncoding:
- is_split_into_words = kwargs.get("is_split_into_words", False)
-
- if not (self.add_prefix_space or not is_split_into_words):
- raise Exception(
- f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"
- " pretokenized inputs."
- )
-
- return super()._encode_plus(*args, **kwargs)
-
- def update_post_processor(self):
- """
- Updates the underlying post processor with the current `bos_token` and `eos_token`.
- """
- bos = self.bos_token
- bos_token_id = self.bos_token_id
- if bos is None and self.add_bos_token:
- raise ValueError("add_bos_token = True but bos_token = None")
-
- eos = self.eos_token
- eos_token_id = self.eos_token_id
- if eos is None and self.add_eos_token:
- raise ValueError("add_eos_token = True but eos_token = None")
-
- single = f"{(bos+':0 ') if self.add_bos_token else ''}$A:0{(' '+eos+':0') if self.add_eos_token else ''}"
- pair = f"{single}{(' '+bos+':1') if self.add_bos_token else ''} $B:1{(' '+eos+':1') if self.add_eos_token else ''}"
-
- special_tokens = []
- if self.add_bos_token:
- special_tokens.append((bos, bos_token_id))
- if self.add_eos_token:
- special_tokens.append((eos, eos_token_id))
- self._tokenizer.post_processor = processors.TemplateProcessing(
- single=single, pair=pair, special_tokens=special_tokens
- )
-
- @property
- def add_eos_token(self):
- return self._add_eos_token
-
- @property
- def add_bos_token(self):
- return self._add_bos_token
-
- @add_eos_token.setter
- def add_eos_token(self, value):
- self._add_eos_token = value
- self.update_post_processor()
-
- @add_bos_token.setter
- def add_bos_token(self, value):
- self._add_bos_token = value
- self.update_post_processor()
-
- @property
- def default_chat_template(self):
- """
- Cohere Tokenizer uses <|START_OF_TURN_TOKEN|> and <|END_OF_TURN_TOKEN|> to indicate each turn in a chat.
- Additioanlly, to indicate the source of the message, <|USER_TOKEN|>, <|CHATBOT_TOKEN|> and <|SYSTEM_TOKEN|>
- for user, assitant and system messages respectively.
-
- The output should look something like:
- <|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{{ preamble }}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>{{ How are you? }}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>{{ I am doing well! }}<|END_OF_TURN_TOKEN|>
-
- Use add_generation_prompt to add a prompt for the model to generate a response:
- >>> from transformers import AutoTokenizer
- >>> tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")
- >>> messages = [{"role": "user", "content": "Hello, how are you?"}]
- >>> tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
- '<|START_OF_TURN_TOKEN|><|USER_TOKEN|>Hello, how are you?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>'
-
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- default_template = (
- "{{ bos_token }}"
- "{% if messages[0]['role'] == 'system' %}"
- "{% set loop_messages = messages[1:] %}" # Extract system message if it's present
- "{% set system_message = messages[0]['content'] %}"
- "{% elif USE_DEFAULT_PROMPT == true %}"
- "{% set loop_messages = messages %}" # Or use the default system message if the flag is set
- "{% set system_message = 'DEFAULT_SYSTEM_MESSAGE' %}"
- "{% else %}"
- "{% set loop_messages = messages %}"
- "{% set system_message = false %}"
- "{% endif %}"
- "{% if system_message != false %}" # Start with system message
- "{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' + system_message + '<|END_OF_TURN_TOKEN|>' }}"
- "{% endif %}"
- "{% for message in loop_messages %}" # Loop over all non-system messages
- "{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}"
- "{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}"
- "{% endif %}"
- "{% set content = message['content'] %}"
- "{% if message['role'] == 'user' %}" # After all of that, handle messages/roles in a fairly normal way
- "{{ '<|START_OF_TURN_TOKEN|><|USER_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}"
- "{% elif message['role'] == 'assistant' %}"
- "{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}"
- "{% endif %}"
- "{% endfor %}"
- "{% if add_generation_prompt %}"
- "{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' }}"
- "{% endif %}"
- )
- default_template = default_template.replace(
- "USE_DEFAULT_PROMPT", "true" if self.use_default_system_prompt else "false"
- )
- default_message = DEFAULT_SYSTEM_PROMPT.replace("\n", "\\n").replace("'", "\\'")
- default_template = default_template.replace("DEFAULT_SYSTEM_MESSAGE", default_message)
-
- tool_use_template = (
- "{{ bos_token }}"
- "{% if messages[0]['role'] == 'system' %}"
- "{% set loop_messages = messages[1:] %}" # Extract system message if it's present
- "{% set system_message = messages[0]['content'] %}"
- "{% else %}"
- "{% set loop_messages = messages %}"
- "{% set system_message = 'DEFAULT_SYSTEM_MESSAGE' %}"
- "{% endif %}"
- "{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' }}"
- "{{ '# Safety Preamble' }}"
- "{{ '\nThe instructions in this section override those in the task description and style guide sections. Don\\'t answer questions that are harmful or immoral.' }}"
- "{{ '\n\n# System Preamble' }}"
- "{{ '\n## Basic Rules' }}"
- "{{ '\nYou are a powerful conversational AI trained by Cohere to help people. You are augmented by a number of tools, and your job is to use and consume the output of these tools to best help the user. You will see a conversation history between yourself and a user, ending with an utterance from the user. You will then see a specific instruction instructing you what kind of response to generate. When you answer the user\\'s requests, you cite your sources in your answers, according to those instructions.' }}"
- "{{ '\n\n# User Preamble' }}"
- "{{ '\n' + system_message }}"
- "{{'\n\n## Available Tools\nHere is a list of tools that you have available to you:\n\n'}}"
- "{% for tool in tools %}"
- "{% if loop.index0 != 0 %}"
- "{{ '\n\n'}}"
- "{% endif %}"
- "{{'```python\ndef ' + tool.name + '('}}"
- "{% for param_name, param_fields in tool.parameter_definitions.items() %}"
- "{% if loop.index0 != 0 %}"
- "{{ ', '}}"
- "{% endif %}"
- "{{param_name}}: "
- "{% if not param_fields.required %}"
- "{{'Optional[' + param_fields.type + '] = None'}}"
- "{% else %}"
- "{{ param_fields.type }}"
- "{% endif %}"
- "{% endfor %}"
- '{{ \') -> List[Dict]:\n """\'}}'
- "{{ tool.description }}"
- "{% if tool.parameter_definitions|length != 0 %}"
- "{{ '\n\n Args:\n '}}"
- "{% for param_name, param_fields in tool.parameter_definitions.items() %}"
- "{% if loop.index0 != 0 %}"
- "{{ '\n ' }}"
- "{% endif %}"
- "{{ param_name + ' ('}}"
- "{% if not param_fields.required %}"
- "{{'Optional[' + param_fields.type + ']'}}"
- "{% else %}"
- "{{ param_fields.type }}"
- "{% endif %}"
- "{{ '): ' + param_fields.description }}"
- "{% endfor %}"
- "{% endif %}"
- '{{ \'\n """\n pass\n```\' }}'
- "{% endfor %}"
- "{{ '<|END_OF_TURN_TOKEN|>'}}"
- "{% for message in loop_messages %}"
- "{% set content = message['content'] %}"
- "{% if message['role'] == 'user' %}"
- "{{ '<|START_OF_TURN_TOKEN|><|USER_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}"
- "{% elif message['role'] == 'system' %}"
- "{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}"
- "{% elif message['role'] == 'assistant' %}"
- "{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}"
- "{% endif %}"
- "{% endfor %}"
- "{{'<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>Write \\'Action:\\' followed by a json-formatted list of actions that you want to perform in order to produce a good response to the user\\'s last input. You can use any of the supplied tools any number of times, but you should aim to execute the minimum number of necessary actions for the input. You should use the `directly-answer` tool if calling the other tools is unnecessary. The list of actions you want to call should be formatted as a list of json objects, for example:\n```json\n[\n {\n \"tool_name\": title of the tool in the specification,\n \"parameters\": a dict of parameters to input into the tool as they are defined in the specs, or {} if it takes no parameters\n }\n]```<|END_OF_TURN_TOKEN|>'}}"
- "{% if add_generation_prompt %}"
- "{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' }}"
- "{% endif %}"
- )
- default_tool_message = DEFAULT_RAG_PREAMBLE.replace("\n", "\\n").replace("'", "\\'")
- tool_use_template = tool_use_template.replace("DEFAULT_SYSTEM_MESSAGE", default_tool_message)
-
- rag_template = (
- "{{ bos_token }}"
- "{% if messages[0]['role'] == 'system' %}"
- "{% set loop_messages = messages[1:] %}" # Extract system message if it's present
- "{% set system_message = messages[0]['content'] %}"
- "{% else %}"
- "{% set loop_messages = messages %}"
- "{% set system_message = 'DEFAULT_SYSTEM_MESSAGE' %}"
- "{% endif %}"
- "{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' }}"
- "{{ '# Safety Preamble' }}"
- "{{ '\nThe instructions in this section override those in the task description and style guide sections. Don\\'t answer questions that are harmful or immoral.' }}"
- "{{ '\n\n# System Preamble' }}"
- "{{ '\n## Basic Rules' }}"
- "{{ '\nYou are a powerful conversational AI trained by Cohere to help people. You are augmented by a number of tools, and your job is to use and consume the output of these tools to best help the user. You will see a conversation history between yourself and a user, ending with an utterance from the user. You will then see a specific instruction instructing you what kind of response to generate. When you answer the user\\'s requests, you cite your sources in your answers, according to those instructions.' }}"
- "{{ '\n\n# User Preamble' }}"
- "{{ '\n' + system_message }}"
- "{{ '<|END_OF_TURN_TOKEN|>'}}"
- "{% for message in loop_messages %}" # Loop over all non-system messages
- "{% set content = message['content'] %}"
- "{% if message['role'] == 'user' %}" # After all of that, handle messages/roles in a fairly normal way
- "{{ '<|START_OF_TURN_TOKEN|><|USER_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}"
- "{% elif message['role'] == 'system' %}"
- "{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}"
- "{% elif message['role'] == 'assistant' %}"
- "{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}"
- "{% endif %}"
- "{% endfor %}"
- "{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>'}}"
- "{{ '' }}"
- "{% for document in documents %}" # Loop over all non-system messages
- "{{ '\nDocument: ' }}"
- "{{ loop.index0 }}\n"
- "{% for key, value in document.items() %}"
- "{{ key }}: {{value}}\n"
- "{% endfor %}"
- "{% endfor %}"
- "{{ ''}}"
- "{{ '<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' }}"
- "{{ 'Carefully perform the following instructions, in order, starting each with a new line.\n' }}"
- "{{ 'Firstly, Decide which of the retrieved documents are relevant to the user\\'s last input by writing \\'Relevant Documents:\\' followed by comma-separated list of document numbers. If none are relevant, you should instead write \\'None\\'.\n' }}"
- "{{ 'Secondly, Decide which of the retrieved documents contain facts that should be cited in a good answer to the user\\'s last input by writing \\'Cited Documents:\\' followed a comma-separated list of document numbers. If you dont want to cite any of them, you should instead write \\'None\\'.\n' }}"
- "{% if citation_mode=='accurate' %}"
- "{{ 'Thirdly, Write \\'Answer:\\' followed by a response to the user\\'s last input in high quality natural english. Use the retrieved documents to help you. Do not insert any citations or grounding markup.\n' }}"
- "{% endif %}"
- "{{ 'Finally, Write \\'Grounded answer:\\' followed by a response to the user\\'s last input in high quality natural english. Use the symbols and to indicate when a fact comes from a document in the search result, e.g my fact for a fact from document 0.' }}"
- "{{ '<|END_OF_TURN_TOKEN|>' }}"
- "{% if add_generation_prompt %}"
- "{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' }}"
- "{% endif %}"
- )
- default_rag_message = DEFAULT_RAG_PREAMBLE.replace("\n", "\\n").replace("'", "\\'")
- rag_template = rag_template.replace("DEFAULT_SYSTEM_MESSAGE", default_rag_message)
-
- return {"default": default_template, "tool_use": tool_use_template, "rag": rag_template}
-
- def apply_tool_use_template(
- self,
- conversation: Union[List[Dict[str, str]], "Conversation"],
- tools: List[Dict],
- **kwargs,
- ) -> Union[str, List[int]]:
- """Create a Command-R tool-use prompt.
-
- Once rendered, the prompt instructs the model to generate a list of actions to perform on a set of user supplied tools
- to help carry out the user's requests.
-
- Conceptually, this works in the same way as `apply_chat_format`, but takes an additional `tools` parameter.
-
- Converts a Conversation object or a list of dictionaries with `"role"` and `"content"` keys and a list of available
- tools for the model to use into a prompt string, or a list of token ids.
- This method will use the tokenizer's `default_tool_use_template` template specified at the class level.
- You can override the default template using the `tool_use_template` kwarg but the quality of your results may decrease.
-
- Args:
- conversation (Union[List[Dict[str, str]], "Conversation"]): A Conversation object or list of dicts
- with "role" and "content" keys, representing the chat history so far.
- tools (List[Dict]): a list of tools to render into the prompt for the model to choose from.
- See an example at the bottom of the docstring.
- The format should be:
- * name (str): The name of the tool to be called. Valid names contain only the characters a-z,
- A-Z, 0-9, _ and must not begin with a digit.
- * description (str): The description of what the tool does, the model uses the description to
- choose when and how to call the function.
- * parameter_definitions (List[Dict]): The input parameters of the tool. Accepts a dictionary
- where the key is the name of the parameter and the value is the parameter spec.
- Valid parameter names contain only the characters a-z, A-Z, 0-9, _ and must not begin with a digit.
- Parameter specs are as follows:
- * description (str): The description of the parameter.
- * type (str): the type of the parameter - most effective for python builtin data types, such as 'str', 'bool'
- * required: boolean: Denotes whether the parameter is always present (required) or not. Defaults to not required.
- add_generation_prompt (bool, *optional*): Whether to end the prompt with the token(s) that indicate
- the start of an assistant message. This is useful when you want to generate a response from the model.
- Note that this argument will be passed to the chat template, and so it must be supported in the
- template for this argument to have any effect.
- tokenize (`bool`, defaults to `True`):
- Whether to tokenize the output. If `False`, the output will be a string.
- padding (`bool`, defaults to `False`):
- Whether to pad sequences to the maximum length. Has no effect if tokenize is `False`.
- truncation (`bool`, defaults to `False`):
- Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`.
- max_length (`int`, *optional*):
- Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If
- not specified, the tokenizer's `max_length` attribute will be used as a default.
- return_tensors (`str` or [`~utils.TensorType`], *optional*):
- If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable
- values are:
- - `'tf'`: Return TensorFlow `tf.Tensor` objects.
- - `'pt'`: Return PyTorch `torch.Tensor` objects.
- - `'np'`: Return NumPy `np.ndarray` objects.
- - `'jax'`: Return JAX `jnp.ndarray` objects.
- return_dict (`bool`, *optional*, defaults to `False`):
- Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`.
- **tokenizer_kwargs: Additional kwargs to pass to the tokenizer.
-
- Returns:
- `str`: A rendered prompt string.
- or if tokenize=True:
- `List[int]`: A list of token ids representing the tokenized chat so far, including control tokens. This
- output is ready to pass to the model, either directly or via methods like `generate()`.
-
- Examples:
-
- ```python
- >> tokenizer = CohereTokenizerFast.from_pretrained("CohereForAI/c4ai-command-r-v01")
- >> tools = [
- {
- "name": "internet_search",
- "description": "Returns a list of relevant document snippets for a textual query retrieved from the internet",
- "parameter_definitions": {
- "query": {
- "description": "Query to search the internet with",
- "type": "str",
- "required": True
- }
- }
- },
- {
- "name': "directly_answer",
- "description": "Calls a standard (un-augmented) AI chatbot to generate a response given the conversation history",
- "parameter_definitions": {}
- }
- ]
- >> conversation = [
- {"role": "user", "content": "Whats the biggest penguin in the world?"}
- ]
- >> # render the prompt, ready for user to inspect, or for input into the model:
- >> prompt = tokenizer.apply_tool_use_template(conversation, tools=tools, tokenize=False, add_generation_prompt=True)
- >> print(prompt)
- <|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># Safety Preamble
- The instructions in this section override those in the task description and style guide sections. Don't answer questions that are harmful or immoral.
-
- # System Preamble
- ## Basic Rules
- You are a powerful conversational AI trained by Cohere to help people. You are augmented by a number of tools, and your job is to use and consume the output of these tools to best help the user. You will see a conversation history between yourself and a user, ending with an utterance from the user. You will then see a specific instruction instructing you what kind of response to generate. When you answer the user's requests, you cite your sources in your answers, according to those instructions.
-
- # User Preamble
- ## Task and Context
- You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.
-
- ## Style Guide
- Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling.
-
- ## Available Tools
- Here is a list of tools that you have available to you:
-
- \\`\\`\\`python
- def internet_search(query: str) -> List[Dict]:
- \"\"\"Returns a list of relevant document snippets for a textual query retrieved from the internet
-
- Args:
- query (str): Query to search the internet with
- \"\"\"
- pass
- \\`\\`\\`
-
- \\`\\`\\`python
- def directly_answer() -> List[Dict]:
- \"\"\"Calls a standard (un-augmented) AI chatbot to generate a response given the conversation history
- \"\"\"
- pass
- \\`\\`\\`<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>Whats the biggest penguin in the world?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>Write 'Action:' followed by a json-formatted list of actions that you want to perform in order to produce a good response to the user's last input. You can use any of the supplied tools any number of times, but you should aim to execute the minimum number of necessary actions for the input. You should use the `directly-answer` tool if calling the other tools is unnecessary. The list of actions you want to call should be formatted as a list of json objects, for example:
- \\`\\`\\`json
- [
- {
- "tool_name": title of the tool in the specification,
- "parameters": a dict of parameters to input into the tool as they are defined in the specs, or {} if it takes no parameters
- }
- ]\\`\\`\\`<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>
- ```
- >> inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors='pt')
- >> outputs = model.generate(inputs, max_new_tokens=128)
- >> print(tokenizer.decode(outputs[0]))
- Action: ```json
- [
- {
- "tool_name": "internet_search",
- "parameters": {
- "query": "biggest penguin in the world"
- }
- }
- ]
- ```
- """
- return self.apply_chat_template(
- conversation,
- chat_template="tool_use",
- tools=tools,
- **kwargs,
- )
-
- def apply_grounded_generation_template(
- self,
- conversation: Union[List[Dict[str, str]], "Conversation"],
- documents: List[Dict],
- citation_mode: Literal["fast", "accurate"] = "accurate",
- **kwargs,
- ) -> Union[str, List[int]]:
- """Create a Command-R grounded generation (aka RAG) prompt.
-
- Once rendered, the prompt instructs the model to generate a response with citations in, based on supplied documents.
-
- Conceptually, this works in the same way as `apply_chat_format`, but takes additional `documents`
- and parameter `citation_mode` parameters.
-
- Converts a Conversation object or a list of dictionaries with `"role"` and `"content"` keys and a list of
- documents for the model to ground its response on into a prompt string, or a list of token ids.
- This method will use the tokenizer's `grounded_generation_template` template specified at the class level.
- You can override the default template using the `grounded_generation_template` kwarg but the quality of your results may decrease.
-
- Args:
- conversation (Union[List[Dict[str, str]], "Conversation"]): A Conversation object or list of dicts
- with "role" and "content" keys, representing the chat history so far.
- documents (List[Dict[str, str]): A list of dicts, representing documents or tool outputs to ground your
- generation on. A document is a semistructured dict, wiht a string to string mapping. Common fields are
- `url`, `title`, `snippet` etc but should be descriptive of the key. They will get rendered into the prompt.
- citation_mode: either "accurate" (prompt the model to generate an answer first, then rewrite it with citation
- spans in) or "fast", where the prompt instructs the model to generate an answer with citations in directly.
- The former has higher quality citations, the latter requires fewer tokens to be generated.
- add_generation_prompt (bool, *optional*): Whether to end the prompt with the token(s) that indicate
- the start of an assistant message. This is useful when you want to generate a response from the model.
- Note that this argument will be passed to the chat template, and so it must be supported in the
- template for this argument to have any effect.
- tokenize (`bool`, defaults to `True`):
- Whether to tokenize the output. If `False`, the output will be a string.
- padding (`bool`, defaults to `False`):
- Whether to pad sequences to the maximum length. Has no effect if tokenize is `False`.
- truncation (`bool`, defaults to `False`):
- Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`.
- max_length (`int`, *optional*):
- Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If
- not specified, the tokenizer's `max_length` attribute will be used as a default.
- return_tensors (`str` or [`~utils.TensorType`], *optional*):
- If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable
- values are:
- - `'tf'`: Return TensorFlow `tf.Tensor` objects.
- - `'pt'`: Return PyTorch `torch.Tensor` objects.
- - `'np'`: Return NumPy `np.ndarray` objects.
- - `'jax'`: Return JAX `jnp.ndarray` objects.
- return_dict (`bool`, *optional*, defaults to `False`):
- Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`.
- **tokenizer_kwargs: Additional kwargs to pass to the tokenizer.
-
- Returns:
- `str`: A rendered prompt string.
- or if tokenize=True:
- `List[int]`: A list of token ids representing the tokenized chat so far, including control tokens. This
- output is ready to pass to the model, either directly or via methods like `generate()`.
-
- Examples:
-
- ```python
- >> tokenizer = CohereTokenizerFast.from_pretrained('CohereForAI/c4ai-command-r-v01')
-
- >> # define documents:
- >> documents = [
- { "title": "Tall penguins", "text": "Emperor penguins are the tallest." },
- { "title": "Penguin habitats", "text": "Emperor penguins only live in Antarctica."}
- ]
- >> # define a conversation:
- >> conversation = [
- {"role": "user", "content": "Whats the biggest penguin in the world?"}
- ]
- >> # render the prompt, ready for user to inspect, or for input into the model:
- >> grounded_generation_prompt = tokenizer.apply_grounded_generation_template(conversation, documents=documents, tokenize=False, add_generation_prompt=True)
- >> print(grounded_generation_prompt)
- <|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># Safety Preamble
- The instructions in this section override those in the task description and style guide sections. Don't answer questions that are harmful or immoral.
-
- ## Basic Rules
- You are a powerful conversational AI trained by Cohere to help people. You are augmented by a number of tools, and your job is to use and consume the output of these tools to best help the user. You will see a conversation history between yourself and a user, ending with an utterance from the user. You will then see a specific instruction instructing you what kind of response to generate. When you answer the user's requests, you cite your sources in your answers, according to those instructions.
-
- # User Preamble
- ## Task and Context
- You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.
-
- ## Style Guide
- Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling.<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>Whats the biggest penguin in the world?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>
- Document: 0
- title: Tall penguins
- text: Emperor penguins are the tallest.
-
- Document: 1
- title: Penguin habitats
- text: Emperor penguins only live in Antarctica.
- <|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>Carefully perform the following instructions, in order, starting each with a new line.
- Firstly, Decide which of the retrieved documents are relevant to the user's last input by writing 'Relevant Documents:' followed by comma-separated list of document numbers. If none are relevant, you should instead write 'None'.
- Secondly, Decide which of the retrieved documents contain facts that should be cited in a good answer to the user's last input by writing 'Cited Documents:' followed a comma-separated list of document numbers. If you dont want to cite any of them, you should instead write 'None'.
- Thirdly, Write 'Answer:' followed by a response to the user's last input in high quality natural english. Use the retrieved documents to help you. Do not insert any citations or grounding markup.
- Finally, Write 'Grounded answer:' followed by a response to the user's last input in high quality natural english. Use the symbols and to indicate when a fact comes from a document in the search result, e.g my fact for a fact from document 0.<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>'''
- ```
- >> inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors='pt')
- >> outputs = model.generate(inputs, max_new_tokens=128)
- >> print(tokenizer.decode(outputs[0]))
- Relevant Documents: 0,1
- Cited Documents: 0,1
- Answer: The Emperor Penguin is the tallest or biggest penguin in the world. It is a bird that lives only in Antarctica and grows to a height of around 122 centimetres.
- Grounded answer: The Emperor Penguin is the tallest or biggest penguin in the world. It is a bird that lives only in Antarctica and grows to a height of around 122 centimetres.
- """
- return self.apply_chat_template(
- conversation,
- chat_template="rag",
- documents=documents,
- citation_mode=citation_mode,
- **kwargs,
- )
-
- # TODO ArthurZ let's rely on the template processor instead, refactor all fast tokenizers
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- bos_token_id = [self.bos_token_id] if self.add_bos_token else []
- eos_token_id = [self.eos_token_id] if self.add_eos_token else []
-
- output = bos_token_id + token_ids_0 + eos_token_id
-
- if token_ids_1 is not None:
- output = output + bos_token_id + token_ids_1 + eos_token_id
-
- return output
diff --git a/transformers/models/conditional_detr/__init__.py b/transformers/models/conditional_detr/__init__.py
deleted file mode 100644
index 565323321160ff80e3abbd120dd591dcc43d0f6c..0000000000000000000000000000000000000000
--- a/transformers/models/conditional_detr/__init__.py
+++ /dev/null
@@ -1,85 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {
- "configuration_conditional_detr": [
- "CONDITIONAL_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "ConditionalDetrConfig",
- "ConditionalDetrOnnxConfig",
- ]
-}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_conditional_detr"] = ["ConditionalDetrFeatureExtractor"]
- _import_structure["image_processing_conditional_detr"] = ["ConditionalDetrImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_conditional_detr"] = [
- "CONDITIONAL_DETR_PRETRAINED_MODEL_ARCHIVE_LIST",
- "ConditionalDetrForObjectDetection",
- "ConditionalDetrForSegmentation",
- "ConditionalDetrModel",
- "ConditionalDetrPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_conditional_detr import (
- CONDITIONAL_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP,
- ConditionalDetrConfig,
- ConditionalDetrOnnxConfig,
- )
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_conditional_detr import ConditionalDetrFeatureExtractor
- from .image_processing_conditional_detr import ConditionalDetrImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_conditional_detr import (
- CONDITIONAL_DETR_PRETRAINED_MODEL_ARCHIVE_LIST,
- ConditionalDetrForObjectDetection,
- ConditionalDetrForSegmentation,
- ConditionalDetrModel,
- ConditionalDetrPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/conditional_detr/__pycache__/__init__.cpython-310.pyc b/transformers/models/conditional_detr/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 5a90523aa43da88c53d64684184449e1eeb4ae6c..0000000000000000000000000000000000000000
Binary files a/transformers/models/conditional_detr/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/conditional_detr/__pycache__/configuration_conditional_detr.cpython-310.pyc b/transformers/models/conditional_detr/__pycache__/configuration_conditional_detr.cpython-310.pyc
deleted file mode 100644
index 655f71e3d9210db8a0f9d5efe9dce313e7d8b8a7..0000000000000000000000000000000000000000
Binary files a/transformers/models/conditional_detr/__pycache__/configuration_conditional_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/conditional_detr/__pycache__/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/conditional_detr/__pycache__/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index b13077f4f4d5ff3501b7492488dea7eb99f5c2e0..0000000000000000000000000000000000000000
Binary files a/transformers/models/conditional_detr/__pycache__/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/conditional_detr/__pycache__/feature_extraction_conditional_detr.cpython-310.pyc b/transformers/models/conditional_detr/__pycache__/feature_extraction_conditional_detr.cpython-310.pyc
deleted file mode 100644
index 2221d347e2253bfcffaec8ab3bddfa2985900b0c..0000000000000000000000000000000000000000
Binary files a/transformers/models/conditional_detr/__pycache__/feature_extraction_conditional_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/conditional_detr/__pycache__/image_processing_conditional_detr.cpython-310.pyc b/transformers/models/conditional_detr/__pycache__/image_processing_conditional_detr.cpython-310.pyc
deleted file mode 100644
index 21658fcc810f76f3e19dbe581e9e0e0bef09fb48..0000000000000000000000000000000000000000
Binary files a/transformers/models/conditional_detr/__pycache__/image_processing_conditional_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/conditional_detr/__pycache__/modeling_conditional_detr.cpython-310.pyc b/transformers/models/conditional_detr/__pycache__/modeling_conditional_detr.cpython-310.pyc
deleted file mode 100644
index bf09bacffb967d626ed97be316132d7de8298c70..0000000000000000000000000000000000000000
Binary files a/transformers/models/conditional_detr/__pycache__/modeling_conditional_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/conditional_detr/configuration_conditional_detr.py b/transformers/models/conditional_detr/configuration_conditional_detr.py
deleted file mode 100644
index 945e5edb32ad30a801d45e6de65edff1bfd2eae4..0000000000000000000000000000000000000000
--- a/transformers/models/conditional_detr/configuration_conditional_detr.py
+++ /dev/null
@@ -1,273 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Conditional DETR model configuration"""
-from collections import OrderedDict
-from typing import Mapping
-
-from packaging import version
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-from ..auto import CONFIG_MAPPING
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import CONDITIONAL_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class ConditionalDetrConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`ConditionalDetrModel`]. It is used to instantiate
- a Conditional DETR model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the Conditional DETR
- [microsoft/conditional-detr-resnet-50](https://huggingface.co/microsoft/conditional-detr-resnet-50) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- use_timm_backbone (`bool`, *optional*, defaults to `True`):
- Whether or not to use the `timm` library for the backbone. If set to `False`, will use the [`AutoBackbone`]
- API.
- backbone_config (`PretrainedConfig` or `dict`, *optional*):
- The configuration of the backbone model. Only used in case `use_timm_backbone` is set to `False` in which
- case it will default to `ResNetConfig()`.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- num_queries (`int`, *optional*, defaults to 100):
- Number of object queries, i.e. detection slots. This is the maximal number of objects
- [`ConditionalDetrModel`] can detect in a single image. For COCO, we recommend 100 queries.
- d_model (`int`, *optional*, defaults to 256):
- Dimension of the layers.
- encoder_layers (`int`, *optional*, defaults to 6):
- Number of encoder layers.
- decoder_layers (`int`, *optional*, defaults to 6):
- Number of decoder layers.
- encoder_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the Transformer encoder.
- decoder_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the Transformer decoder.
- decoder_ffn_dim (`int`, *optional*, defaults to 2048):
- Dimension of the "intermediate" (often named feed-forward) layer in decoder.
- encoder_ffn_dim (`int`, *optional*, defaults to 2048):
- Dimension of the "intermediate" (often named feed-forward) layer in decoder.
- activation_function (`str` or `function`, *optional*, defaults to `"relu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- activation_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for activations inside the fully connected layer.
- init_std (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- init_xavier_std (`float`, *optional*, defaults to 1):
- The scaling factor used for the Xavier initialization gain in the HM Attention map module.
- encoder_layerdrop (`float`, *optional*, defaults to 0.0):
- The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
- for more details.
- decoder_layerdrop (`float`, *optional*, defaults to 0.0):
- The LayerDrop probability for the decoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
- for more details.
- auxiliary_loss (`bool`, *optional*, defaults to `False`):
- Whether auxiliary decoding losses (loss at each decoder layer) are to be used.
- position_embedding_type (`str`, *optional*, defaults to `"sine"`):
- Type of position embeddings to be used on top of the image features. One of `"sine"` or `"learned"`.
- backbone (`str`, *optional*, defaults to `"resnet50"`):
- Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
- will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
- is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
- use_pretrained_backbone (`bool`, *optional*, defaults to `True`):
- Whether to use pretrained weights for the backbone.
- backbone_kwargs (`dict`, *optional*):
- Keyword arguments to be passed to AutoBackbone when loading from a checkpoint
- e.g. `{'out_indices': (0, 1, 2, 3)}`. Cannot be specified if `backbone_config` is set.
- dilation (`bool`, *optional*, defaults to `False`):
- Whether to replace stride with dilation in the last convolutional block (DC5). Only supported when
- `use_timm_backbone` = `True`.
- class_cost (`float`, *optional*, defaults to 1):
- Relative weight of the classification error in the Hungarian matching cost.
- bbox_cost (`float`, *optional*, defaults to 5):
- Relative weight of the L1 error of the bounding box coordinates in the Hungarian matching cost.
- giou_cost (`float`, *optional*, defaults to 2):
- Relative weight of the generalized IoU loss of the bounding box in the Hungarian matching cost.
- mask_loss_coefficient (`float`, *optional*, defaults to 1):
- Relative weight of the Focal loss in the panoptic segmentation loss.
- dice_loss_coefficient (`float`, *optional*, defaults to 1):
- Relative weight of the DICE/F-1 loss in the panoptic segmentation loss.
- bbox_loss_coefficient (`float`, *optional*, defaults to 5):
- Relative weight of the L1 bounding box loss in the object detection loss.
- giou_loss_coefficient (`float`, *optional*, defaults to 2):
- Relative weight of the generalized IoU loss in the object detection loss.
- eos_coefficient (`float`, *optional*, defaults to 0.1):
- Relative classification weight of the 'no-object' class in the object detection loss.
- focal_alpha (`float`, *optional*, defaults to 0.25):
- Alpha parameter in the focal loss.
-
- Examples:
-
- ```python
- >>> from transformers import ConditionalDetrConfig, ConditionalDetrModel
-
- >>> # Initializing a Conditional DETR microsoft/conditional-detr-resnet-50 style configuration
- >>> configuration = ConditionalDetrConfig()
-
- >>> # Initializing a model (with random weights) from the microsoft/conditional-detr-resnet-50 style configuration
- >>> model = ConditionalDetrModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "conditional_detr"
- keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {
- "hidden_size": "d_model",
- "num_attention_heads": "encoder_attention_heads",
- }
-
- def __init__(
- self,
- use_timm_backbone=True,
- backbone_config=None,
- num_channels=3,
- num_queries=300,
- encoder_layers=6,
- encoder_ffn_dim=2048,
- encoder_attention_heads=8,
- decoder_layers=6,
- decoder_ffn_dim=2048,
- decoder_attention_heads=8,
- encoder_layerdrop=0.0,
- decoder_layerdrop=0.0,
- is_encoder_decoder=True,
- activation_function="relu",
- d_model=256,
- dropout=0.1,
- attention_dropout=0.0,
- activation_dropout=0.0,
- init_std=0.02,
- init_xavier_std=1.0,
- auxiliary_loss=False,
- position_embedding_type="sine",
- backbone="resnet50",
- use_pretrained_backbone=True,
- backbone_kwargs=None,
- dilation=False,
- class_cost=2,
- bbox_cost=5,
- giou_cost=2,
- mask_loss_coefficient=1,
- dice_loss_coefficient=1,
- cls_loss_coefficient=2,
- bbox_loss_coefficient=5,
- giou_loss_coefficient=2,
- focal_alpha=0.25,
- **kwargs,
- ):
- if not use_timm_backbone and use_pretrained_backbone:
- raise ValueError(
- "Loading pretrained backbone weights from the transformers library is not supported yet. `use_timm_backbone` must be set to `True` when `use_pretrained_backbone=True`"
- )
-
- if backbone_config is not None and backbone is not None:
- raise ValueError("You can't specify both `backbone` and `backbone_config`.")
-
- if backbone_config is not None and use_timm_backbone:
- raise ValueError("You can't specify both `backbone_config` and `use_timm_backbone`.")
-
- if backbone_kwargs is not None and backbone_kwargs and backbone_config is not None:
- raise ValueError("You can't specify both `backbone_kwargs` and `backbone_config`.")
-
- if not use_timm_backbone:
- if backbone_config is None:
- logger.info("`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.")
- backbone_config = CONFIG_MAPPING["resnet"](out_features=["stage4"])
- elif isinstance(backbone_config, dict):
- backbone_model_type = backbone_config.get("model_type")
- config_class = CONFIG_MAPPING[backbone_model_type]
- backbone_config = config_class.from_dict(backbone_config)
-
- self.use_timm_backbone = use_timm_backbone
- self.backbone_config = backbone_config
- self.num_channels = num_channels
- self.num_queries = num_queries
- self.d_model = d_model
- self.encoder_ffn_dim = encoder_ffn_dim
- self.encoder_layers = encoder_layers
- self.encoder_attention_heads = encoder_attention_heads
- self.decoder_ffn_dim = decoder_ffn_dim
- self.decoder_layers = decoder_layers
- self.decoder_attention_heads = decoder_attention_heads
- self.dropout = dropout
- self.attention_dropout = attention_dropout
- self.activation_dropout = activation_dropout
- self.activation_function = activation_function
- self.init_std = init_std
- self.init_xavier_std = init_xavier_std
- self.encoder_layerdrop = encoder_layerdrop
- self.decoder_layerdrop = decoder_layerdrop
- self.num_hidden_layers = encoder_layers
- self.auxiliary_loss = auxiliary_loss
- self.position_embedding_type = position_embedding_type
- self.backbone = backbone
- self.use_pretrained_backbone = use_pretrained_backbone
- self.backbone_kwargs = backbone_kwargs
- self.dilation = dilation
- # Hungarian matcher
- self.class_cost = class_cost
- self.bbox_cost = bbox_cost
- self.giou_cost = giou_cost
- # Loss coefficients
- self.mask_loss_coefficient = mask_loss_coefficient
- self.dice_loss_coefficient = dice_loss_coefficient
- self.cls_loss_coefficient = cls_loss_coefficient
- self.bbox_loss_coefficient = bbox_loss_coefficient
- self.giou_loss_coefficient = giou_loss_coefficient
- self.focal_alpha = focal_alpha
- super().__init__(is_encoder_decoder=is_encoder_decoder, **kwargs)
-
- @property
- def num_attention_heads(self) -> int:
- return self.encoder_attention_heads
-
- @property
- def hidden_size(self) -> int:
- return self.d_model
-
-
-class ConditionalDetrOnnxConfig(OnnxConfig):
- torch_onnx_minimum_version = version.parse("1.11")
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- return OrderedDict(
- [
- ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
- ("pixel_mask", {0: "batch"}),
- ]
- )
-
- @property
- def atol_for_validation(self) -> float:
- return 1e-5
-
- @property
- def default_onnx_opset(self) -> int:
- return 12
diff --git a/transformers/models/conditional_detr/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/conditional_detr/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index b1a1b1c817ae702ac8f99513ebbe4c90eefdece6..0000000000000000000000000000000000000000
--- a/transformers/models/conditional_detr/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,325 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Conditional DETR checkpoints."""
-
-
-import argparse
-import json
-from collections import OrderedDict
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from transformers import (
- ConditionalDetrConfig,
- ConditionalDetrForObjectDetection,
- ConditionalDetrForSegmentation,
- ConditionalDetrImageProcessor,
-)
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-rename_keys = []
-for i in range(6):
- # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.self_attn.out_proj.weight", f"encoder.layers.{i}.self_attn.out_proj.weight")
- )
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.self_attn.out_proj.bias", f"encoder.layers.{i}.self_attn.out_proj.bias")
- )
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.weight", f"encoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.bias", f"encoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.weight", f"encoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.bias", f"encoder.layers.{i}.fc2.bias"))
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.norm1.weight", f"encoder.layers.{i}.self_attn_layer_norm.weight")
- )
- rename_keys.append((f"transformer.encoder.layers.{i}.norm1.bias", f"encoder.layers.{i}.self_attn_layer_norm.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.weight", f"encoder.layers.{i}.final_layer_norm.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.bias", f"encoder.layers.{i}.final_layer_norm.bias"))
- # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.self_attn.out_proj.weight", f"decoder.layers.{i}.self_attn.out_proj.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.self_attn.out_proj.bias", f"decoder.layers.{i}.self_attn.out_proj.bias")
- )
- rename_keys.append(
- (
- f"transformer.decoder.layers.{i}.cross_attn.out_proj.weight",
- f"decoder.layers.{i}.encoder_attn.out_proj.weight",
- )
- )
- rename_keys.append(
- (
- f"transformer.decoder.layers.{i}.cross_attn.out_proj.bias",
- f"decoder.layers.{i}.encoder_attn.out_proj.bias",
- )
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.weight", f"decoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.bias", f"decoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.weight", f"decoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.bias", f"decoder.layers.{i}.fc2.bias"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm1.weight", f"decoder.layers.{i}.self_attn_layer_norm.weight")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.norm1.bias", f"decoder.layers.{i}.self_attn_layer_norm.bias"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm2.weight", f"decoder.layers.{i}.encoder_attn_layer_norm.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm2.bias", f"decoder.layers.{i}.encoder_attn_layer_norm.bias")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.weight", f"decoder.layers.{i}.final_layer_norm.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.bias", f"decoder.layers.{i}.final_layer_norm.bias"))
-
- # q, k, v projections in self/cross-attention in decoder for conditional DETR
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.sa_qcontent_proj.weight", f"decoder.layers.{i}.sa_qcontent_proj.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.sa_kcontent_proj.weight", f"decoder.layers.{i}.sa_kcontent_proj.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.sa_qpos_proj.weight", f"decoder.layers.{i}.sa_qpos_proj.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.sa_kpos_proj.weight", f"decoder.layers.{i}.sa_kpos_proj.weight")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.sa_v_proj.weight", f"decoder.layers.{i}.sa_v_proj.weight"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.ca_qcontent_proj.weight", f"decoder.layers.{i}.ca_qcontent_proj.weight")
- )
- # rename_keys.append((f"transformer.decoder.layers.{i}.ca_qpos_proj.weight", f"decoder.layers.{i}.ca_qpos_proj.weight"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.ca_kcontent_proj.weight", f"decoder.layers.{i}.ca_kcontent_proj.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.ca_kpos_proj.weight", f"decoder.layers.{i}.ca_kpos_proj.weight")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.ca_v_proj.weight", f"decoder.layers.{i}.ca_v_proj.weight"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.ca_qpos_sine_proj.weight", f"decoder.layers.{i}.ca_qpos_sine_proj.weight")
- )
-
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.sa_qcontent_proj.bias", f"decoder.layers.{i}.sa_qcontent_proj.bias")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.sa_kcontent_proj.bias", f"decoder.layers.{i}.sa_kcontent_proj.bias")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.sa_qpos_proj.bias", f"decoder.layers.{i}.sa_qpos_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.sa_kpos_proj.bias", f"decoder.layers.{i}.sa_kpos_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.sa_v_proj.bias", f"decoder.layers.{i}.sa_v_proj.bias"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.ca_qcontent_proj.bias", f"decoder.layers.{i}.ca_qcontent_proj.bias")
- )
- # rename_keys.append((f"transformer.decoder.layers.{i}.ca_qpos_proj.bias", f"decoder.layers.{i}.ca_qpos_proj.bias"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.ca_kcontent_proj.bias", f"decoder.layers.{i}.ca_kcontent_proj.bias")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.ca_kpos_proj.bias", f"decoder.layers.{i}.ca_kpos_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.ca_v_proj.bias", f"decoder.layers.{i}.ca_v_proj.bias"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.ca_qpos_sine_proj.bias", f"decoder.layers.{i}.ca_qpos_sine_proj.bias")
- )
-
-# convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
-# for conditional DETR, also convert reference point head and query scale MLP
-rename_keys.extend(
- [
- ("input_proj.weight", "input_projection.weight"),
- ("input_proj.bias", "input_projection.bias"),
- ("query_embed.weight", "query_position_embeddings.weight"),
- ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
- ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
- ("class_embed.weight", "class_labels_classifier.weight"),
- ("class_embed.bias", "class_labels_classifier.bias"),
- ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
- ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
- ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
- ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
- ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
- ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
- ("transformer.decoder.ref_point_head.layers.0.weight", "decoder.ref_point_head.layers.0.weight"),
- ("transformer.decoder.ref_point_head.layers.0.bias", "decoder.ref_point_head.layers.0.bias"),
- ("transformer.decoder.ref_point_head.layers.1.weight", "decoder.ref_point_head.layers.1.weight"),
- ("transformer.decoder.ref_point_head.layers.1.bias", "decoder.ref_point_head.layers.1.bias"),
- ("transformer.decoder.query_scale.layers.0.weight", "decoder.query_scale.layers.0.weight"),
- ("transformer.decoder.query_scale.layers.0.bias", "decoder.query_scale.layers.0.bias"),
- ("transformer.decoder.query_scale.layers.1.weight", "decoder.query_scale.layers.1.weight"),
- ("transformer.decoder.query_scale.layers.1.bias", "decoder.query_scale.layers.1.bias"),
- ("transformer.decoder.layers.0.ca_qpos_proj.weight", "decoder.layers.0.ca_qpos_proj.weight"),
- ("transformer.decoder.layers.0.ca_qpos_proj.bias", "decoder.layers.0.ca_qpos_proj.bias"),
- ]
-)
-
-
-def rename_key(state_dict, old, new):
- val = state_dict.pop(old)
- state_dict[new] = val
-
-
-def rename_backbone_keys(state_dict):
- new_state_dict = OrderedDict()
- for key, value in state_dict.items():
- if "backbone.0.body" in key:
- new_key = key.replace("backbone.0.body", "backbone.conv_encoder.model")
- new_state_dict[new_key] = value
- else:
- new_state_dict[key] = value
-
- return new_state_dict
-
-
-def read_in_q_k_v(state_dict, is_panoptic=False):
- prefix = ""
- if is_panoptic:
- prefix = "conditional_detr."
-
- # first: transformer encoder
- for i in range(6):
- # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"encoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
- state_dict[f"encoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
- state_dict[f"encoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
- state_dict[f"encoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
- state_dict[f"encoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
- state_dict[f"encoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
-
- return im
-
-
-@torch.no_grad()
-def convert_conditional_detr_checkpoint(model_name, pytorch_dump_folder_path):
- """
- Copy/paste/tweak model's weights to our CONDITIONAL_DETR structure.
- """
-
- # load default config
- config = ConditionalDetrConfig()
- # set backbone and dilation attributes
- if "resnet101" in model_name:
- config.backbone = "resnet101"
- if "dc5" in model_name:
- config.dilation = True
- is_panoptic = "panoptic" in model_name
- if is_panoptic:
- config.num_labels = 250
- else:
- config.num_labels = 91
- repo_id = "huggingface/label-files"
- filename = "coco-detection-id2label.json"
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
-
- # load image processor
- format = "coco_panoptic" if is_panoptic else "coco_detection"
- image_processor = ConditionalDetrImageProcessor(format=format)
-
- # prepare image
- img = prepare_img()
- encoding = image_processor(images=img, return_tensors="pt")
- pixel_values = encoding["pixel_values"]
-
- logger.info(f"Converting model {model_name}...")
-
- # load original model from torch hub
- conditional_detr = torch.hub.load("DeppMeng/ConditionalDETR", model_name, pretrained=True).eval()
- state_dict = conditional_detr.state_dict()
- # rename keys
- for src, dest in rename_keys:
- if is_panoptic:
- src = "conditional_detr." + src
- rename_key(state_dict, src, dest)
- state_dict = rename_backbone_keys(state_dict)
- # query, key and value matrices need special treatment
- read_in_q_k_v(state_dict, is_panoptic=is_panoptic)
- # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
- prefix = "conditional_detr.model." if is_panoptic else "model."
- for key in state_dict.copy().keys():
- if is_panoptic:
- if (
- key.startswith("conditional_detr")
- and not key.startswith("class_labels_classifier")
- and not key.startswith("bbox_predictor")
- ):
- val = state_dict.pop(key)
- state_dict["conditional_detr.model" + key[4:]] = val
- elif "class_labels_classifier" in key or "bbox_predictor" in key:
- val = state_dict.pop(key)
- state_dict["conditional_detr." + key] = val
- elif key.startswith("bbox_attention") or key.startswith("mask_head"):
- continue
- else:
- val = state_dict.pop(key)
- state_dict[prefix + key] = val
- else:
- if not key.startswith("class_labels_classifier") and not key.startswith("bbox_predictor"):
- val = state_dict.pop(key)
- state_dict[prefix + key] = val
- # finally, create HuggingFace model and load state dict
- model = ConditionalDetrForSegmentation(config) if is_panoptic else ConditionalDetrForObjectDetection(config)
- model.load_state_dict(state_dict)
- model.eval()
- model.push_to_hub(repo_id=model_name, organization="DepuMeng", commit_message="Add model")
- # verify our conversion
- original_outputs = conditional_detr(pixel_values)
- outputs = model(pixel_values)
- assert torch.allclose(outputs.logits, original_outputs["pred_logits"], atol=1e-4)
- assert torch.allclose(outputs.pred_boxes, original_outputs["pred_boxes"], atol=1e-4)
- if is_panoptic:
- assert torch.allclose(outputs.pred_masks, original_outputs["pred_masks"], atol=1e-4)
-
- # Save model and image processor
- logger.info(f"Saving PyTorch model and image processor to {pytorch_dump_folder_path}...")
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- model.save_pretrained(pytorch_dump_folder_path)
- image_processor.save_pretrained(pytorch_dump_folder_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- "--model_name",
- default="conditional_detr_resnet50",
- type=str,
- help="Name of the CONDITIONAL_DETR model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
- )
- args = parser.parse_args()
- convert_conditional_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path)
diff --git a/transformers/models/conditional_detr/feature_extraction_conditional_detr.py b/transformers/models/conditional_detr/feature_extraction_conditional_detr.py
deleted file mode 100644
index bfdec373f865c5fcbaccfd6b3c906eb690942ddc..0000000000000000000000000000000000000000
--- a/transformers/models/conditional_detr/feature_extraction_conditional_detr.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for Conditional DETR."""
-
-import warnings
-
-from ...image_transforms import rgb_to_id as _rgb_to_id
-from ...utils import logging
-from .image_processing_conditional_detr import ConditionalDetrImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-def rgb_to_id(x):
- warnings.warn(
- "rgb_to_id has moved and will not be importable from this module from v5. "
- "Please import from transformers.image_transforms instead.",
- FutureWarning,
- )
- return _rgb_to_id(x)
-
-
-class ConditionalDetrFeatureExtractor(ConditionalDetrImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class ConditionalDetrFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
- " Please use ConditionalDetrImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/conditional_detr/image_processing_conditional_detr.py b/transformers/models/conditional_detr/image_processing_conditional_detr.py
deleted file mode 100644
index e88bfc8fe230dffeb9d21e4057f15398e90c0222..0000000000000000000000000000000000000000
--- a/transformers/models/conditional_detr/image_processing_conditional_detr.py
+++ /dev/null
@@ -1,1777 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for Conditional DETR."""
-
-import io
-import pathlib
-from collections import defaultdict
-from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
-
-import numpy as np
-
-from ...feature_extraction_utils import BatchFeature
-from ...image_processing_utils import BaseImageProcessor, get_size_dict
-from ...image_transforms import (
- PaddingMode,
- center_to_corners_format,
- corners_to_center_format,
- id_to_rgb,
- pad,
- rescale,
- resize,
- rgb_to_id,
- to_channel_dimension_format,
-)
-from ...image_utils import (
- IMAGENET_DEFAULT_MEAN,
- IMAGENET_DEFAULT_STD,
- AnnotationFormat,
- AnnotationType,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- get_image_size,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_annotations,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import (
- TensorType,
- is_flax_available,
- is_jax_tensor,
- is_scipy_available,
- is_tf_available,
- is_tf_tensor,
- is_torch_available,
- is_torch_tensor,
- is_vision_available,
- logging,
-)
-
-
-if is_torch_available():
- import torch
- from torch import nn
-
-
-if is_vision_available():
- import PIL
-
-
-if is_scipy_available():
- import scipy.special
- import scipy.stats
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-
-SUPPORTED_ANNOTATION_FORMATS = (AnnotationFormat.COCO_DETECTION, AnnotationFormat.COCO_PANOPTIC)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_size_with_aspect_ratio
-def get_size_with_aspect_ratio(image_size, size, max_size=None) -> Tuple[int, int]:
- """
- Computes the output image size given the input image size and the desired output size.
-
- Args:
- image_size (`Tuple[int, int]`):
- The input image size.
- size (`int`):
- The desired output size.
- max_size (`int`, *optional*):
- The maximum allowed output size.
- """
- height, width = image_size
- if max_size is not None:
- min_original_size = float(min((height, width)))
- max_original_size = float(max((height, width)))
- if max_original_size / min_original_size * size > max_size:
- size = int(round(max_size * min_original_size / max_original_size))
-
- if (height <= width and height == size) or (width <= height and width == size):
- return height, width
-
- if width < height:
- ow = size
- oh = int(size * height / width)
- else:
- oh = size
- ow = int(size * width / height)
- return (oh, ow)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_resize_output_image_size
-def get_resize_output_image_size(
- input_image: np.ndarray,
- size: Union[int, Tuple[int, int], List[int]],
- max_size: Optional[int] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
-) -> Tuple[int, int]:
- """
- Computes the output image size given the input image size and the desired output size. If the desired output size
- is a tuple or list, the output image size is returned as is. If the desired output size is an integer, the output
- image size is computed by keeping the aspect ratio of the input image size.
-
- Args:
- input_image (`np.ndarray`):
- The image to resize.
- size (`int` or `Tuple[int, int]` or `List[int]`):
- The desired output size.
- max_size (`int`, *optional*):
- The maximum allowed output size.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred from the input image.
- """
- image_size = get_image_size(input_image, input_data_format)
- if isinstance(size, (list, tuple)):
- return size
-
- return get_size_with_aspect_ratio(image_size, size, max_size)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_numpy_to_framework_fn
-def get_numpy_to_framework_fn(arr) -> Callable:
- """
- Returns a function that converts a numpy array to the framework of the input array.
-
- Args:
- arr (`np.ndarray`): The array to convert.
- """
- if isinstance(arr, np.ndarray):
- return np.array
- if is_tf_available() and is_tf_tensor(arr):
- import tensorflow as tf
-
- return tf.convert_to_tensor
- if is_torch_available() and is_torch_tensor(arr):
- import torch
-
- return torch.tensor
- if is_flax_available() and is_jax_tensor(arr):
- import jax.numpy as jnp
-
- return jnp.array
- raise ValueError(f"Cannot convert arrays of type {type(arr)}")
-
-
-# Copied from transformers.models.detr.image_processing_detr.safe_squeeze
-def safe_squeeze(arr: np.ndarray, axis: Optional[int] = None) -> np.ndarray:
- """
- Squeezes an array, but only if the axis specified has dim 1.
- """
- if axis is None:
- return arr.squeeze()
-
- try:
- return arr.squeeze(axis=axis)
- except ValueError:
- return arr
-
-
-# Copied from transformers.models.detr.image_processing_detr.normalize_annotation
-def normalize_annotation(annotation: Dict, image_size: Tuple[int, int]) -> Dict:
- image_height, image_width = image_size
- norm_annotation = {}
- for key, value in annotation.items():
- if key == "boxes":
- boxes = value
- boxes = corners_to_center_format(boxes)
- boxes /= np.asarray([image_width, image_height, image_width, image_height], dtype=np.float32)
- norm_annotation[key] = boxes
- else:
- norm_annotation[key] = value
- return norm_annotation
-
-
-# Copied from transformers.models.detr.image_processing_detr.max_across_indices
-def max_across_indices(values: Iterable[Any]) -> List[Any]:
- """
- Return the maximum value across all indices of an iterable of values.
- """
- return [max(values_i) for values_i in zip(*values)]
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_max_height_width
-def get_max_height_width(
- images: List[np.ndarray], input_data_format: Optional[Union[str, ChannelDimension]] = None
-) -> List[int]:
- """
- Get the maximum height and width across all images in a batch.
- """
- if input_data_format is None:
- input_data_format = infer_channel_dimension_format(images[0])
-
- if input_data_format == ChannelDimension.FIRST:
- _, max_height, max_width = max_across_indices([img.shape for img in images])
- elif input_data_format == ChannelDimension.LAST:
- max_height, max_width, _ = max_across_indices([img.shape for img in images])
- else:
- raise ValueError(f"Invalid channel dimension format: {input_data_format}")
- return (max_height, max_width)
-
-
-# Copied from transformers.models.detr.image_processing_detr.make_pixel_mask
-def make_pixel_mask(
- image: np.ndarray, output_size: Tuple[int, int], input_data_format: Optional[Union[str, ChannelDimension]] = None
-) -> np.ndarray:
- """
- Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.
-
- Args:
- image (`np.ndarray`):
- Image to make the pixel mask for.
- output_size (`Tuple[int, int]`):
- Output size of the mask.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- mask = np.zeros(output_size, dtype=np.int64)
- mask[:input_height, :input_width] = 1
- return mask
-
-
-# Copied from transformers.models.detr.image_processing_detr.convert_coco_poly_to_mask
-def convert_coco_poly_to_mask(segmentations, height: int, width: int) -> np.ndarray:
- """
- Convert a COCO polygon annotation to a mask.
-
- Args:
- segmentations (`List[List[float]]`):
- List of polygons, each polygon represented by a list of x-y coordinates.
- height (`int`):
- Height of the mask.
- width (`int`):
- Width of the mask.
- """
- try:
- from pycocotools import mask as coco_mask
- except ImportError:
- raise ImportError("Pycocotools is not installed in your environment.")
-
- masks = []
- for polygons in segmentations:
- rles = coco_mask.frPyObjects(polygons, height, width)
- mask = coco_mask.decode(rles)
- if len(mask.shape) < 3:
- mask = mask[..., None]
- mask = np.asarray(mask, dtype=np.uint8)
- mask = np.any(mask, axis=2)
- masks.append(mask)
- if masks:
- masks = np.stack(masks, axis=0)
- else:
- masks = np.zeros((0, height, width), dtype=np.uint8)
-
- return masks
-
-
-# Copied from transformers.models.detr.image_processing_detr.prepare_coco_detection_annotation with DETR->ConditionalDetr
-def prepare_coco_detection_annotation(
- image,
- target,
- return_segmentation_masks: bool = False,
- input_data_format: Optional[Union[ChannelDimension, str]] = None,
-):
- """
- Convert the target in COCO format into the format expected by ConditionalDetr.
- """
- image_height, image_width = get_image_size(image, channel_dim=input_data_format)
-
- image_id = target["image_id"]
- image_id = np.asarray([image_id], dtype=np.int64)
-
- # Get all COCO annotations for the given image.
- annotations = target["annotations"]
- annotations = [obj for obj in annotations if "iscrowd" not in obj or obj["iscrowd"] == 0]
-
- classes = [obj["category_id"] for obj in annotations]
- classes = np.asarray(classes, dtype=np.int64)
-
- # for conversion to coco api
- area = np.asarray([obj["area"] for obj in annotations], dtype=np.float32)
- iscrowd = np.asarray([obj["iscrowd"] if "iscrowd" in obj else 0 for obj in annotations], dtype=np.int64)
-
- boxes = [obj["bbox"] for obj in annotations]
- # guard against no boxes via resizing
- boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4)
- boxes[:, 2:] += boxes[:, :2]
- boxes[:, 0::2] = boxes[:, 0::2].clip(min=0, max=image_width)
- boxes[:, 1::2] = boxes[:, 1::2].clip(min=0, max=image_height)
-
- keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
-
- new_target = {}
- new_target["image_id"] = image_id
- new_target["class_labels"] = classes[keep]
- new_target["boxes"] = boxes[keep]
- new_target["area"] = area[keep]
- new_target["iscrowd"] = iscrowd[keep]
- new_target["orig_size"] = np.asarray([int(image_height), int(image_width)], dtype=np.int64)
-
- if annotations and "keypoints" in annotations[0]:
- keypoints = [obj["keypoints"] for obj in annotations]
- # Converting the filtered keypoints list to a numpy array
- keypoints = np.asarray(keypoints, dtype=np.float32)
- # Apply the keep mask here to filter the relevant annotations
- keypoints = keypoints[keep]
- num_keypoints = keypoints.shape[0]
- keypoints = keypoints.reshape((-1, 3)) if num_keypoints else keypoints
- new_target["keypoints"] = keypoints
-
- if return_segmentation_masks:
- segmentation_masks = [obj["segmentation"] for obj in annotations]
- masks = convert_coco_poly_to_mask(segmentation_masks, image_height, image_width)
- new_target["masks"] = masks[keep]
-
- return new_target
-
-
-# Copied from transformers.models.detr.image_processing_detr.masks_to_boxes
-def masks_to_boxes(masks: np.ndarray) -> np.ndarray:
- """
- Compute the bounding boxes around the provided panoptic segmentation masks.
-
- Args:
- masks: masks in format `[number_masks, height, width]` where N is the number of masks
-
- Returns:
- boxes: bounding boxes in format `[number_masks, 4]` in xyxy format
- """
- if masks.size == 0:
- return np.zeros((0, 4))
-
- h, w = masks.shape[-2:]
- y = np.arange(0, h, dtype=np.float32)
- x = np.arange(0, w, dtype=np.float32)
- # see https://github.com/pytorch/pytorch/issues/50276
- y, x = np.meshgrid(y, x, indexing="ij")
-
- x_mask = masks * np.expand_dims(x, axis=0)
- x_max = x_mask.reshape(x_mask.shape[0], -1).max(-1)
- x = np.ma.array(x_mask, mask=~(np.array(masks, dtype=bool)))
- x_min = x.filled(fill_value=1e8)
- x_min = x_min.reshape(x_min.shape[0], -1).min(-1)
-
- y_mask = masks * np.expand_dims(y, axis=0)
- y_max = y_mask.reshape(x_mask.shape[0], -1).max(-1)
- y = np.ma.array(y_mask, mask=~(np.array(masks, dtype=bool)))
- y_min = y.filled(fill_value=1e8)
- y_min = y_min.reshape(y_min.shape[0], -1).min(-1)
-
- return np.stack([x_min, y_min, x_max, y_max], 1)
-
-
-# Copied from transformers.models.detr.image_processing_detr.prepare_coco_panoptic_annotation with DETR->ConditionalDetr
-def prepare_coco_panoptic_annotation(
- image: np.ndarray,
- target: Dict,
- masks_path: Union[str, pathlib.Path],
- return_masks: bool = True,
- input_data_format: Union[ChannelDimension, str] = None,
-) -> Dict:
- """
- Prepare a coco panoptic annotation for ConditionalDetr.
- """
- image_height, image_width = get_image_size(image, channel_dim=input_data_format)
- annotation_path = pathlib.Path(masks_path) / target["file_name"]
-
- new_target = {}
- new_target["image_id"] = np.asarray([target["image_id"] if "image_id" in target else target["id"]], dtype=np.int64)
- new_target["size"] = np.asarray([image_height, image_width], dtype=np.int64)
- new_target["orig_size"] = np.asarray([image_height, image_width], dtype=np.int64)
-
- if "segments_info" in target:
- masks = np.asarray(PIL.Image.open(annotation_path), dtype=np.uint32)
- masks = rgb_to_id(masks)
-
- ids = np.array([segment_info["id"] for segment_info in target["segments_info"]])
- masks = masks == ids[:, None, None]
- masks = masks.astype(np.uint8)
- if return_masks:
- new_target["masks"] = masks
- new_target["boxes"] = masks_to_boxes(masks)
- new_target["class_labels"] = np.array(
- [segment_info["category_id"] for segment_info in target["segments_info"]], dtype=np.int64
- )
- new_target["iscrowd"] = np.asarray(
- [segment_info["iscrowd"] for segment_info in target["segments_info"]], dtype=np.int64
- )
- new_target["area"] = np.asarray(
- [segment_info["area"] for segment_info in target["segments_info"]], dtype=np.float32
- )
-
- return new_target
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_segmentation_image
-def get_segmentation_image(
- masks: np.ndarray, input_size: Tuple, target_size: Tuple, stuff_equiv_classes, deduplicate=False
-):
- h, w = input_size
- final_h, final_w = target_size
-
- m_id = scipy.special.softmax(masks.transpose(0, 1), -1)
-
- if m_id.shape[-1] == 0:
- # We didn't detect any mask :(
- m_id = np.zeros((h, w), dtype=np.int64)
- else:
- m_id = m_id.argmax(-1).reshape(h, w)
-
- if deduplicate:
- # Merge the masks corresponding to the same stuff class
- for equiv in stuff_equiv_classes.values():
- for eq_id in equiv:
- m_id[m_id == eq_id] = equiv[0]
-
- seg_img = id_to_rgb(m_id)
- seg_img = resize(seg_img, (final_w, final_h), resample=PILImageResampling.NEAREST)
- return seg_img
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_mask_area
-def get_mask_area(seg_img: np.ndarray, target_size: Tuple[int, int], n_classes: int) -> np.ndarray:
- final_h, final_w = target_size
- np_seg_img = seg_img.astype(np.uint8)
- np_seg_img = np_seg_img.reshape(final_h, final_w, 3)
- m_id = rgb_to_id(np_seg_img)
- area = [(m_id == i).sum() for i in range(n_classes)]
- return area
-
-
-# Copied from transformers.models.detr.image_processing_detr.score_labels_from_class_probabilities
-def score_labels_from_class_probabilities(logits: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
- probs = scipy.special.softmax(logits, axis=-1)
- labels = probs.argmax(-1, keepdims=True)
- scores = np.take_along_axis(probs, labels, axis=-1)
- scores, labels = scores.squeeze(-1), labels.squeeze(-1)
- return scores, labels
-
-
-# Copied from transformers.models.detr.image_processing_detr.post_process_panoptic_sample with DetrForSegmentation->ConditionalDetrForSegmentation
-def post_process_panoptic_sample(
- out_logits: np.ndarray,
- masks: np.ndarray,
- boxes: np.ndarray,
- processed_size: Tuple[int, int],
- target_size: Tuple[int, int],
- is_thing_map: Dict,
- threshold=0.85,
-) -> Dict:
- """
- Converts the output of [`ConditionalDetrForSegmentation`] into panoptic segmentation predictions for a single sample.
-
- Args:
- out_logits (`torch.Tensor`):
- The logits for this sample.
- masks (`torch.Tensor`):
- The predicted segmentation masks for this sample.
- boxes (`torch.Tensor`):
- The prediced bounding boxes for this sample. The boxes are in the normalized format `(center_x, center_y,
- width, height)` and values between `[0, 1]`, relative to the size the image (disregarding padding).
- processed_size (`Tuple[int, int]`):
- The processed size of the image `(height, width)`, as returned by the preprocessing step i.e. the size
- after data augmentation but before batching.
- target_size (`Tuple[int, int]`):
- The target size of the image, `(height, width)` corresponding to the requested final size of the
- prediction.
- is_thing_map (`Dict`):
- A dictionary mapping class indices to a boolean value indicating whether the class is a thing or not.
- threshold (`float`, *optional*, defaults to 0.85):
- The threshold used to binarize the segmentation masks.
- """
- # we filter empty queries and detection below threshold
- scores, labels = score_labels_from_class_probabilities(out_logits)
- keep = (labels != out_logits.shape[-1] - 1) & (scores > threshold)
-
- cur_scores = scores[keep]
- cur_classes = labels[keep]
- cur_boxes = center_to_corners_format(boxes[keep])
-
- if len(cur_boxes) != len(cur_classes):
- raise ValueError("Not as many boxes as there are classes")
-
- cur_masks = masks[keep]
- cur_masks = resize(cur_masks[:, None], processed_size, resample=PILImageResampling.BILINEAR)
- cur_masks = safe_squeeze(cur_masks, 1)
- b, h, w = cur_masks.shape
-
- # It may be that we have several predicted masks for the same stuff class.
- # In the following, we track the list of masks ids for each stuff class (they are merged later on)
- cur_masks = cur_masks.reshape(b, -1)
- stuff_equiv_classes = defaultdict(list)
- for k, label in enumerate(cur_classes):
- if not is_thing_map[label]:
- stuff_equiv_classes[label].append(k)
-
- seg_img = get_segmentation_image(cur_masks, processed_size, target_size, stuff_equiv_classes, deduplicate=True)
- area = get_mask_area(cur_masks, processed_size, n_classes=len(cur_scores))
-
- # We filter out any mask that is too small
- if cur_classes.size() > 0:
- # We know filter empty masks as long as we find some
- filtered_small = np.array([a <= 4 for a in area], dtype=bool)
- while filtered_small.any():
- cur_masks = cur_masks[~filtered_small]
- cur_scores = cur_scores[~filtered_small]
- cur_classes = cur_classes[~filtered_small]
- seg_img = get_segmentation_image(cur_masks, (h, w), target_size, stuff_equiv_classes, deduplicate=True)
- area = get_mask_area(seg_img, target_size, n_classes=len(cur_scores))
- filtered_small = np.array([a <= 4 for a in area], dtype=bool)
- else:
- cur_classes = np.ones((1, 1), dtype=np.int64)
-
- segments_info = [
- {"id": i, "isthing": is_thing_map[cat], "category_id": int(cat), "area": a}
- for i, (cat, a) in enumerate(zip(cur_classes, area))
- ]
- del cur_classes
-
- with io.BytesIO() as out:
- PIL.Image.fromarray(seg_img).save(out, format="PNG")
- predictions = {"png_string": out.getvalue(), "segments_info": segments_info}
-
- return predictions
-
-
-# Copied from transformers.models.detr.image_processing_detr.resize_annotation
-def resize_annotation(
- annotation: Dict[str, Any],
- orig_size: Tuple[int, int],
- target_size: Tuple[int, int],
- threshold: float = 0.5,
- resample: PILImageResampling = PILImageResampling.NEAREST,
-):
- """
- Resizes an annotation to a target size.
-
- Args:
- annotation (`Dict[str, Any]`):
- The annotation dictionary.
- orig_size (`Tuple[int, int]`):
- The original size of the input image.
- target_size (`Tuple[int, int]`):
- The target size of the image, as returned by the preprocessing `resize` step.
- threshold (`float`, *optional*, defaults to 0.5):
- The threshold used to binarize the segmentation masks.
- resample (`PILImageResampling`, defaults to `PILImageResampling.NEAREST`):
- The resampling filter to use when resizing the masks.
- """
- ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(target_size, orig_size))
- ratio_height, ratio_width = ratios
-
- new_annotation = {}
- new_annotation["size"] = target_size
-
- for key, value in annotation.items():
- if key == "boxes":
- boxes = value
- scaled_boxes = boxes * np.asarray([ratio_width, ratio_height, ratio_width, ratio_height], dtype=np.float32)
- new_annotation["boxes"] = scaled_boxes
- elif key == "area":
- area = value
- scaled_area = area * (ratio_width * ratio_height)
- new_annotation["area"] = scaled_area
- elif key == "masks":
- masks = value[:, None]
- masks = np.array([resize(mask, target_size, resample=resample) for mask in masks])
- masks = masks.astype(np.float32)
- masks = masks[:, 0] > threshold
- new_annotation["masks"] = masks
- elif key == "size":
- new_annotation["size"] = target_size
- else:
- new_annotation[key] = value
-
- return new_annotation
-
-
-# Copied from transformers.models.detr.image_processing_detr.binary_mask_to_rle
-def binary_mask_to_rle(mask):
- """
- Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format.
-
- Args:
- mask (`torch.Tensor` or `numpy.array`):
- A binary mask tensor of shape `(height, width)` where 0 denotes background and 1 denotes the target
- segment_id or class_id.
- Returns:
- `List`: Run-length encoded list of the binary mask. Refer to COCO API for more information about the RLE
- format.
- """
- if is_torch_tensor(mask):
- mask = mask.numpy()
-
- pixels = mask.flatten()
- pixels = np.concatenate([[0], pixels, [0]])
- runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
- runs[1::2] -= runs[::2]
- return list(runs)
-
-
-# Copied from transformers.models.detr.image_processing_detr.convert_segmentation_to_rle
-def convert_segmentation_to_rle(segmentation):
- """
- Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format.
-
- Args:
- segmentation (`torch.Tensor` or `numpy.array`):
- A segmentation map of shape `(height, width)` where each value denotes a segment or class id.
- Returns:
- `List[List]`: A list of lists, where each list is the run-length encoding of a segment / class id.
- """
- segment_ids = torch.unique(segmentation)
-
- run_length_encodings = []
- for idx in segment_ids:
- mask = torch.where(segmentation == idx, 1, 0)
- rle = binary_mask_to_rle(mask)
- run_length_encodings.append(rle)
-
- return run_length_encodings
-
-
-# Copied from transformers.models.detr.image_processing_detr.remove_low_and_no_objects
-def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels):
- """
- Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and
- `labels`.
-
- Args:
- masks (`torch.Tensor`):
- A tensor of shape `(num_queries, height, width)`.
- scores (`torch.Tensor`):
- A tensor of shape `(num_queries)`.
- labels (`torch.Tensor`):
- A tensor of shape `(num_queries)`.
- object_mask_threshold (`float`):
- A number between 0 and 1 used to binarize the masks.
- Raises:
- `ValueError`: Raised when the first dimension doesn't match in all input tensors.
- Returns:
- `Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region
- < `object_mask_threshold`.
- """
- if not (masks.shape[0] == scores.shape[0] == labels.shape[0]):
- raise ValueError("mask, scores and labels must have the same shape!")
-
- to_keep = labels.ne(num_labels) & (scores > object_mask_threshold)
-
- return masks[to_keep], scores[to_keep], labels[to_keep]
-
-
-# Copied from transformers.models.detr.image_processing_detr.check_segment_validity
-def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8):
- # Get the mask associated with the k class
- mask_k = mask_labels == k
- mask_k_area = mask_k.sum()
-
- # Compute the area of all the stuff in query k
- original_area = (mask_probs[k] >= mask_threshold).sum()
- mask_exists = mask_k_area > 0 and original_area > 0
-
- # Eliminate disconnected tiny segments
- if mask_exists:
- area_ratio = mask_k_area / original_area
- if not area_ratio.item() > overlap_mask_area_threshold:
- mask_exists = False
-
- return mask_exists, mask_k
-
-
-# Copied from transformers.models.detr.image_processing_detr.compute_segments
-def compute_segments(
- mask_probs,
- pred_scores,
- pred_labels,
- mask_threshold: float = 0.5,
- overlap_mask_area_threshold: float = 0.8,
- label_ids_to_fuse: Optional[Set[int]] = None,
- target_size: Tuple[int, int] = None,
-):
- height = mask_probs.shape[1] if target_size is None else target_size[0]
- width = mask_probs.shape[2] if target_size is None else target_size[1]
-
- segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device)
- segments: List[Dict] = []
-
- if target_size is not None:
- mask_probs = nn.functional.interpolate(
- mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False
- )[0]
-
- current_segment_id = 0
-
- # Weigh each mask by its prediction score
- mask_probs *= pred_scores.view(-1, 1, 1)
- mask_labels = mask_probs.argmax(0) # [height, width]
-
- # Keep track of instances of each class
- stuff_memory_list: Dict[str, int] = {}
- for k in range(pred_labels.shape[0]):
- pred_class = pred_labels[k].item()
- should_fuse = pred_class in label_ids_to_fuse
-
- # Check if mask exists and large enough to be a segment
- mask_exists, mask_k = check_segment_validity(
- mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold
- )
-
- if mask_exists:
- if pred_class in stuff_memory_list:
- current_segment_id = stuff_memory_list[pred_class]
- else:
- current_segment_id += 1
-
- # Add current object segment to final segmentation map
- segmentation[mask_k] = current_segment_id
- segment_score = round(pred_scores[k].item(), 6)
- segments.append(
- {
- "id": current_segment_id,
- "label_id": pred_class,
- "was_fused": should_fuse,
- "score": segment_score,
- }
- )
- if should_fuse:
- stuff_memory_list[pred_class] = current_segment_id
-
- return segmentation, segments
-
-
-class ConditionalDetrImageProcessor(BaseImageProcessor):
- r"""
- Constructs a Conditional Detr image processor.
-
- Args:
- format (`str`, *optional*, defaults to `"coco_detection"`):
- Data format of the annotations. One of "coco_detection" or "coco_panoptic".
- do_resize (`bool`, *optional*, defaults to `True`):
- Controls whether to resize the image's (height, width) dimensions to the specified `size`. Can be
- overridden by the `do_resize` parameter in the `preprocess` method.
- size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 800, "longest_edge": 1333}`):
- Size of the image's (height, width) dimensions after resizing. Can be overridden by the `size` parameter in
- the `preprocess` method.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Controls whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the
- `do_rescale` parameter in the `preprocess` method.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the
- `preprocess` method.
- do_normalize:
- Controls whether to normalize the image. Can be overridden by the `do_normalize` parameter in the
- `preprocess` method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_DEFAULT_MEAN`):
- Mean values to use when normalizing the image. Can be a single value or a list of values, one for each
- channel. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_DEFAULT_STD`):
- Standard deviation values to use when normalizing the image. Can be a single value or a list of values, one
- for each channel. Can be overridden by the `image_std` parameter in the `preprocess` method.
- do_convert_annotations (`bool`, *optional*, defaults to `True`):
- Controls whether to convert the annotations to the format expected by the DETR model. Converts the
- bounding boxes to the format `(center_x, center_y, width, height)` and in the range `[0, 1]`.
- Can be overridden by the `do_convert_annotations` parameter in the `preprocess` method.
- do_pad (`bool`, *optional*, defaults to `True`):
- Controls whether to pad the image. Can be overridden by the `do_pad` parameter in the `preprocess`
- method. If `True` will pad the images in the batch to the largest height and width in the batch.
- Padding will be applied to the bottom and right of the image with zeros.
- """
-
- model_input_names = ["pixel_values", "pixel_mask"]
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.__init__
- def __init__(
- self,
- format: Union[str, AnnotationFormat] = AnnotationFormat.COCO_DETECTION,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- do_normalize: bool = True,
- image_mean: Union[float, List[float]] = None,
- image_std: Union[float, List[float]] = None,
- do_convert_annotations: Optional[bool] = None,
- do_pad: bool = True,
- **kwargs,
- ) -> None:
- if "pad_and_return_pixel_mask" in kwargs:
- do_pad = kwargs.pop("pad_and_return_pixel_mask")
-
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` parameter is deprecated and will be removed in v4.26. "
- "Please specify in `size['longest_edge'] instead`.",
- )
- max_size = kwargs.pop("max_size")
- else:
- max_size = None if size is None else 1333
-
- size = size if size is not None else {"shortest_edge": 800, "longest_edge": 1333}
- size = get_size_dict(size, max_size=max_size, default_to_square=False)
-
- # Backwards compatibility
- if do_convert_annotations is None:
- do_convert_annotations = do_normalize
-
- super().__init__(**kwargs)
- self.format = format
- self.do_resize = do_resize
- self.size = size
- self.resample = resample
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_normalize = do_normalize
- self.do_convert_annotations = do_convert_annotations
- self.image_mean = image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
- self.do_pad = do_pad
- self._valid_processor_keys = [
- "images",
- "annotations",
- "return_segmentation_masks",
- "masks_path",
- "do_resize",
- "size",
- "resample",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "do_convert_annotations",
- "image_mean",
- "image_std",
- "do_pad",
- "format",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- @classmethod
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.from_dict with Detr->ConditionalDetr
- def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
- """
- Overrides the `from_dict` method from the base class to make sure parameters are updated if image processor is
- created using from_dict and kwargs e.g. `ConditionalDetrImageProcessor.from_pretrained(checkpoint, size=600,
- max_size=800)`
- """
- image_processor_dict = image_processor_dict.copy()
- if "max_size" in kwargs:
- image_processor_dict["max_size"] = kwargs.pop("max_size")
- if "pad_and_return_pixel_mask" in kwargs:
- image_processor_dict["pad_and_return_pixel_mask"] = kwargs.pop("pad_and_return_pixel_mask")
- return super().from_dict(image_processor_dict, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_annotation with DETR->ConditionalDetr
- def prepare_annotation(
- self,
- image: np.ndarray,
- target: Dict,
- format: Optional[AnnotationFormat] = None,
- return_segmentation_masks: bool = None,
- masks_path: Optional[Union[str, pathlib.Path]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> Dict:
- """
- Prepare an annotation for feeding into ConditionalDetr model.
- """
- format = format if format is not None else self.format
-
- if format == AnnotationFormat.COCO_DETECTION:
- return_segmentation_masks = False if return_segmentation_masks is None else return_segmentation_masks
- target = prepare_coco_detection_annotation(
- image, target, return_segmentation_masks, input_data_format=input_data_format
- )
- elif format == AnnotationFormat.COCO_PANOPTIC:
- return_segmentation_masks = True if return_segmentation_masks is None else return_segmentation_masks
- target = prepare_coco_panoptic_annotation(
- image,
- target,
- masks_path=masks_path,
- return_masks=return_segmentation_masks,
- input_data_format=input_data_format,
- )
- else:
- raise ValueError(f"Format {format} is not supported.")
- return target
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare
- def prepare(self, image, target, return_segmentation_masks=None, masks_path=None):
- logger.warning_once(
- "The `prepare` method is deprecated and will be removed in a v4.33. "
- "Please use `prepare_annotation` instead. Note: the `prepare_annotation` method "
- "does not return the image anymore.",
- )
- target = self.prepare_annotation(image, target, return_segmentation_masks, masks_path, self.format)
- return image, target
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.convert_coco_poly_to_mask
- def convert_coco_poly_to_mask(self, *args, **kwargs):
- logger.warning_once("The `convert_coco_poly_to_mask` method is deprecated and will be removed in v4.33. ")
- return convert_coco_poly_to_mask(*args, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_coco_detection with DETR->ConditionalDetr
- def prepare_coco_detection(self, *args, **kwargs):
- logger.warning_once("The `prepare_coco_detection` method is deprecated and will be removed in v4.33. ")
- return prepare_coco_detection_annotation(*args, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_coco_panoptic
- def prepare_coco_panoptic(self, *args, **kwargs):
- logger.warning_once("The `prepare_coco_panoptic` method is deprecated and will be removed in v4.33. ")
- return prepare_coco_panoptic_annotation(*args, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.resize
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize the image to the given size. Size can be `min_size` (scalar) or `(height, width)` tuple. If size is an
- int, smaller edge of the image will be matched to this number.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary containing the size to resize to. Can contain the keys `shortest_edge` and `longest_edge` or
- `height` and `width`.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` parameter is deprecated and will be removed in v4.26. "
- "Please specify in `size['longest_edge'] instead`.",
- )
- max_size = kwargs.pop("max_size")
- else:
- max_size = None
- size = get_size_dict(size, max_size=max_size, default_to_square=False)
- if "shortest_edge" in size and "longest_edge" in size:
- size = get_resize_output_image_size(
- image, size["shortest_edge"], size["longest_edge"], input_data_format=input_data_format
- )
- elif "height" in size and "width" in size:
- size = (size["height"], size["width"])
- else:
- raise ValueError(
- "Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got"
- f" {size.keys()}."
- )
- image = resize(
- image, size=size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs
- )
- return image
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.resize_annotation
- def resize_annotation(
- self,
- annotation,
- orig_size,
- size,
- resample: PILImageResampling = PILImageResampling.NEAREST,
- ) -> Dict:
- """
- Resize the annotation to match the resized image. If size is an int, smaller edge of the mask will be matched
- to this number.
- """
- return resize_annotation(annotation, orig_size=orig_size, target_size=size, resample=resample)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.rescale
- def rescale(
- self,
- image: np.ndarray,
- rescale_factor: float,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.ndarray:
- """
- Rescale the image by the given factor. image = image * rescale_factor.
-
- Args:
- image (`np.ndarray`):
- Image to rescale.
- rescale_factor (`float`):
- The value to use for rescaling.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the input image. If unset, is inferred from the input image. Can be
- one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- """
- return rescale(image, rescale_factor, data_format=data_format, input_data_format=input_data_format)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.normalize_annotation
- def normalize_annotation(self, annotation: Dict, image_size: Tuple[int, int]) -> Dict:
- """
- Normalize the boxes in the annotation from `[top_left_x, top_left_y, bottom_right_x, bottom_right_y]` to
- `[center_x, center_y, width, height]` format and from absolute to relative pixel values.
- """
- return normalize_annotation(annotation, image_size=image_size)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor._update_annotation_for_padded_image
- def _update_annotation_for_padded_image(
- self,
- annotation: Dict,
- input_image_size: Tuple[int, int],
- output_image_size: Tuple[int, int],
- padding,
- update_bboxes,
- ) -> Dict:
- """
- Update the annotation for a padded image.
- """
- new_annotation = {}
- new_annotation["size"] = output_image_size
-
- for key, value in annotation.items():
- if key == "masks":
- masks = value
- masks = pad(
- masks,
- padding,
- mode=PaddingMode.CONSTANT,
- constant_values=0,
- input_data_format=ChannelDimension.FIRST,
- )
- masks = safe_squeeze(masks, 1)
- new_annotation["masks"] = masks
- elif key == "boxes" and update_bboxes:
- boxes = value
- boxes *= np.asarray(
- [
- input_image_size[1] / output_image_size[1],
- input_image_size[0] / output_image_size[0],
- input_image_size[1] / output_image_size[1],
- input_image_size[0] / output_image_size[0],
- ]
- )
- new_annotation["boxes"] = boxes
- elif key == "size":
- new_annotation["size"] = output_image_size
- else:
- new_annotation[key] = value
- return new_annotation
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor._pad_image
- def _pad_image(
- self,
- image: np.ndarray,
- output_size: Tuple[int, int],
- annotation: Optional[Dict[str, Any]] = None,
- constant_values: Union[float, Iterable[float]] = 0,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- update_bboxes: bool = True,
- ) -> np.ndarray:
- """
- Pad an image with zeros to the given size.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- output_height, output_width = output_size
-
- pad_bottom = output_height - input_height
- pad_right = output_width - input_width
- padding = ((0, pad_bottom), (0, pad_right))
- padded_image = pad(
- image,
- padding,
- mode=PaddingMode.CONSTANT,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- )
- if annotation is not None:
- annotation = self._update_annotation_for_padded_image(
- annotation, (input_height, input_width), (output_height, output_width), padding, update_bboxes
- )
- return padded_image, annotation
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.pad
- def pad(
- self,
- images: List[np.ndarray],
- annotations: Optional[Union[AnnotationType, List[AnnotationType]]] = None,
- constant_values: Union[float, Iterable[float]] = 0,
- return_pixel_mask: bool = True,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- update_bboxes: bool = True,
- ) -> BatchFeature:
- """
- Pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width
- in the batch and optionally returns their corresponding pixel mask.
-
- Args:
- images (List[`np.ndarray`]):
- Images to pad.
- annotations (`AnnotationType` or `List[AnnotationType]`, *optional*):
- Annotations to transform according to the padding that is applied to the images.
- constant_values (`float` or `Iterable[float]`, *optional*):
- The value to use for the padding if `mode` is `"constant"`.
- return_pixel_mask (`bool`, *optional*, defaults to `True`):
- Whether to return a pixel mask.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the image. If not provided, it will be the same as the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- update_bboxes (`bool`, *optional*, defaults to `True`):
- Whether to update the bounding boxes in the annotations to match the padded images. If the
- bounding boxes have not been converted to relative coordinates and `(centre_x, centre_y, width, height)`
- format, the bounding boxes will not be updated.
- """
- pad_size = get_max_height_width(images, input_data_format=input_data_format)
-
- annotation_list = annotations if annotations is not None else [None] * len(images)
- padded_images = []
- padded_annotations = []
- for image, annotation in zip(images, annotation_list):
- padded_image, padded_annotation = self._pad_image(
- image,
- pad_size,
- annotation,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- update_bboxes=update_bboxes,
- )
- padded_images.append(padded_image)
- padded_annotations.append(padded_annotation)
-
- data = {"pixel_values": padded_images}
-
- if return_pixel_mask:
- masks = [
- make_pixel_mask(image=image, output_size=pad_size, input_data_format=input_data_format)
- for image in images
- ]
- data["pixel_mask"] = masks
-
- encoded_inputs = BatchFeature(data=data, tensor_type=return_tensors)
-
- if annotations is not None:
- encoded_inputs["labels"] = [
- BatchFeature(annotation, tensor_type=return_tensors) for annotation in padded_annotations
- ]
-
- return encoded_inputs
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.preprocess
- def preprocess(
- self,
- images: ImageInput,
- annotations: Optional[Union[AnnotationType, List[AnnotationType]]] = None,
- return_segmentation_masks: bool = None,
- masks_path: Optional[Union[str, pathlib.Path]] = None,
- do_resize: Optional[bool] = None,
- size: Optional[Dict[str, int]] = None,
- resample=None, # PILImageResampling
- do_rescale: Optional[bool] = None,
- rescale_factor: Optional[Union[int, float]] = None,
- do_normalize: Optional[bool] = None,
- do_convert_annotations: Optional[bool] = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- do_pad: Optional[bool] = None,
- format: Optional[Union[str, AnnotationFormat]] = None,
- return_tensors: Optional[Union[TensorType, str]] = None,
- data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> BatchFeature:
- """
- Preprocess an image or a batch of images so that it can be used by the model.
-
- Args:
- images (`ImageInput`):
- Image or batch of images to preprocess. Expects a single or batch of images with pixel values ranging
- from 0 to 255. If passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- annotations (`AnnotationType` or `List[AnnotationType]`, *optional*):
- List of annotations associated with the image or batch of images. If annotation is for object
- detection, the annotations should be a dictionary with the following keys:
- - "image_id" (`int`): The image id.
- - "annotations" (`List[Dict]`): List of annotations for an image. Each annotation should be a
- dictionary. An image can have no annotations, in which case the list should be empty.
- If annotation is for segmentation, the annotations should be a dictionary with the following keys:
- - "image_id" (`int`): The image id.
- - "segments_info" (`List[Dict]`): List of segments for an image. Each segment should be a dictionary.
- An image can have no segments, in which case the list should be empty.
- - "file_name" (`str`): The file name of the image.
- return_segmentation_masks (`bool`, *optional*, defaults to self.return_segmentation_masks):
- Whether to return segmentation masks.
- masks_path (`str` or `pathlib.Path`, *optional*):
- Path to the directory containing the segmentation masks.
- do_resize (`bool`, *optional*, defaults to self.do_resize):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to self.size):
- Size of the image after resizing.
- resample (`PILImageResampling`, *optional*, defaults to self.resample):
- Resampling filter to use when resizing the image.
- do_rescale (`bool`, *optional*, defaults to self.do_rescale):
- Whether to rescale the image.
- rescale_factor (`float`, *optional*, defaults to self.rescale_factor):
- Rescale factor to use when rescaling the image.
- do_normalize (`bool`, *optional*, defaults to self.do_normalize):
- Whether to normalize the image.
- do_convert_annotations (`bool`, *optional*, defaults to self.do_convert_annotations):
- Whether to convert the annotations to the format expected by the model. Converts the bounding
- boxes from the format `(top_left_x, top_left_y, width, height)` to `(center_x, center_y, width, height)`
- and in relative coordinates.
- image_mean (`float` or `List[float]`, *optional*, defaults to self.image_mean):
- Mean to use when normalizing the image.
- image_std (`float` or `List[float]`, *optional*, defaults to self.image_std):
- Standard deviation to use when normalizing the image.
- do_pad (`bool`, *optional*, defaults to self.do_pad):
- Whether to pad the image. If `True` will pad the images in the batch to the largest image in the batch
- and create a pixel mask. Padding will be applied to the bottom and right of the image with zeros.
- format (`str` or `AnnotationFormat`, *optional*, defaults to self.format):
- Format of the annotations.
- return_tensors (`str` or `TensorType`, *optional*, defaults to self.return_tensors):
- Type of tensors to return. If `None`, will return the list of images.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - Unset: Use the channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- if "pad_and_return_pixel_mask" in kwargs:
- logger.warning_once(
- "The `pad_and_return_pixel_mask` argument is deprecated and will be removed in a future version, "
- "use `do_pad` instead."
- )
- do_pad = kwargs.pop("pad_and_return_pixel_mask")
-
- max_size = None
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` argument is deprecated and will be removed in a future version, use"
- " `size['longest_edge']` instead."
- )
- size = kwargs.pop("max_size")
-
- do_resize = self.do_resize if do_resize is None else do_resize
- size = self.size if size is None else size
- size = get_size_dict(size=size, max_size=max_size, default_to_square=False)
- resample = self.resample if resample is None else resample
- do_rescale = self.do_rescale if do_rescale is None else do_rescale
- rescale_factor = self.rescale_factor if rescale_factor is None else rescale_factor
- do_normalize = self.do_normalize if do_normalize is None else do_normalize
- image_mean = self.image_mean if image_mean is None else image_mean
- image_std = self.image_std if image_std is None else image_std
- do_convert_annotations = (
- self.do_convert_annotations if do_convert_annotations is None else do_convert_annotations
- )
- do_pad = self.do_pad if do_pad is None else do_pad
- format = self.format if format is None else format
-
- images = make_list_of_images(images)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- # Here, the pad() method pads to the maximum of (width, height). It does not need to be validated.
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
-
- if annotations is not None and isinstance(annotations, dict):
- annotations = [annotations]
-
- if annotations is not None and len(images) != len(annotations):
- raise ValueError(
- f"The number of images ({len(images)}) and annotations ({len(annotations)}) do not match."
- )
-
- format = AnnotationFormat(format)
- if annotations is not None:
- validate_annotations(format, SUPPORTED_ANNOTATION_FORMATS, annotations)
-
- if (
- masks_path is not None
- and format == AnnotationFormat.COCO_PANOPTIC
- and not isinstance(masks_path, (pathlib.Path, str))
- ):
- raise ValueError(
- "The path to the directory containing the mask PNG files should be provided as a"
- f" `pathlib.Path` or string object, but is {type(masks_path)} instead."
- )
-
- # All transformations expect numpy arrays
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- # prepare (COCO annotations as a list of Dict -> DETR target as a single Dict per image)
- if annotations is not None:
- prepared_images = []
- prepared_annotations = []
- for image, target in zip(images, annotations):
- target = self.prepare_annotation(
- image,
- target,
- format,
- return_segmentation_masks=return_segmentation_masks,
- masks_path=masks_path,
- input_data_format=input_data_format,
- )
- prepared_images.append(image)
- prepared_annotations.append(target)
- images = prepared_images
- annotations = prepared_annotations
- del prepared_images, prepared_annotations
-
- # transformations
- if do_resize:
- if annotations is not None:
- resized_images, resized_annotations = [], []
- for image, target in zip(images, annotations):
- orig_size = get_image_size(image, input_data_format)
- resized_image = self.resize(
- image, size=size, max_size=max_size, resample=resample, input_data_format=input_data_format
- )
- resized_annotation = self.resize_annotation(
- target, orig_size, get_image_size(resized_image, input_data_format)
- )
- resized_images.append(resized_image)
- resized_annotations.append(resized_annotation)
- images = resized_images
- annotations = resized_annotations
- del resized_images, resized_annotations
- else:
- images = [
- self.resize(image, size=size, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_rescale:
- images = [self.rescale(image, rescale_factor, input_data_format=input_data_format) for image in images]
-
- if do_normalize:
- images = [
- self.normalize(image, image_mean, image_std, input_data_format=input_data_format) for image in images
- ]
-
- if do_convert_annotations and annotations is not None:
- annotations = [
- self.normalize_annotation(annotation, get_image_size(image, input_data_format))
- for annotation, image in zip(annotations, images)
- ]
-
- if do_pad:
- # Pads images and returns their mask: {'pixel_values': ..., 'pixel_mask': ...}
- encoded_inputs = self.pad(
- images,
- annotations=annotations,
- return_pixel_mask=True,
- data_format=data_format,
- input_data_format=input_data_format,
- update_bboxes=do_convert_annotations,
- return_tensors=return_tensors,
- )
- else:
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
- for image in images
- ]
- encoded_inputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)
- if annotations is not None:
- encoded_inputs["labels"] = [
- BatchFeature(annotation, tensor_type=return_tensors) for annotation in annotations
- ]
-
- return encoded_inputs
-
- # POSTPROCESSING METHODS - TODO: add support for other frameworks
- def post_process(self, outputs, target_sizes):
- """
- Converts the output of [`ConditionalDetrForObjectDetection`] into the format expected by the Pascal VOC format (xmin, ymin, xmax, ymax).
- Only supports PyTorch.
-
- Args:
- outputs ([`ConditionalDetrObjectDetectionOutput`]):
- Raw outputs of the model.
- target_sizes (`torch.Tensor` of shape `(batch_size, 2)`):
- Tensor containing the size (h, w) of each image of the batch. For evaluation, this must be the original
- image size (before any data augmentation). For visualization, this should be the image size after data
- augment, but before padding.
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
- in the batch as predicted by the model.
- """
- logging.warning_once(
- "`post_process` is deprecated and will be removed in v5 of Transformers, please use"
- " `post_process_object_detection` instead, with `threshold=0.` for equivalent results.",
- )
-
- out_logits, out_bbox = outputs.logits, outputs.pred_boxes
-
- if len(out_logits) != len(target_sizes):
- raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the logits")
- if target_sizes.shape[1] != 2:
- raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch")
-
- prob = out_logits.sigmoid()
- topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), 300, dim=1)
- scores = topk_values
- topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode="floor")
- labels = topk_indexes % out_logits.shape[2]
- boxes = center_to_corners_format(out_bbox)
- boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
-
- # and from relative [0, 1] to absolute [0, height] coordinates
- img_h, img_w = target_sizes.unbind(1)
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
- boxes = boxes * scale_fct[:, None, :]
-
- results = [{"scores": s, "labels": l, "boxes": b} for s, l, b in zip(scores, labels, boxes)]
-
- return results
-
- # Copied from transformers.models.deformable_detr.image_processing_deformable_detr.DeformableDetrImageProcessor.post_process_object_detection with DeformableDetr->ConditionalDetr
- def post_process_object_detection(
- self, outputs, threshold: float = 0.5, target_sizes: Union[TensorType, List[Tuple]] = None, top_k: int = 100
- ):
- """
- Converts the raw output of [`ConditionalDetrForObjectDetection`] into final bounding boxes in (top_left_x,
- top_left_y, bottom_right_x, bottom_right_y) format. Only supports PyTorch.
-
- Args:
- outputs ([`DetrObjectDetectionOutput`]):
- Raw outputs of the model.
- threshold (`float`, *optional*):
- Score threshold to keep object detection predictions.
- target_sizes (`torch.Tensor` or `List[Tuple[int, int]]`, *optional*):
- Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
- (height, width) of each image in the batch. If left to None, predictions will not be resized.
- top_k (`int`, *optional*, defaults to 100):
- Keep only top k bounding boxes before filtering by thresholding.
-
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
- in the batch as predicted by the model.
- """
- out_logits, out_bbox = outputs.logits, outputs.pred_boxes
-
- if target_sizes is not None:
- if len(out_logits) != len(target_sizes):
- raise ValueError(
- "Make sure that you pass in as many target sizes as the batch dimension of the logits"
- )
-
- prob = out_logits.sigmoid()
- prob = prob.view(out_logits.shape[0], -1)
- k_value = min(top_k, prob.size(1))
- topk_values, topk_indexes = torch.topk(prob, k_value, dim=1)
- scores = topk_values
- topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode="floor")
- labels = topk_indexes % out_logits.shape[2]
- boxes = center_to_corners_format(out_bbox)
- boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
-
- # and from relative [0, 1] to absolute [0, height] coordinates
- if target_sizes is not None:
- if isinstance(target_sizes, List):
- img_h = torch.Tensor([i[0] for i in target_sizes])
- img_w = torch.Tensor([i[1] for i in target_sizes])
- else:
- img_h, img_w = target_sizes.unbind(1)
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device)
- boxes = boxes * scale_fct[:, None, :]
-
- results = []
- for s, l, b in zip(scores, labels, boxes):
- score = s[s > threshold]
- label = l[s > threshold]
- box = b[s > threshold]
- results.append({"scores": score, "labels": label, "boxes": box})
-
- return results
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.post_process_semantic_segmentation with Detr->ConditionalDetr
- def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple[int, int]] = None):
- """
- Converts the output of [`ConditionalDetrForSegmentation`] into semantic segmentation maps. Only supports PyTorch.
-
- Args:
- outputs ([`ConditionalDetrForSegmentation`]):
- Raw outputs of the model.
- target_sizes (`List[Tuple[int, int]]`, *optional*):
- A list of tuples (`Tuple[int, int]`) containing the target size (height, width) of each image in the
- batch. If unset, predictions will not be resized.
- Returns:
- `List[torch.Tensor]`:
- A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width)
- corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each
- `torch.Tensor` correspond to a semantic class id.
- """
- class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1]
- masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width]
-
- # Remove the null class `[..., :-1]`
- masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
- masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width]
-
- # Semantic segmentation logits of shape (batch_size, num_classes, height, width)
- segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
- batch_size = class_queries_logits.shape[0]
-
- # Resize logits and compute semantic segmentation maps
- if target_sizes is not None:
- if batch_size != len(target_sizes):
- raise ValueError(
- "Make sure that you pass in as many target sizes as the batch dimension of the logits"
- )
-
- semantic_segmentation = []
- for idx in range(batch_size):
- resized_logits = nn.functional.interpolate(
- segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
- )
- semantic_map = resized_logits[0].argmax(dim=0)
- semantic_segmentation.append(semantic_map)
- else:
- semantic_segmentation = segmentation.argmax(dim=1)
- semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]
-
- return semantic_segmentation
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.post_process_instance_segmentation with Detr->ConditionalDetr
- def post_process_instance_segmentation(
- self,
- outputs,
- threshold: float = 0.5,
- mask_threshold: float = 0.5,
- overlap_mask_area_threshold: float = 0.8,
- target_sizes: Optional[List[Tuple[int, int]]] = None,
- return_coco_annotation: Optional[bool] = False,
- ) -> List[Dict]:
- """
- Converts the output of [`ConditionalDetrForSegmentation`] into instance segmentation predictions. Only supports PyTorch.
-
- Args:
- outputs ([`ConditionalDetrForSegmentation`]):
- Raw outputs of the model.
- threshold (`float`, *optional*, defaults to 0.5):
- The probability score threshold to keep predicted instance masks.
- mask_threshold (`float`, *optional*, defaults to 0.5):
- Threshold to use when turning the predicted masks into binary values.
- overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8):
- The overlap mask area threshold to merge or discard small disconnected parts within each binary
- instance mask.
- target_sizes (`List[Tuple]`, *optional*):
- List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested
- final size (height, width) of each prediction. If unset, predictions will not be resized.
- return_coco_annotation (`bool`, *optional*):
- Defaults to `False`. If set to `True`, segmentation maps are returned in COCO run-length encoding (RLE)
- format.
- Returns:
- `List[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys:
- - **segmentation** -- A tensor of shape `(height, width)` where each pixel represents a `segment_id` or
- `List[List]` run-length encoding (RLE) of the segmentation map if return_coco_annotation is set to
- `True`. Set to `None` if no mask if found above `threshold`.
- - **segments_info** -- A dictionary that contains additional information on each segment.
- - **id** -- An integer representing the `segment_id`.
- - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`.
- - **score** -- Prediction score of segment with `segment_id`.
- """
- class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1]
- masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width]
-
- batch_size = class_queries_logits.shape[0]
- num_labels = class_queries_logits.shape[-1] - 1
-
- mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width]
-
- # Predicted label and score of each query (batch_size, num_queries)
- pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1)
-
- # Loop over items in batch size
- results: List[Dict[str, TensorType]] = []
-
- for i in range(batch_size):
- mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects(
- mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels
- )
-
- # No mask found
- if mask_probs_item.shape[0] <= 0:
- height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:]
- segmentation = torch.zeros((height, width)) - 1
- results.append({"segmentation": segmentation, "segments_info": []})
- continue
-
- # Get segmentation map and segment information of batch item
- target_size = target_sizes[i] if target_sizes is not None else None
- segmentation, segments = compute_segments(
- mask_probs=mask_probs_item,
- pred_scores=pred_scores_item,
- pred_labels=pred_labels_item,
- mask_threshold=mask_threshold,
- overlap_mask_area_threshold=overlap_mask_area_threshold,
- label_ids_to_fuse=[],
- target_size=target_size,
- )
-
- # Return segmentation map in run-length encoding (RLE) format
- if return_coco_annotation:
- segmentation = convert_segmentation_to_rle(segmentation)
-
- results.append({"segmentation": segmentation, "segments_info": segments})
- return results
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.post_process_panoptic_segmentation with Detr->ConditionalDetr
- def post_process_panoptic_segmentation(
- self,
- outputs,
- threshold: float = 0.5,
- mask_threshold: float = 0.5,
- overlap_mask_area_threshold: float = 0.8,
- label_ids_to_fuse: Optional[Set[int]] = None,
- target_sizes: Optional[List[Tuple[int, int]]] = None,
- ) -> List[Dict]:
- """
- Converts the output of [`ConditionalDetrForSegmentation`] into image panoptic segmentation predictions. Only supports
- PyTorch.
-
- Args:
- outputs ([`ConditionalDetrForSegmentation`]):
- The outputs from [`ConditionalDetrForSegmentation`].
- threshold (`float`, *optional*, defaults to 0.5):
- The probability score threshold to keep predicted instance masks.
- mask_threshold (`float`, *optional*, defaults to 0.5):
- Threshold to use when turning the predicted masks into binary values.
- overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8):
- The overlap mask area threshold to merge or discard small disconnected parts within each binary
- instance mask.
- label_ids_to_fuse (`Set[int]`, *optional*):
- The labels in this state will have all their instances be fused together. For instance we could say
- there can only be one sky in an image, but several persons, so the label ID for sky would be in that
- set, but not the one for person.
- target_sizes (`List[Tuple]`, *optional*):
- List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested
- final size (height, width) of each prediction in batch. If unset, predictions will not be resized.
- Returns:
- `List[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys:
- - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id` or
- `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized to
- the corresponding `target_sizes` entry.
- - **segments_info** -- A dictionary that contains additional information on each segment.
- - **id** -- an integer representing the `segment_id`.
- - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`.
- - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise.
- Multiple instances of the same class / label were fused and assigned a single `segment_id`.
- - **score** -- Prediction score of segment with `segment_id`.
- """
-
- if label_ids_to_fuse is None:
- logger.warning_once("`label_ids_to_fuse` unset. No instance will be fused.")
- label_ids_to_fuse = set()
-
- class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1]
- masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width]
-
- batch_size = class_queries_logits.shape[0]
- num_labels = class_queries_logits.shape[-1] - 1
-
- mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width]
-
- # Predicted label and score of each query (batch_size, num_queries)
- pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1)
-
- # Loop over items in batch size
- results: List[Dict[str, TensorType]] = []
-
- for i in range(batch_size):
- mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects(
- mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels
- )
-
- # No mask found
- if mask_probs_item.shape[0] <= 0:
- height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:]
- segmentation = torch.zeros((height, width)) - 1
- results.append({"segmentation": segmentation, "segments_info": []})
- continue
-
- # Get segmentation map and segment information of batch item
- target_size = target_sizes[i] if target_sizes is not None else None
- segmentation, segments = compute_segments(
- mask_probs=mask_probs_item,
- pred_scores=pred_scores_item,
- pred_labels=pred_labels_item,
- mask_threshold=mask_threshold,
- overlap_mask_area_threshold=overlap_mask_area_threshold,
- label_ids_to_fuse=label_ids_to_fuse,
- target_size=target_size,
- )
-
- results.append({"segmentation": segmentation, "segments_info": segments})
- return results
diff --git a/transformers/models/conditional_detr/modeling_conditional_detr.py b/transformers/models/conditional_detr/modeling_conditional_detr.py
deleted file mode 100644
index d8ff371fad77d1d148179054a0eea2d811fea56c..0000000000000000000000000000000000000000
--- a/transformers/models/conditional_detr/modeling_conditional_detr.py
+++ /dev/null
@@ -1,2759 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Microsoft Research Asia and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Conditional DETR model."""
-
-
-import math
-from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple, Union
-
-import torch
-from torch import Tensor, nn
-
-from ...activations import ACT2FN
-from ...modeling_attn_mask_utils import _prepare_4d_attention_mask
-from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithCrossAttentions, Seq2SeqModelOutput
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_accelerate_available,
- is_scipy_available,
- is_timm_available,
- is_vision_available,
- logging,
- replace_return_docstrings,
- requires_backends,
-)
-from ...utils.backbone_utils import load_backbone
-from .configuration_conditional_detr import ConditionalDetrConfig
-
-
-if is_accelerate_available():
- from accelerate import PartialState
- from accelerate.utils import reduce
-
-if is_scipy_available():
- from scipy.optimize import linear_sum_assignment
-
-if is_timm_available():
- from timm import create_model
-
-if is_vision_available():
- from ...image_transforms import center_to_corners_format
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "ConditionalDetrConfig"
-_CHECKPOINT_FOR_DOC = "microsoft/conditional-detr-resnet-50"
-
-
-from ..deprecated._archive_maps import CONDITIONAL_DETR_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class ConditionalDetrDecoderOutput(BaseModelOutputWithCrossAttentions):
- """
- Base class for outputs of the Conditional DETR decoder. This class adds one attribute to
- BaseModelOutputWithCrossAttentions, namely an optional stack of intermediate decoder activations, i.e. the output
- of each decoder layer, each of them gone through a layernorm. This is useful when training the model with auxiliary
- decoding losses.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
- the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`):
- Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a
- layernorm.
- """
-
- intermediate_hidden_states: Optional[torch.FloatTensor] = None
- reference_points: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class ConditionalDetrModelOutput(Seq2SeqModelOutput):
- """
- Base class for outputs of the Conditional DETR encoder-decoder model. This class adds one attribute to
- Seq2SeqModelOutput, namely an optional stack of intermediate decoder activations, i.e. the output of each decoder
- layer, each of them gone through a layernorm. This is useful when training the model with auxiliary decoding
- losses.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the decoder at the output of each
- layer plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the encoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, sequence_length, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`):
- Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a
- layernorm.
- """
-
- intermediate_hidden_states: Optional[torch.FloatTensor] = None
- reference_points: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-# Copied from transformers.models.detr.modeling_detr.DetrObjectDetectionOutput with Detr->ConditionalDetr
-class ConditionalDetrObjectDetectionOutput(ModelOutput):
- """
- Output type of [`ConditionalDetrForObjectDetection`].
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
- Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
- bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
- scale-invariant IoU loss.
- loss_dict (`Dict`, *optional*):
- A dictionary containing the individual losses. Useful for logging.
- logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
- Classification logits (including no-object) for all queries.
- pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
- values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
- possible padding). You can use [`~ConditionalDetrImageProcessor.post_process_object_detection`] to retrieve the
- unnormalized bounding boxes.
- auxiliary_outputs (`list[Dict]`, *optional*):
- Optional, only returned when auxilary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
- and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
- `pred_boxes`) for each decoder layer.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the decoder at the output of each
- layer plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the encoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- loss_dict: Optional[Dict] = None
- logits: torch.FloatTensor = None
- pred_boxes: torch.FloatTensor = None
- auxiliary_outputs: Optional[List[Dict]] = None
- last_hidden_state: Optional[torch.FloatTensor] = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-# Copied from transformers.models.detr.modeling_detr.DetrSegmentationOutput with Detr->ConditionalDetr
-class ConditionalDetrSegmentationOutput(ModelOutput):
- """
- Output type of [`ConditionalDetrForSegmentation`].
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
- Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
- bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
- scale-invariant IoU loss.
- loss_dict (`Dict`, *optional*):
- A dictionary containing the individual losses. Useful for logging.
- logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
- Classification logits (including no-object) for all queries.
- pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
- values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
- possible padding). You can use [`~ConditionalDetrImageProcessor.post_process_object_detection`] to retrieve the
- unnormalized bounding boxes.
- pred_masks (`torch.FloatTensor` of shape `(batch_size, num_queries, height/4, width/4)`):
- Segmentation masks logits for all queries. See also
- [`~ConditionalDetrImageProcessor.post_process_semantic_segmentation`] or
- [`~ConditionalDetrImageProcessor.post_process_instance_segmentation`]
- [`~ConditionalDetrImageProcessor.post_process_panoptic_segmentation`] to evaluate semantic, instance and panoptic
- segmentation masks respectively.
- auxiliary_outputs (`list[Dict]`, *optional*):
- Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
- and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
- `pred_boxes`) for each decoder layer.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the decoder at the output of each
- layer plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the encoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- loss_dict: Optional[Dict] = None
- logits: torch.FloatTensor = None
- pred_boxes: torch.FloatTensor = None
- pred_masks: torch.FloatTensor = None
- auxiliary_outputs: Optional[List[Dict]] = None
- last_hidden_state: Optional[torch.FloatTensor] = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrFrozenBatchNorm2d with Detr->ConditionalDetr
-class ConditionalDetrFrozenBatchNorm2d(nn.Module):
- """
- BatchNorm2d where the batch statistics and the affine parameters are fixed.
-
- Copy-paste from torchvision.misc.ops with added eps before rqsrt, without which any other models than
- torchvision.models.resnet[18,34,50,101] produce nans.
- """
-
- def __init__(self, n):
- super().__init__()
- self.register_buffer("weight", torch.ones(n))
- self.register_buffer("bias", torch.zeros(n))
- self.register_buffer("running_mean", torch.zeros(n))
- self.register_buffer("running_var", torch.ones(n))
-
- def _load_from_state_dict(
- self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
- ):
- num_batches_tracked_key = prefix + "num_batches_tracked"
- if num_batches_tracked_key in state_dict:
- del state_dict[num_batches_tracked_key]
-
- super()._load_from_state_dict(
- state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
- )
-
- def forward(self, x):
- # move reshapes to the beginning
- # to make it user-friendly
- weight = self.weight.reshape(1, -1, 1, 1)
- bias = self.bias.reshape(1, -1, 1, 1)
- running_var = self.running_var.reshape(1, -1, 1, 1)
- running_mean = self.running_mean.reshape(1, -1, 1, 1)
- epsilon = 1e-5
- scale = weight * (running_var + epsilon).rsqrt()
- bias = bias - running_mean * scale
- return x * scale + bias
-
-
-# Copied from transformers.models.detr.modeling_detr.replace_batch_norm with Detr->ConditionalDetr
-def replace_batch_norm(model):
- r"""
- Recursively replace all `torch.nn.BatchNorm2d` with `ConditionalDetrFrozenBatchNorm2d`.
-
- Args:
- model (torch.nn.Module):
- input model
- """
- for name, module in model.named_children():
- if isinstance(module, nn.BatchNorm2d):
- new_module = ConditionalDetrFrozenBatchNorm2d(module.num_features)
-
- if not module.weight.device == torch.device("meta"):
- new_module.weight.data.copy_(module.weight)
- new_module.bias.data.copy_(module.bias)
- new_module.running_mean.data.copy_(module.running_mean)
- new_module.running_var.data.copy_(module.running_var)
-
- model._modules[name] = new_module
-
- if len(list(module.children())) > 0:
- replace_batch_norm(module)
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrConvEncoder
-class ConditionalDetrConvEncoder(nn.Module):
- """
- Convolutional backbone, using either the AutoBackbone API or one from the timm library.
-
- nn.BatchNorm2d layers are replaced by DetrFrozenBatchNorm2d as defined above.
-
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.config = config
-
- if config.use_timm_backbone:
- requires_backends(self, ["timm"])
- kwargs = {}
- if config.dilation:
- kwargs["output_stride"] = 16
- backbone = create_model(
- config.backbone,
- pretrained=config.use_pretrained_backbone,
- features_only=True,
- out_indices=(1, 2, 3, 4),
- in_chans=config.num_channels,
- **kwargs,
- )
- else:
- backbone = load_backbone(config)
-
- # replace batch norm by frozen batch norm
- with torch.no_grad():
- replace_batch_norm(backbone)
- self.model = backbone
- self.intermediate_channel_sizes = (
- self.model.feature_info.channels() if config.use_timm_backbone else self.model.channels
- )
-
- backbone_model_type = config.backbone if config.use_timm_backbone else config.backbone_config.model_type
- if "resnet" in backbone_model_type:
- for name, parameter in self.model.named_parameters():
- if config.use_timm_backbone:
- if "layer2" not in name and "layer3" not in name and "layer4" not in name:
- parameter.requires_grad_(False)
- else:
- if "stage.1" not in name and "stage.2" not in name and "stage.3" not in name:
- parameter.requires_grad_(False)
-
- def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
- # send pixel_values through the model to get list of feature maps
- features = self.model(pixel_values) if self.config.use_timm_backbone else self.model(pixel_values).feature_maps
-
- out = []
- for feature_map in features:
- # downsample pixel_mask to match shape of corresponding feature_map
- mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
- out.append((feature_map, mask))
- return out
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrConvModel with Detr->ConditionalDetr
-class ConditionalDetrConvModel(nn.Module):
- """
- This module adds 2D position embeddings to all intermediate feature maps of the convolutional encoder.
- """
-
- def __init__(self, conv_encoder, position_embedding):
- super().__init__()
- self.conv_encoder = conv_encoder
- self.position_embedding = position_embedding
-
- def forward(self, pixel_values, pixel_mask):
- # send pixel_values and pixel_mask through backbone to get list of (feature_map, pixel_mask) tuples
- out = self.conv_encoder(pixel_values, pixel_mask)
- pos = []
- for feature_map, mask in out:
- # position encoding
- pos.append(self.position_embedding(feature_map, mask).to(feature_map.dtype))
-
- return out, pos
-
-
-class ConditionalDetrSinePositionEmbedding(nn.Module):
- """
- This is a more standard version of the position embedding, very similar to the one used by the Attention is all you
- need paper, generalized to work on images.
- """
-
- def __init__(self, embedding_dim=64, temperature=10000, normalize=False, scale=None):
- super().__init__()
- self.embedding_dim = embedding_dim
- self.temperature = temperature
- self.normalize = normalize
- if scale is not None and normalize is False:
- raise ValueError("normalize should be True if scale is passed")
- if scale is None:
- scale = 2 * math.pi
- self.scale = scale
-
- def forward(self, pixel_values, pixel_mask):
- if pixel_mask is None:
- raise ValueError("No pixel mask provided")
- y_embed = pixel_mask.cumsum(1, dtype=torch.float32)
- x_embed = pixel_mask.cumsum(2, dtype=torch.float32)
- if self.normalize:
- y_embed = y_embed / (y_embed[:, -1:, :] + 1e-6) * self.scale
- x_embed = x_embed / (x_embed[:, :, -1:] + 1e-6) * self.scale
-
- dim_t = torch.arange(self.embedding_dim, dtype=torch.int64, device=pixel_values.device).float()
- dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.embedding_dim)
-
- pos_x = x_embed[:, :, :, None] / dim_t
- pos_y = y_embed[:, :, :, None] / dim_t
- pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
- return pos
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrLearnedPositionEmbedding with Detr->ConditionalDetr
-class ConditionalDetrLearnedPositionEmbedding(nn.Module):
- """
- This module learns positional embeddings up to a fixed maximum size.
- """
-
- def __init__(self, embedding_dim=256):
- super().__init__()
- self.row_embeddings = nn.Embedding(50, embedding_dim)
- self.column_embeddings = nn.Embedding(50, embedding_dim)
-
- def forward(self, pixel_values, pixel_mask=None):
- height, width = pixel_values.shape[-2:]
- width_values = torch.arange(width, device=pixel_values.device)
- height_values = torch.arange(height, device=pixel_values.device)
- x_emb = self.column_embeddings(width_values)
- y_emb = self.row_embeddings(height_values)
- pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1)
- pos = pos.permute(2, 0, 1)
- pos = pos.unsqueeze(0)
- pos = pos.repeat(pixel_values.shape[0], 1, 1, 1)
- return pos
-
-
-# Copied from transformers.models.detr.modeling_detr.build_position_encoding with Detr->ConditionalDetr
-def build_position_encoding(config):
- n_steps = config.d_model // 2
- if config.position_embedding_type == "sine":
- # TODO find a better way of exposing other arguments
- position_embedding = ConditionalDetrSinePositionEmbedding(n_steps, normalize=True)
- elif config.position_embedding_type == "learned":
- position_embedding = ConditionalDetrLearnedPositionEmbedding(n_steps)
- else:
- raise ValueError(f"Not supported {config.position_embedding_type}")
-
- return position_embedding
-
-
-# function to generate sine positional embedding for 2d coordinates
-def gen_sine_position_embeddings(pos_tensor, d_model):
- scale = 2 * math.pi
- dim = d_model // 2
- dim_t = torch.arange(dim, dtype=torch.float32, device=pos_tensor.device)
- dim_t = 10000 ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / dim)
- x_embed = pos_tensor[:, :, 0] * scale
- y_embed = pos_tensor[:, :, 1] * scale
- pos_x = x_embed[:, :, None] / dim_t
- pos_y = y_embed[:, :, None] / dim_t
- pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)
- pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2)
- pos = torch.cat((pos_y, pos_x), dim=2)
- return pos
-
-
-def inverse_sigmoid(x, eps=1e-5):
- x = x.clamp(min=0, max=1)
- x1 = x.clamp(min=eps)
- x2 = (1 - x).clamp(min=eps)
- return torch.log(x1 / x2)
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrAttention
-class DetrAttention(nn.Module):
- """
- Multi-headed attention from 'Attention Is All You Need' paper.
-
- Here, we add position embeddings to the queries and keys (as explained in the DETR paper).
- """
-
- def __init__(
- self,
- embed_dim: int,
- num_heads: int,
- dropout: float = 0.0,
- bias: bool = True,
- ):
- super().__init__()
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.dropout = dropout
- self.head_dim = embed_dim // num_heads
- if self.head_dim * num_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {num_heads})."
- )
- self.scaling = self.head_dim**-0.5
-
- self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
-
- def _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int):
- return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def with_pos_embed(self, tensor: torch.Tensor, object_queries: Optional[Tensor], **kwargs):
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- return tensor if object_queries is None else tensor + object_queries
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- object_queries: Optional[torch.Tensor] = None,
- key_value_states: Optional[torch.Tensor] = None,
- spatial_position_embeddings: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- **kwargs,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel"""
-
- position_embeddings = kwargs.pop("position_ebmeddings", None)
- key_value_position_embeddings = kwargs.pop("key_value_position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if key_value_position_embeddings is not None and spatial_position_embeddings is not None:
- raise ValueError(
- "Cannot specify both key_value_position_embeddings and spatial_position_embeddings. Please use just spatial_position_embeddings"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- if key_value_position_embeddings is not None:
- logger.warning_once(
- "key_value_position_embeddings has been deprecated and will be removed in v4.34. Please use spatial_position_embeddings instead"
- )
- spatial_position_embeddings = key_value_position_embeddings
-
- # if key_value_states are provided this layer is used as a cross-attention layer
- # for the decoder
- is_cross_attention = key_value_states is not None
- batch_size, target_len, embed_dim = hidden_states.size()
-
- # add position embeddings to the hidden states before projecting to queries and keys
- if object_queries is not None:
- hidden_states_original = hidden_states
- hidden_states = self.with_pos_embed(hidden_states, object_queries)
-
- # add key-value position embeddings to the key value states
- if spatial_position_embeddings is not None:
- key_value_states_original = key_value_states
- key_value_states = self.with_pos_embed(key_value_states, spatial_position_embeddings)
-
- # get query proj
- query_states = self.q_proj(hidden_states) * self.scaling
- # get key, value proj
- if is_cross_attention:
- # cross_attentions
- key_states = self._shape(self.k_proj(key_value_states), -1, batch_size)
- value_states = self._shape(self.v_proj(key_value_states_original), -1, batch_size)
- else:
- # self_attention
- key_states = self._shape(self.k_proj(hidden_states), -1, batch_size)
- value_states = self._shape(self.v_proj(hidden_states_original), -1, batch_size)
-
- proj_shape = (batch_size * self.num_heads, -1, self.head_dim)
- query_states = self._shape(query_states, target_len, batch_size).view(*proj_shape)
- key_states = key_states.view(*proj_shape)
- value_states = value_states.view(*proj_shape)
-
- source_len = key_states.size(1)
-
- attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-
- if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len):
- raise ValueError(
- f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- if attention_mask is not None:
- if attention_mask.size() != (batch_size, 1, target_len, source_len):
- raise ValueError(
- f"Attention mask should be of size {(batch_size, 1, target_len, source_len)}, but is"
- f" {attention_mask.size()}"
- )
- attn_weights = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attention_mask
- attn_weights = attn_weights.view(batch_size * self.num_heads, target_len, source_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if output_attentions:
- # this operation is a bit awkward, but it's required to
- # make sure that attn_weights keeps its gradient.
- # In order to do so, attn_weights have to reshaped
- # twice and have to be reused in the following
- attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len)
- attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-
- attn_output = torch.bmm(attn_probs, value_states)
-
- if attn_output.size() != (batch_size * self.num_heads, target_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.head_dim)
- attn_output = attn_output.transpose(1, 2)
- attn_output = attn_output.reshape(batch_size, target_len, embed_dim)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped
-
-
-class ConditionalDetrAttention(nn.Module):
- """
- Cross-Attention used in Conditional DETR 'Conditional DETR for Fast Training Convergence' paper.
-
- The key q_proj, k_proj, v_proj are defined outside the attention. This attention allows the dim of q, k to be
- different to v.
- """
-
- def __init__(
- self,
- embed_dim: int,
- out_dim: int,
- num_heads: int,
- dropout: float = 0.0,
- bias: bool = True,
- ):
- super().__init__()
- self.embed_dim = embed_dim
- self.out_dim = out_dim
- self.num_heads = num_heads
- self.dropout = dropout
- self.head_dim = embed_dim // num_heads
- if self.head_dim * num_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {num_heads})."
- )
- # head dimension of values
- self.v_head_dim = out_dim // num_heads
- if self.v_head_dim * num_heads != self.out_dim:
- raise ValueError(
- f"out_dim must be divisible by num_heads (got `out_dim`: {self.out_dim} and `num_heads`: {num_heads})."
- )
- self.scaling = self.head_dim**-0.5
-
- self.out_proj = nn.Linear(out_dim, out_dim, bias=bias)
-
- def _qk_shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int):
- return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def _v_shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int):
- return tensor.view(batch_size, seq_len, self.num_heads, self.v_head_dim).transpose(1, 2).contiguous()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- key_states: Optional[torch.Tensor] = None,
- value_states: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel"""
-
- batch_size, target_len, _ = hidden_states.size()
-
- # get query proj
- query_states = hidden_states * self.scaling
- # get key, value proj
- key_states = self._qk_shape(key_states, -1, batch_size)
- value_states = self._v_shape(value_states, -1, batch_size)
-
- proj_shape = (batch_size * self.num_heads, -1, self.head_dim)
- v_proj_shape = (batch_size * self.num_heads, -1, self.v_head_dim)
- query_states = self._qk_shape(query_states, target_len, batch_size).view(*proj_shape)
- key_states = key_states.view(*proj_shape)
- value_states = value_states.view(*v_proj_shape)
-
- source_len = key_states.size(1)
-
- attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-
- if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len):
- raise ValueError(
- f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- if attention_mask is not None:
- if attention_mask.size() != (batch_size, 1, target_len, source_len):
- raise ValueError(
- f"Attention mask should be of size {(batch_size, 1, target_len, source_len)}, but is"
- f" {attention_mask.size()}"
- )
- attn_weights = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attention_mask
- attn_weights = attn_weights.view(batch_size * self.num_heads, target_len, source_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if output_attentions:
- # this operation is a bit awkward, but it's required to
- # make sure that attn_weights keeps its gradient.
- # In order to do so, attn_weights have to reshaped
- # twice and have to be reused in the following
- attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len)
- attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-
- attn_output = torch.bmm(attn_probs, value_states)
-
- if attn_output.size() != (batch_size * self.num_heads, target_len, self.v_head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.v_head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.v_head_dim)
- attn_output = attn_output.transpose(1, 2)
- attn_output = attn_output.reshape(batch_size, target_len, self.out_dim)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrEncoderLayer with DetrEncoderLayer->ConditionalDetrEncoderLayer,DetrConfig->ConditionalDetrConfig
-class ConditionalDetrEncoderLayer(nn.Module):
- def __init__(self, config: ConditionalDetrConfig):
- super().__init__()
- self.embed_dim = config.d_model
- self.self_attn = DetrAttention(
- embed_dim=self.embed_dim,
- num_heads=config.encoder_attention_heads,
- dropout=config.attention_dropout,
- )
- self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
- self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
- self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: torch.Tensor,
- object_queries: torch.Tensor = None,
- output_attentions: bool = False,
- **kwargs,
- ):
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`): attention mask of size
- `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
- values.
- object_queries (`torch.FloatTensor`, *optional*):
- Object queries (also called content embeddings), to be added to the hidden states.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- residual = hidden_states
- hidden_states, attn_weights = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- object_queries=object_queries,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- residual = hidden_states
- hidden_states = self.activation_fn(self.fc1(hidden_states))
- hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
-
- hidden_states = self.fc2(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
- hidden_states = residual + hidden_states
- hidden_states = self.final_layer_norm(hidden_states)
-
- if self.training:
- if torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any():
- clamp_value = torch.finfo(hidden_states.dtype).max - 1000
- hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
-
-class ConditionalDetrDecoderLayer(nn.Module):
- def __init__(self, config: ConditionalDetrConfig):
- super().__init__()
- self.embed_dim = config.d_model
-
- d_model = config.d_model
- # Decoder Self-Attention projections
- self.sa_qcontent_proj = nn.Linear(d_model, d_model)
- self.sa_qpos_proj = nn.Linear(d_model, d_model)
- self.sa_kcontent_proj = nn.Linear(d_model, d_model)
- self.sa_kpos_proj = nn.Linear(d_model, d_model)
- self.sa_v_proj = nn.Linear(d_model, d_model)
-
- self.self_attn = ConditionalDetrAttention(
- embed_dim=self.embed_dim,
- out_dim=self.embed_dim,
- num_heads=config.decoder_attention_heads,
- dropout=config.attention_dropout,
- )
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
-
- self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
-
- # Decoder Cross-Attention projections
- self.ca_qcontent_proj = nn.Linear(d_model, d_model)
- self.ca_qpos_proj = nn.Linear(d_model, d_model)
- self.ca_kcontent_proj = nn.Linear(d_model, d_model)
- self.ca_kpos_proj = nn.Linear(d_model, d_model)
- self.ca_v_proj = nn.Linear(d_model, d_model)
- self.ca_qpos_sine_proj = nn.Linear(d_model, d_model)
-
- self.encoder_attn = ConditionalDetrAttention(
- self.embed_dim * 2, self.embed_dim, config.decoder_attention_heads, dropout=config.attention_dropout
- )
- self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
- self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = nn.LayerNorm(self.embed_dim)
- self.nhead = config.decoder_attention_heads
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- object_queries: Optional[torch.Tensor] = None,
- query_position_embeddings: Optional[torch.Tensor] = None,
- query_sine_embed: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- is_first: Optional[bool] = False,
- **kwargs,
- ):
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)`
- attention_mask (`torch.FloatTensor`): attention mask of size
- `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
- values.
- object_queries (`torch.FloatTensor`, *optional*):
- object_queries that are added to the queries and keys
- in the cross-attention layer.
- query_position_embeddings (`torch.FloatTensor`, *optional*):
- object_queries that are added to the queries and keys
- in the self-attention layer.
- encoder_hidden_states (`torch.FloatTensor`):
- cross attention input to the layer of shape `(seq_len, batch, embed_dim)`
- encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
- `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
- values.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- residual = hidden_states
-
- # ========== Begin of Self-Attention =============
- # Apply projections here
- # shape: num_queries x batch_size x 256
- q_content = self.sa_qcontent_proj(
- hidden_states
- ) # target is the input of the first decoder layer. zero by default.
- q_pos = self.sa_qpos_proj(query_position_embeddings)
- k_content = self.sa_kcontent_proj(hidden_states)
- k_pos = self.sa_kpos_proj(query_position_embeddings)
- v = self.sa_v_proj(hidden_states)
-
- _, num_queries, n_model = q_content.shape
-
- q = q_content + q_pos
- k = k_content + k_pos
- hidden_states, self_attn_weights = self.self_attn(
- hidden_states=q,
- attention_mask=attention_mask,
- key_states=k,
- value_states=v,
- output_attentions=output_attentions,
- )
- # ============ End of Self-Attention =============
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- # ========== Begin of Cross-Attention =============
- # Apply projections here
- # shape: num_queries x batch_size x 256
- q_content = self.ca_qcontent_proj(hidden_states)
- k_content = self.ca_kcontent_proj(encoder_hidden_states)
- v = self.ca_v_proj(encoder_hidden_states)
-
- batch_size, num_queries, n_model = q_content.shape
- _, source_len, _ = k_content.shape
-
- k_pos = self.ca_kpos_proj(object_queries)
-
- # For the first decoder layer, we concatenate the positional embedding predicted from
- # the object query (the positional embedding) into the original query (key) in DETR.
- if is_first:
- q_pos = self.ca_qpos_proj(query_position_embeddings)
- q = q_content + q_pos
- k = k_content + k_pos
- else:
- q = q_content
- k = k_content
-
- q = q.view(batch_size, num_queries, self.nhead, n_model // self.nhead)
- query_sine_embed = self.ca_qpos_sine_proj(query_sine_embed)
- query_sine_embed = query_sine_embed.view(batch_size, num_queries, self.nhead, n_model // self.nhead)
- q = torch.cat([q, query_sine_embed], dim=3).view(batch_size, num_queries, n_model * 2)
- k = k.view(batch_size, source_len, self.nhead, n_model // self.nhead)
- k_pos = k_pos.view(batch_size, source_len, self.nhead, n_model // self.nhead)
- k = torch.cat([k, k_pos], dim=3).view(batch_size, source_len, n_model * 2)
-
- # Cross-Attention Block
- cross_attn_weights = None
- if encoder_hidden_states is not None:
- residual = hidden_states
-
- hidden_states, cross_attn_weights = self.encoder_attn(
- hidden_states=q,
- attention_mask=encoder_attention_mask,
- key_states=k,
- value_states=v,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.encoder_attn_layer_norm(hidden_states)
-
- # ============ End of Cross-Attention =============
-
- # Fully Connected
- residual = hidden_states
- hidden_states = self.activation_fn(self.fc1(hidden_states))
- hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
- hidden_states = self.fc2(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.final_layer_norm(hidden_states)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights, cross_attn_weights)
-
- return outputs
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrClassificationHead with Detr->ConditionalDetr
-class ConditionalDetrClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, input_dim: int, inner_dim: int, num_classes: int, pooler_dropout: float):
- super().__init__()
- self.dense = nn.Linear(input_dim, inner_dim)
- self.dropout = nn.Dropout(p=pooler_dropout)
- self.out_proj = nn.Linear(inner_dim, num_classes)
-
- def forward(self, hidden_states: torch.Tensor):
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.dense(hidden_states)
- hidden_states = torch.tanh(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.out_proj(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrMLPPredictionHead with DetrMLPPredictionHead->MLP
-class MLP(nn.Module):
- """
- Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates,
- height and width of a bounding box w.r.t. an image.
-
- Copied from https://github.com/facebookresearch/detr/blob/master/models/detr.py
-
- """
-
- def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
- super().__init__()
- self.num_layers = num_layers
- h = [hidden_dim] * (num_layers - 1)
- self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
-
- def forward(self, x):
- for i, layer in enumerate(self.layers):
- x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
- return x
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrPreTrainedModel with Detr->ConditionalDetr
-class ConditionalDetrPreTrainedModel(PreTrainedModel):
- config_class = ConditionalDetrConfig
- base_model_prefix = "model"
- main_input_name = "pixel_values"
- _no_split_modules = [r"ConditionalDetrConvEncoder", r"ConditionalDetrEncoderLayer", r"ConditionalDetrDecoderLayer"]
-
- def _init_weights(self, module):
- std = self.config.init_std
- xavier_std = self.config.init_xavier_std
-
- if isinstance(module, ConditionalDetrMHAttentionMap):
- nn.init.zeros_(module.k_linear.bias)
- nn.init.zeros_(module.q_linear.bias)
- nn.init.xavier_uniform_(module.k_linear.weight, gain=xavier_std)
- nn.init.xavier_uniform_(module.q_linear.weight, gain=xavier_std)
- elif isinstance(module, ConditionalDetrLearnedPositionEmbedding):
- nn.init.uniform_(module.row_embeddings.weight)
- nn.init.uniform_(module.column_embeddings.weight)
- if isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
-
-CONDITIONAL_DETR_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`ConditionalDetrConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CONDITIONAL_DETR_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Padding will be ignored by default should you provide it.
-
- Pixel values can be obtained using [`AutoImageProcessor`]. See [`ConditionalDetrImageProcessor.__call__`]
- for details.
-
- pixel_mask (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
- Mask to avoid performing attention on padding pixel values. Mask values selected in `[0, 1]`:
-
- - 1 for pixels that are real (i.e. **not masked**),
- - 0 for pixels that are padding (i.e. **masked**).
-
- [What are attention masks?](../glossary#attention-mask)
-
- decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*):
- Not used by default. Can be used to mask object queries.
- encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*):
- Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
- `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of
- hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you
- can choose to directly pass a flattened representation of an image.
- decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
- Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an
- embedded representation.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrEncoder with Detr->ConditionalDetr,DETR->ConditionalDETR
-class ConditionalDetrEncoder(ConditionalDetrPreTrainedModel):
- """
- Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
- [`ConditionalDetrEncoderLayer`].
-
- The encoder updates the flattened feature map through multiple self-attention layers.
-
- Small tweak for ConditionalDETR:
-
- - object_queries are added to the forward pass.
-
- Args:
- config: ConditionalDetrConfig
- """
-
- def __init__(self, config: ConditionalDetrConfig):
- super().__init__(config)
-
- self.dropout = config.dropout
- self.layerdrop = config.encoder_layerdrop
-
- self.layers = nn.ModuleList([ConditionalDetrEncoderLayer(config) for _ in range(config.encoder_layers)])
-
- # in the original ConditionalDETR, no layernorm is used at the end of the encoder, as "normalize_before" is set to False by default
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def forward(
- self,
- inputs_embeds=None,
- attention_mask=None,
- object_queries=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- **kwargs,
- ):
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Flattened feature map (output of the backbone + projection layer) that is passed to the encoder.
-
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`:
-
- - 1 for pixel features that are real (i.e. **not masked**),
- - 0 for pixel features that are padding (i.e. **masked**).
-
- [What are attention masks?](../glossary#attention-mask)
-
- object_queries (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Object queries that are added to the queries in each self-attention layer.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- """
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- hidden_states = inputs_embeds
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
- # expand attention_mask
- if attention_mask is not None:
- # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len]
- attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
-
- encoder_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
- for i, encoder_layer in enumerate(self.layers):
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
- # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
- to_drop = False
- if self.training:
- dropout_probability = torch.rand([])
- if dropout_probability < self.layerdrop: # skip the layer
- to_drop = True
-
- if to_drop:
- layer_outputs = (None, None)
- else:
- # we add object_queries as extra input to the encoder_layer
- layer_outputs = encoder_layer(
- hidden_states,
- attention_mask,
- object_queries=object_queries,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
- )
-
-
-class ConditionalDetrDecoder(ConditionalDetrPreTrainedModel):
- """
- Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`ConditionalDetrDecoderLayer`].
-
- The decoder updates the query embeddings through multiple self-attention and cross-attention layers.
-
- Some small tweaks for Conditional DETR:
-
- - object_queries and query_position_embeddings are added to the forward pass.
- - if self.config.auxiliary_loss is set to True, also returns a stack of activations from all decoding layers.
-
- Args:
- config: ConditionalDetrConfig
- """
-
- def __init__(self, config: ConditionalDetrConfig):
- super().__init__(config)
- self.dropout = config.dropout
- self.layerdrop = config.decoder_layerdrop
-
- self.layers = nn.ModuleList([ConditionalDetrDecoderLayer(config) for _ in range(config.decoder_layers)])
- # in Conditional DETR, the decoder uses layernorm after the last decoder layer output
- self.layernorm = nn.LayerNorm(config.d_model)
- d_model = config.d_model
- self.gradient_checkpointing = False
-
- # query_scale is the FFN applied on f to generate transformation T
- self.query_scale = MLP(d_model, d_model, d_model, 2)
- self.ref_point_head = MLP(d_model, d_model, 2, 2)
- for layer_id in range(config.decoder_layers - 1):
- self.layers[layer_id + 1].ca_qpos_proj = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def forward(
- self,
- inputs_embeds=None,
- attention_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- object_queries=None,
- query_position_embeddings=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- **kwargs,
- ):
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- The query embeddings that are passed into the decoder.
-
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on certain queries. Mask values selected in `[0, 1]`:
-
- - 1 for queries that are **not masked**,
- - 0 for queries that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
- of the decoder.
- encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
- Mask to avoid performing cross-attention on padding pixel_values of the encoder. Mask values selected
- in `[0, 1]`:
-
- - 1 for pixels that are real (i.e. **not masked**),
- - 0 for pixels that are padding (i.e. **masked**).
-
- object_queries (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Position embeddings that are added to the queries and keys in each cross-attention layer.
- query_position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
- , *optional*): Position embeddings that are added to the queries and keys in each self-attention layer.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- """
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if inputs_embeds is not None:
- hidden_states = inputs_embeds
- input_shape = inputs_embeds.size()[:-1]
-
- # expand encoder attention mask
- if encoder_hidden_states is not None and encoder_attention_mask is not None:
- # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len]
- encoder_attention_mask = _prepare_4d_attention_mask(
- encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
- )
-
- # optional intermediate hidden states
- intermediate = () if self.config.auxiliary_loss else None
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
-
- reference_points_before_sigmoid = self.ref_point_head(
- query_position_embeddings
- ) # [num_queries, batch_size, 2]
- reference_points = reference_points_before_sigmoid.sigmoid().transpose(0, 1)
- obj_center = reference_points[..., :2].transpose(0, 1)
- # get sine embedding for the query vector
- query_sine_embed_before_transformation = gen_sine_position_embeddings(obj_center, self.config.d_model)
-
- for idx, decoder_layer in enumerate(self.layers):
- # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
- if self.training:
- dropout_probability = torch.rand([])
- if dropout_probability < self.layerdrop:
- continue
- if idx == 0:
- pos_transformation = 1
- else:
- pos_transformation = self.query_scale(hidden_states)
- # apply transformation
- query_sine_embed = query_sine_embed_before_transformation * pos_transformation
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- decoder_layer.__call__,
- hidden_states,
- None,
- object_queries,
- query_position_embeddings,
- query_sine_embed,
- encoder_hidden_states,
- encoder_attention_mask,
- None,
- None,
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- attention_mask=None,
- object_queries=object_queries,
- query_position_embeddings=query_position_embeddings,
- query_sine_embed=query_sine_embed,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- is_first=(idx == 0),
- )
-
- hidden_states = layer_outputs[0]
-
- if self.config.auxiliary_loss:
- hidden_states = self.layernorm(hidden_states)
- intermediate += (hidden_states,)
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- if encoder_hidden_states is not None:
- all_cross_attentions += (layer_outputs[2],)
-
- # finally, apply layernorm
- hidden_states = self.layernorm(hidden_states)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- # stack intermediate decoder activations
- if self.config.auxiliary_loss:
- intermediate = torch.stack(intermediate)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- all_hidden_states,
- all_self_attns,
- all_cross_attentions,
- intermediate,
- reference_points,
- ]
- if v is not None
- )
- return ConditionalDetrDecoderOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- cross_attentions=all_cross_attentions,
- intermediate_hidden_states=intermediate,
- reference_points=reference_points,
- )
-
-
-@add_start_docstrings(
- """
- The bare Conditional DETR Model (consisting of a backbone and encoder-decoder Transformer) outputting raw
- hidden-states without any specific head on top.
- """,
- CONDITIONAL_DETR_START_DOCSTRING,
-)
-class ConditionalDetrModel(ConditionalDetrPreTrainedModel):
- def __init__(self, config: ConditionalDetrConfig):
- super().__init__(config)
-
- # Create backbone + positional encoding
- backbone = ConditionalDetrConvEncoder(config)
- object_queries = build_position_encoding(config)
- self.backbone = ConditionalDetrConvModel(backbone, object_queries)
-
- # Create projection layer
- self.input_projection = nn.Conv2d(backbone.intermediate_channel_sizes[-1], config.d_model, kernel_size=1)
-
- self.query_position_embeddings = nn.Embedding(config.num_queries, config.d_model)
-
- self.encoder = ConditionalDetrEncoder(config)
- self.decoder = ConditionalDetrDecoder(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_encoder(self):
- return self.encoder
-
- def get_decoder(self):
- return self.decoder
-
- def freeze_backbone(self):
- for name, param in self.backbone.conv_encoder.model.named_parameters():
- param.requires_grad_(False)
-
- def unfreeze_backbone(self):
- for name, param in self.backbone.conv_encoder.model.named_parameters():
- param.requires_grad_(True)
-
- @add_start_docstrings_to_model_forward(CONDITIONAL_DETR_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=ConditionalDetrModelOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.LongTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], ConditionalDetrModelOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, AutoModel
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50")
- >>> model = AutoModel.from_pretrained("microsoft/conditional-detr-resnet-50")
-
- >>> # prepare image for the model
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> # forward pass
- >>> outputs = model(**inputs)
-
- >>> # the last hidden states are the final query embeddings of the Transformer decoder
- >>> # these are of shape (batch_size, num_queries, hidden_size)
- >>> last_hidden_states = outputs.last_hidden_state
- >>> list(last_hidden_states.shape)
- [1, 300, 256]
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- batch_size, num_channels, height, width = pixel_values.shape
- device = pixel_values.device
-
- if pixel_mask is None:
- pixel_mask = torch.ones(((batch_size, height, width)), device=device)
-
- # First, sent pixel_values + pixel_mask through Backbone to obtain the features
- # pixel_values should be of shape (batch_size, num_channels, height, width)
- # pixel_mask should be of shape (batch_size, height, width)
- features, object_queries_list = self.backbone(pixel_values, pixel_mask)
-
- # get final feature map and downsampled mask
- feature_map, mask = features[-1]
-
- if mask is None:
- raise ValueError("Backbone does not return downsampled pixel mask")
-
- # Second, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
- projected_feature_map = self.input_projection(feature_map)
-
- # Third, flatten the feature map + object_queries of shape NxCxHxW to NxCxHW, and permute it to NxHWxC
- # In other words, turn their shape into (batch_size, sequence_length, hidden_size)
- flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1)
- object_queries = object_queries_list[-1].flatten(2).permute(0, 2, 1)
-
- flattened_mask = mask.flatten(1)
-
- # Fourth, sent flattened_features + flattened_mask + object_queries through encoder
- # flattened_features is a Tensor of shape (batch_size, heigth*width, hidden_size)
- # flattened_mask is a Tensor of shape (batch_size, heigth*width)
- if encoder_outputs is None:
- encoder_outputs = self.encoder(
- inputs_embeds=flattened_features,
- attention_mask=flattened_mask,
- object_queries=object_queries,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
- elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
- encoder_outputs = BaseModelOutput(
- last_hidden_state=encoder_outputs[0],
- hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
- attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
- )
-
- # Fifth, sent query embeddings + object_queries through the decoder (which is conditioned on the encoder output)
- query_position_embeddings = self.query_position_embeddings.weight.unsqueeze(0).repeat(batch_size, 1, 1)
- queries = torch.zeros_like(query_position_embeddings)
-
- # decoder outputs consists of (dec_features, dec_hidden, dec_attn)
- decoder_outputs = self.decoder(
- inputs_embeds=queries,
- attention_mask=None,
- object_queries=object_queries,
- query_position_embeddings=query_position_embeddings,
- encoder_hidden_states=encoder_outputs[0],
- encoder_attention_mask=flattened_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- return decoder_outputs + encoder_outputs
-
- return ConditionalDetrModelOutput(
- last_hidden_state=decoder_outputs.last_hidden_state,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
- reference_points=decoder_outputs.reference_points,
- )
-
-
-@add_start_docstrings(
- """
- CONDITIONAL_DETR Model (consisting of a backbone and encoder-decoder Transformer) with object detection heads on
- top, for tasks such as COCO detection.
- """,
- CONDITIONAL_DETR_START_DOCSTRING,
-)
-class ConditionalDetrForObjectDetection(ConditionalDetrPreTrainedModel):
- def __init__(self, config: ConditionalDetrConfig):
- super().__init__(config)
-
- # CONDITIONAL DETR encoder-decoder model
- self.model = ConditionalDetrModel(config)
-
- # Object detection heads
- self.class_labels_classifier = nn.Linear(
- config.d_model, config.num_labels
- ) # We add one for the "no object" class
- self.bbox_predictor = ConditionalDetrMLPPredictionHead(
- input_dim=config.d_model, hidden_dim=config.d_model, output_dim=4, num_layers=3
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- # taken from https://github.com/Atten4Vis/conditionalDETR/blob/master/models/conditional_detr.py
- @torch.jit.unused
- def _set_aux_loss(self, outputs_class, outputs_coord):
- # this is a workaround to make torchscript happy, as torchscript
- # doesn't support dictionary with non-homogeneous values, such
- # as a dict having both a Tensor and a list.
- return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
-
- @add_start_docstrings_to_model_forward(CONDITIONAL_DETR_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=ConditionalDetrObjectDetectionOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.LongTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[List[dict]] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], ConditionalDetrObjectDetectionOutput]:
- r"""
- labels (`List[Dict]` of len `(batch_size,)`, *optional*):
- Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
- following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
- respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
- in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, AutoModelForObjectDetection
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50")
- >>> model = AutoModelForObjectDetection.from_pretrained("microsoft/conditional-detr-resnet-50")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> outputs = model(**inputs)
-
- >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
- >>> target_sizes = torch.tensor([image.size[::-1]])
- >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[
- ... 0
- ... ]
- >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
- ... box = [round(i, 2) for i in box.tolist()]
- ... print(
- ... f"Detected {model.config.id2label[label.item()]} with confidence "
- ... f"{round(score.item(), 3)} at location {box}"
- ... )
- Detected remote with confidence 0.833 at location [38.31, 72.1, 177.63, 118.45]
- Detected cat with confidence 0.831 at location [9.2, 51.38, 321.13, 469.0]
- Detected cat with confidence 0.804 at location [340.3, 16.85, 642.93, 370.95]
- Detected remote with confidence 0.683 at location [334.48, 73.49, 366.37, 190.01]
- Detected couch with confidence 0.535 at location [0.52, 1.19, 640.35, 475.1]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # First, sent images through CONDITIONAL_DETR base model to obtain encoder + decoder outputs
- outputs = self.model(
- pixel_values,
- pixel_mask=pixel_mask,
- decoder_attention_mask=decoder_attention_mask,
- encoder_outputs=encoder_outputs,
- inputs_embeds=inputs_embeds,
- decoder_inputs_embeds=decoder_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- # class logits + predicted bounding boxes
- logits = self.class_labels_classifier(sequence_output)
-
- reference = outputs.reference_points if return_dict else outputs[-1]
- reference_before_sigmoid = inverse_sigmoid(reference).transpose(0, 1)
- outputs_coords = []
- hs = sequence_output
- tmp = self.bbox_predictor(hs)
- tmp[..., :2] += reference_before_sigmoid
- pred_boxes = tmp.sigmoid()
- # pred_boxes = self.bbox_predictor(sequence_output).sigmoid()
-
- loss, loss_dict, auxiliary_outputs = None, None, None
- if labels is not None:
- # First: create the matcher
- matcher = ConditionalDetrHungarianMatcher(
- class_cost=self.config.class_cost, bbox_cost=self.config.bbox_cost, giou_cost=self.config.giou_cost
- )
- # Second: create the criterion
- losses = ["labels", "boxes", "cardinality"]
- criterion = ConditionalDetrLoss(
- matcher=matcher,
- num_classes=self.config.num_labels,
- focal_alpha=self.config.focal_alpha,
- losses=losses,
- )
- criterion.to(self.device)
- # Third: compute the losses, based on outputs and labels
- outputs_loss = {}
- outputs_loss["logits"] = logits
- outputs_loss["pred_boxes"] = pred_boxes
- if self.config.auxiliary_loss:
- intermediate = outputs.intermediate_hidden_states if return_dict else outputs[4]
- outputs_class = self.class_labels_classifier(intermediate)
-
- for lvl in range(intermediate.shape[0]):
- tmp = self.bbox_predictor(intermediate[lvl])
- tmp[..., :2] += reference_before_sigmoid
- outputs_coord = tmp.sigmoid()
- outputs_coords.append(outputs_coord)
- outputs_coord = torch.stack(outputs_coords)
-
- auxiliary_outputs = self._set_aux_loss(outputs_class, outputs_coord)
- outputs_loss["auxiliary_outputs"] = auxiliary_outputs
-
- loss_dict = criterion(outputs_loss, labels)
- # Fourth: compute total loss, as a weighted sum of the various losses
- weight_dict = {"loss_ce": self.config.cls_loss_coefficient, "loss_bbox": self.config.bbox_loss_coefficient}
- weight_dict["loss_giou"] = self.config.giou_loss_coefficient
- if self.config.auxiliary_loss:
- aux_weight_dict = {}
- for i in range(self.config.decoder_layers - 1):
- aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
- weight_dict.update(aux_weight_dict)
- loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
-
- if not return_dict:
- if auxiliary_outputs is not None:
- output = (logits, pred_boxes) + auxiliary_outputs + outputs
- else:
- output = (logits, pred_boxes) + outputs
- return ((loss, loss_dict) + output) if loss is not None else output
-
- return ConditionalDetrObjectDetectionOutput(
- loss=loss,
- loss_dict=loss_dict,
- logits=logits,
- pred_boxes=pred_boxes,
- auxiliary_outputs=auxiliary_outputs,
- last_hidden_state=outputs.last_hidden_state,
- decoder_hidden_states=outputs.decoder_hidden_states,
- decoder_attentions=outputs.decoder_attentions,
- cross_attentions=outputs.cross_attentions,
- encoder_last_hidden_state=outputs.encoder_last_hidden_state,
- encoder_hidden_states=outputs.encoder_hidden_states,
- encoder_attentions=outputs.encoder_attentions,
- )
-
-
-@add_start_docstrings(
- """
- CONDITIONAL_DETR Model (consisting of a backbone and encoder-decoder Transformer) with a segmentation head on top,
- for tasks such as COCO panoptic.
-
- """,
- CONDITIONAL_DETR_START_DOCSTRING,
-)
-class ConditionalDetrForSegmentation(ConditionalDetrPreTrainedModel):
- def __init__(self, config: ConditionalDetrConfig):
- super().__init__(config)
-
- # object detection model
- self.conditional_detr = ConditionalDetrForObjectDetection(config)
-
- # segmentation head
- hidden_size, number_of_heads = config.d_model, config.encoder_attention_heads
- intermediate_channel_sizes = self.conditional_detr.model.backbone.conv_encoder.intermediate_channel_sizes
-
- self.mask_head = ConditionalDetrMaskHeadSmallConv(
- hidden_size + number_of_heads, intermediate_channel_sizes[::-1][-3:], hidden_size
- )
-
- self.bbox_attention = ConditionalDetrMHAttentionMap(
- hidden_size, hidden_size, number_of_heads, dropout=0.0, std=config.init_xavier_std
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONDITIONAL_DETR_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=ConditionalDetrSegmentationOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.FloatTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[List[dict]] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], ConditionalDetrSegmentationOutput]:
- r"""
- labels (`List[Dict]` of len `(batch_size,)`, *optional*):
- Labels for computing the bipartite matching loss, DICE/F-1 loss and Focal loss. List of dicts, each
- dictionary containing at least the following 3 keys: 'class_labels', 'boxes' and 'masks' (the class labels,
- bounding boxes and segmentation masks of an image in the batch respectively). The class labels themselves
- should be a `torch.LongTensor` of len `(number of bounding boxes in the image,)`, the boxes a
- `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)` and the masks a
- `torch.FloatTensor` of shape `(number of bounding boxes in the image, height, width)`.
-
- Returns:
-
- Examples:
-
- ```python
- >>> import io
- >>> import requests
- >>> from PIL import Image
- >>> import torch
- >>> import numpy
-
- >>> from transformers import (
- ... AutoImageProcessor,
- ... ConditionalDetrConfig,
- ... ConditionalDetrForSegmentation,
- ... )
- >>> from transformers.image_transforms import rgb_to_id
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50")
-
- >>> # randomly initialize all weights of the model
- >>> config = ConditionalDetrConfig()
- >>> model = ConditionalDetrForSegmentation(config)
-
- >>> # prepare image for the model
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> # forward pass
- >>> outputs = model(**inputs)
-
- >>> # Use the `post_process_panoptic_segmentation` method of the `image_processor` to retrieve post-processed panoptic segmentation maps
- >>> # Segmentation results are returned as a list of dictionaries
- >>> result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[(300, 500)])
- >>> # A tensor of shape (height, width) where each value denotes a segment id, filled with -1 if no segment is found
- >>> panoptic_seg = result[0]["segmentation"]
- >>> # Get prediction score and segment_id to class_id mapping of each segment
- >>> panoptic_segments_info = result[0]["segments_info"]
- ```"""
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- batch_size, num_channels, height, width = pixel_values.shape
- device = pixel_values.device
-
- if pixel_mask is None:
- pixel_mask = torch.ones((batch_size, height, width), device=device)
-
- # First, get list of feature maps and object_queries
- features, object_queries_list = self.conditional_detr.model.backbone(pixel_values, pixel_mask=pixel_mask)
-
- # Second, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
- feature_map, mask = features[-1]
- batch_size, num_channels, height, width = feature_map.shape
- projected_feature_map = self.conditional_detr.model.input_projection(feature_map)
-
- # Third, flatten the feature map + object_queries of shape NxCxHxW to NxCxHW, and permute it to NxHWxC
- # In other words, turn their shape into (batch_size, sequence_length, hidden_size)
- flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1)
- object_queries = object_queries_list[-1].flatten(2).permute(0, 2, 1)
-
- flattened_mask = mask.flatten(1)
-
- # Fourth, sent flattened_features + flattened_mask + object_queries through encoder
- # flattened_features is a Tensor of shape (batch_size, heigth*width, hidden_size)
- # flattened_mask is a Tensor of shape (batch_size, heigth*width)
- if encoder_outputs is None:
- encoder_outputs = self.conditional_detr.model.encoder(
- inputs_embeds=flattened_features,
- attention_mask=flattened_mask,
- object_queries=object_queries,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
- elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
- encoder_outputs = BaseModelOutput(
- last_hidden_state=encoder_outputs[0],
- hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
- attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
- )
-
- # Fifth, sent query embeddings + object_queries through the decoder (which is conditioned on the encoder output)
- query_position_embeddings = self.conditional_detr.model.query_position_embeddings.weight.unsqueeze(0).repeat(
- batch_size, 1, 1
- )
- queries = torch.zeros_like(query_position_embeddings)
-
- # decoder outputs consists of (dec_features, dec_hidden, dec_attn)
- decoder_outputs = self.conditional_detr.model.decoder(
- inputs_embeds=queries,
- attention_mask=None,
- object_queries=object_queries,
- query_position_embeddings=query_position_embeddings,
- encoder_hidden_states=encoder_outputs[0],
- encoder_attention_mask=flattened_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = decoder_outputs[0]
-
- # Sixth, compute logits, pred_boxes and pred_masks
- logits = self.conditional_detr.class_labels_classifier(sequence_output)
- pred_boxes = self.conditional_detr.bbox_predictor(sequence_output).sigmoid()
-
- memory = encoder_outputs[0].permute(0, 2, 1).view(batch_size, self.config.d_model, height, width)
- mask = flattened_mask.view(batch_size, height, width)
-
- # FIXME h_boxes takes the last one computed, keep this in mind
- # important: we need to reverse the mask, since in the original implementation the mask works reversed
- # bbox_mask is of shape (batch_size, num_queries, number_of_attention_heads in bbox_attention, height/32, width/32)
- bbox_mask = self.bbox_attention(sequence_output, memory, mask=~mask)
-
- seg_masks = self.mask_head(projected_feature_map, bbox_mask, [features[2][0], features[1][0], features[0][0]])
-
- pred_masks = seg_masks.view(
- batch_size, self.conditional_detr.config.num_queries, seg_masks.shape[-2], seg_masks.shape[-1]
- )
-
- loss, loss_dict, auxiliary_outputs = None, None, None
- if labels is not None:
- # First: create the matcher
- matcher = ConditionalDetrHungarianMatcher(
- class_cost=self.config.class_cost, bbox_cost=self.config.bbox_cost, giou_cost=self.config.giou_cost
- )
- # Second: create the criterion
- losses = ["labels", "boxes", "cardinality", "masks"]
- criterion = ConditionalDetrLoss(
- matcher=matcher,
- num_classes=self.config.num_labels,
- focal_alpha=self.config.focal_alpha,
- losses=losses,
- )
- criterion.to(self.device)
- # Third: compute the losses, based on outputs and labels
- outputs_loss = {}
- outputs_loss["logits"] = logits
- outputs_loss["pred_boxes"] = pred_boxes
- outputs_loss["pred_masks"] = pred_masks
- if self.config.auxiliary_loss:
- intermediate = decoder_outputs.intermediate_hidden_states if return_dict else decoder_outputs[-1]
- outputs_class = self.conditional_detr.class_labels_classifier(intermediate)
- outputs_coord = self.conditional_detr.bbox_predictor(intermediate).sigmoid()
- auxiliary_outputs = self.conditional_detr._set_aux_loss(outputs_class, outputs_coord)
- outputs_loss["auxiliary_outputs"] = auxiliary_outputs
-
- loss_dict = criterion(outputs_loss, labels)
- # Fourth: compute total loss, as a weighted sum of the various losses
- weight_dict = {"loss_ce": 1, "loss_bbox": self.config.bbox_loss_coefficient}
- weight_dict["loss_giou"] = self.config.giou_loss_coefficient
- weight_dict["loss_mask"] = self.config.mask_loss_coefficient
- weight_dict["loss_dice"] = self.config.dice_loss_coefficient
- if self.config.auxiliary_loss:
- aux_weight_dict = {}
- for i in range(self.config.decoder_layers - 1):
- aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
- weight_dict.update(aux_weight_dict)
- loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
-
- if not return_dict:
- if auxiliary_outputs is not None:
- output = (logits, pred_boxes, pred_masks) + auxiliary_outputs + decoder_outputs + encoder_outputs
- else:
- output = (logits, pred_boxes, pred_masks) + decoder_outputs + encoder_outputs
- return ((loss, loss_dict) + output) if loss is not None else output
-
- return ConditionalDetrSegmentationOutput(
- loss=loss,
- loss_dict=loss_dict,
- logits=logits,
- pred_boxes=pred_boxes,
- pred_masks=pred_masks,
- auxiliary_outputs=auxiliary_outputs,
- last_hidden_state=decoder_outputs.last_hidden_state,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- )
-
-
-def _expand(tensor, length: int):
- return tensor.unsqueeze(1).repeat(1, int(length), 1, 1, 1).flatten(0, 1)
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrMaskHeadSmallConv with Detr->ConditionalDetr
-class ConditionalDetrMaskHeadSmallConv(nn.Module):
- """
- Simple convolutional head, using group norm. Upsampling is done using a FPN approach
- """
-
- def __init__(self, dim, fpn_dims, context_dim):
- super().__init__()
-
- if dim % 8 != 0:
- raise ValueError(
- "The hidden_size + number of attention heads must be divisible by 8 as the number of groups in"
- " GroupNorm is set to 8"
- )
-
- inter_dims = [dim, context_dim // 2, context_dim // 4, context_dim // 8, context_dim // 16, context_dim // 64]
-
- self.lay1 = nn.Conv2d(dim, dim, 3, padding=1)
- self.gn1 = nn.GroupNorm(8, dim)
- self.lay2 = nn.Conv2d(dim, inter_dims[1], 3, padding=1)
- self.gn2 = nn.GroupNorm(min(8, inter_dims[1]), inter_dims[1])
- self.lay3 = nn.Conv2d(inter_dims[1], inter_dims[2], 3, padding=1)
- self.gn3 = nn.GroupNorm(min(8, inter_dims[2]), inter_dims[2])
- self.lay4 = nn.Conv2d(inter_dims[2], inter_dims[3], 3, padding=1)
- self.gn4 = nn.GroupNorm(min(8, inter_dims[3]), inter_dims[3])
- self.lay5 = nn.Conv2d(inter_dims[3], inter_dims[4], 3, padding=1)
- self.gn5 = nn.GroupNorm(min(8, inter_dims[4]), inter_dims[4])
- self.out_lay = nn.Conv2d(inter_dims[4], 1, 3, padding=1)
-
- self.dim = dim
-
- self.adapter1 = nn.Conv2d(fpn_dims[0], inter_dims[1], 1)
- self.adapter2 = nn.Conv2d(fpn_dims[1], inter_dims[2], 1)
- self.adapter3 = nn.Conv2d(fpn_dims[2], inter_dims[3], 1)
-
- for m in self.modules():
- if isinstance(m, nn.Conv2d):
- nn.init.kaiming_uniform_(m.weight, a=1)
- nn.init.constant_(m.bias, 0)
-
- def forward(self, x: Tensor, bbox_mask: Tensor, fpns: List[Tensor]):
- # here we concatenate x, the projected feature map, of shape (batch_size, d_model, heigth/32, width/32) with
- # the bbox_mask = the attention maps of shape (batch_size, n_queries, n_heads, height/32, width/32).
- # We expand the projected feature map to match the number of heads.
- x = torch.cat([_expand(x, bbox_mask.shape[1]), bbox_mask.flatten(0, 1)], 1)
-
- x = self.lay1(x)
- x = self.gn1(x)
- x = nn.functional.relu(x)
- x = self.lay2(x)
- x = self.gn2(x)
- x = nn.functional.relu(x)
-
- cur_fpn = self.adapter1(fpns[0])
- if cur_fpn.size(0) != x.size(0):
- cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
- x = cur_fpn + nn.functional.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
- x = self.lay3(x)
- x = self.gn3(x)
- x = nn.functional.relu(x)
-
- cur_fpn = self.adapter2(fpns[1])
- if cur_fpn.size(0) != x.size(0):
- cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
- x = cur_fpn + nn.functional.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
- x = self.lay4(x)
- x = self.gn4(x)
- x = nn.functional.relu(x)
-
- cur_fpn = self.adapter3(fpns[2])
- if cur_fpn.size(0) != x.size(0):
- cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
- x = cur_fpn + nn.functional.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
- x = self.lay5(x)
- x = self.gn5(x)
- x = nn.functional.relu(x)
-
- x = self.out_lay(x)
- return x
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrMHAttentionMap with Detr->ConditionalDetr
-class ConditionalDetrMHAttentionMap(nn.Module):
- """This is a 2D attention module, which only returns the attention softmax (no multiplication by value)"""
-
- def __init__(self, query_dim, hidden_dim, num_heads, dropout=0.0, bias=True, std=None):
- super().__init__()
- self.num_heads = num_heads
- self.hidden_dim = hidden_dim
- self.dropout = nn.Dropout(dropout)
-
- self.q_linear = nn.Linear(query_dim, hidden_dim, bias=bias)
- self.k_linear = nn.Linear(query_dim, hidden_dim, bias=bias)
-
- self.normalize_fact = float(hidden_dim / self.num_heads) ** -0.5
-
- def forward(self, q, k, mask: Optional[Tensor] = None):
- q = self.q_linear(q)
- k = nn.functional.conv2d(k, self.k_linear.weight.unsqueeze(-1).unsqueeze(-1), self.k_linear.bias)
- queries_per_head = q.view(q.shape[0], q.shape[1], self.num_heads, self.hidden_dim // self.num_heads)
- keys_per_head = k.view(k.shape[0], self.num_heads, self.hidden_dim // self.num_heads, k.shape[-2], k.shape[-1])
- weights = torch.einsum("bqnc,bnchw->bqnhw", queries_per_head * self.normalize_fact, keys_per_head)
-
- if mask is not None:
- weights.masked_fill_(mask.unsqueeze(1).unsqueeze(1), torch.finfo(weights.dtype).min)
- weights = nn.functional.softmax(weights.flatten(2), dim=-1).view(weights.size())
- weights = self.dropout(weights)
- return weights
-
-
-# Copied from transformers.models.detr.modeling_detr.dice_loss
-def dice_loss(inputs, targets, num_boxes):
- """
- Compute the DICE loss, similar to generalized IOU for masks
-
- Args:
- inputs: A float tensor of arbitrary shape.
- The predictions for each example.
- targets: A float tensor with the same shape as inputs. Stores the binary
- classification label for each element in inputs (0 for the negative class and 1 for the positive
- class).
- """
- inputs = inputs.sigmoid()
- inputs = inputs.flatten(1)
- numerator = 2 * (inputs * targets).sum(1)
- denominator = inputs.sum(-1) + targets.sum(-1)
- loss = 1 - (numerator + 1) / (denominator + 1)
- return loss.sum() / num_boxes
-
-
-# Copied from transformers.models.detr.modeling_detr.sigmoid_focal_loss
-def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2):
- """
- Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
-
- Args:
- inputs (`torch.FloatTensor` of arbitrary shape):
- The predictions for each example.
- targets (`torch.FloatTensor` with the same shape as `inputs`)
- A tensor storing the binary classification label for each element in the `inputs` (0 for the negative class
- and 1 for the positive class).
- alpha (`float`, *optional*, defaults to `0.25`):
- Optional weighting factor in the range (0,1) to balance positive vs. negative examples.
- gamma (`int`, *optional*, defaults to `2`):
- Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples.
-
- Returns:
- Loss tensor
- """
- prob = inputs.sigmoid()
- ce_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
- # add modulating factor
- p_t = prob * targets + (1 - prob) * (1 - targets)
- loss = ce_loss * ((1 - p_t) ** gamma)
-
- if alpha >= 0:
- alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
- loss = alpha_t * loss
-
- return loss.mean(1).sum() / num_boxes
-
-
-class ConditionalDetrLoss(nn.Module):
- """
- This class computes the losses for ConditionalDetrForObjectDetection/ConditionalDetrForSegmentation. The process
- happens in two steps: 1) we compute hungarian assignment between ground truth boxes and the outputs of the model 2)
- we supervise each pair of matched ground-truth / prediction (supervise class and box).
-
- Args:
- matcher (`ConditionalDetrHungarianMatcher`):
- Module able to compute a matching between targets and proposals.
- num_classes (`int`):
- Number of object categories, omitting the special no-object category.
- focal_alpha (`float`):
- Alpha parameter in focal loss.
- losses (`List[str]`):
- List of all the losses to be applied. See `get_loss` for a list of all available losses.
- """
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss.__init__
- def __init__(self, matcher, num_classes, focal_alpha, losses):
- super().__init__()
- self.matcher = matcher
- self.num_classes = num_classes
- self.focal_alpha = focal_alpha
- self.losses = losses
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss.loss_labels
- def loss_labels(self, outputs, targets, indices, num_boxes):
- """
- Classification loss (Binary focal loss) targets dicts must contain the key "class_labels" containing a tensor
- of dim [nb_target_boxes]
- """
- if "logits" not in outputs:
- raise KeyError("No logits were found in the outputs")
- source_logits = outputs["logits"]
-
- idx = self._get_source_permutation_idx(indices)
- target_classes_o = torch.cat([t["class_labels"][J] for t, (_, J) in zip(targets, indices)])
- target_classes = torch.full(
- source_logits.shape[:2], self.num_classes, dtype=torch.int64, device=source_logits.device
- )
- target_classes[idx] = target_classes_o
-
- target_classes_onehot = torch.zeros(
- [source_logits.shape[0], source_logits.shape[1], source_logits.shape[2] + 1],
- dtype=source_logits.dtype,
- layout=source_logits.layout,
- device=source_logits.device,
- )
- target_classes_onehot.scatter_(2, target_classes.unsqueeze(-1), 1)
-
- target_classes_onehot = target_classes_onehot[:, :, :-1]
- loss_ce = (
- sigmoid_focal_loss(source_logits, target_classes_onehot, num_boxes, alpha=self.focal_alpha, gamma=2)
- * source_logits.shape[1]
- )
- losses = {"loss_ce": loss_ce}
-
- return losses
-
- @torch.no_grad()
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss.loss_cardinality
- def loss_cardinality(self, outputs, targets, indices, num_boxes):
- """
- Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes.
-
- This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients.
- """
- logits = outputs["logits"]
- device = logits.device
- target_lengths = torch.as_tensor([len(v["class_labels"]) for v in targets], device=device)
- # Count the number of predictions that are NOT "no-object" (which is the last class)
- card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1)
- card_err = nn.functional.l1_loss(card_pred.float(), target_lengths.float())
- losses = {"cardinality_error": card_err}
- return losses
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss.loss_boxes
- def loss_boxes(self, outputs, targets, indices, num_boxes):
- """
- Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss.
-
- Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes
- are expected in format (center_x, center_y, w, h), normalized by the image size.
- """
- if "pred_boxes" not in outputs:
- raise KeyError("No predicted boxes found in outputs")
- idx = self._get_source_permutation_idx(indices)
- source_boxes = outputs["pred_boxes"][idx]
- target_boxes = torch.cat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0)
-
- loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none")
-
- losses = {}
- losses["loss_bbox"] = loss_bbox.sum() / num_boxes
-
- loss_giou = 1 - torch.diag(
- generalized_box_iou(center_to_corners_format(source_boxes), center_to_corners_format(target_boxes))
- )
- losses["loss_giou"] = loss_giou.sum() / num_boxes
- return losses
-
- # Copied from transformers.models.detr.modeling_detr.DetrLoss.loss_masks
- def loss_masks(self, outputs, targets, indices, num_boxes):
- """
- Compute the losses related to the masks: the focal loss and the dice loss.
-
- Targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w].
- """
- if "pred_masks" not in outputs:
- raise KeyError("No predicted masks found in outputs")
-
- source_idx = self._get_source_permutation_idx(indices)
- target_idx = self._get_target_permutation_idx(indices)
- source_masks = outputs["pred_masks"]
- source_masks = source_masks[source_idx]
- masks = [t["masks"] for t in targets]
- # TODO use valid to mask invalid areas due to padding in loss
- target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
- target_masks = target_masks.to(source_masks)
- target_masks = target_masks[target_idx]
-
- # upsample predictions to the target size
- source_masks = nn.functional.interpolate(
- source_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False
- )
- source_masks = source_masks[:, 0].flatten(1)
-
- target_masks = target_masks.flatten(1)
- target_masks = target_masks.view(source_masks.shape)
- losses = {
- "loss_mask": sigmoid_focal_loss(source_masks, target_masks, num_boxes),
- "loss_dice": dice_loss(source_masks, target_masks, num_boxes),
- }
- return losses
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss._get_source_permutation_idx
- def _get_source_permutation_idx(self, indices):
- # permute predictions following indices
- batch_idx = torch.cat([torch.full_like(source, i) for i, (source, _) in enumerate(indices)])
- source_idx = torch.cat([source for (source, _) in indices])
- return batch_idx, source_idx
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss._get_target_permutation_idx
- def _get_target_permutation_idx(self, indices):
- # permute targets following indices
- batch_idx = torch.cat([torch.full_like(target, i) for i, (_, target) in enumerate(indices)])
- target_idx = torch.cat([target for (_, target) in indices])
- return batch_idx, target_idx
-
- # Copied from transformers.models.detr.modeling_detr.DetrLoss.get_loss
- def get_loss(self, loss, outputs, targets, indices, num_boxes):
- loss_map = {
- "labels": self.loss_labels,
- "cardinality": self.loss_cardinality,
- "boxes": self.loss_boxes,
- "masks": self.loss_masks,
- }
- if loss not in loss_map:
- raise ValueError(f"Loss {loss} not supported")
- return loss_map[loss](outputs, targets, indices, num_boxes)
-
- # Copied from transformers.models.detr.modeling_detr.DetrLoss.forward
- def forward(self, outputs, targets):
- """
- This performs the loss computation.
-
- Args:
- outputs (`dict`, *optional*):
- Dictionary of tensors, see the output specification of the model for the format.
- targets (`List[dict]`, *optional*):
- List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the
- losses applied, see each loss' doc.
- """
- outputs_without_aux = {k: v for k, v in outputs.items() if k != "auxiliary_outputs"}
-
- # Retrieve the matching between the outputs of the last layer and the targets
- indices = self.matcher(outputs_without_aux, targets)
-
- # Compute the average number of target boxes across all nodes, for normalization purposes
- num_boxes = sum(len(t["class_labels"]) for t in targets)
- num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
-
- world_size = 1
- if is_accelerate_available():
- if PartialState._shared_state != {}:
- num_boxes = reduce(num_boxes)
- world_size = PartialState().num_processes
- num_boxes = torch.clamp(num_boxes / world_size, min=1).item()
-
- # Compute all the requested losses
- losses = {}
- for loss in self.losses:
- losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))
-
- # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
- if "auxiliary_outputs" in outputs:
- for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]):
- indices = self.matcher(auxiliary_outputs, targets)
- for loss in self.losses:
- if loss == "masks":
- # Intermediate masks losses are too costly to compute, we ignore them.
- continue
- l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes)
- l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
- losses.update(l_dict)
-
- return losses
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrMLPPredictionHead with Detr->ConditionalDetr
-class ConditionalDetrMLPPredictionHead(nn.Module):
- """
- Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates,
- height and width of a bounding box w.r.t. an image.
-
- Copied from https://github.com/facebookresearch/detr/blob/master/models/detr.py
-
- """
-
- def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
- super().__init__()
- self.num_layers = num_layers
- h = [hidden_dim] * (num_layers - 1)
- self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
-
- def forward(self, x):
- for i, layer in enumerate(self.layers):
- x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
- return x
-
-
-# Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrHungarianMatcher with DeformableDetr->ConditionalDetr
-class ConditionalDetrHungarianMatcher(nn.Module):
- """
- This class computes an assignment between the targets and the predictions of the network.
-
- For efficiency reasons, the targets don't include the no_object. Because of this, in general, there are more
- predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, while the others are
- un-matched (and thus treated as non-objects).
-
- Args:
- class_cost:
- The relative weight of the classification error in the matching cost.
- bbox_cost:
- The relative weight of the L1 error of the bounding box coordinates in the matching cost.
- giou_cost:
- The relative weight of the giou loss of the bounding box in the matching cost.
- """
-
- def __init__(self, class_cost: float = 1, bbox_cost: float = 1, giou_cost: float = 1):
- super().__init__()
- requires_backends(self, ["scipy"])
-
- self.class_cost = class_cost
- self.bbox_cost = bbox_cost
- self.giou_cost = giou_cost
- if class_cost == 0 and bbox_cost == 0 and giou_cost == 0:
- raise ValueError("All costs of the Matcher can't be 0")
-
- @torch.no_grad()
- def forward(self, outputs, targets):
- """
- Args:
- outputs (`dict`):
- A dictionary that contains at least these entries:
- * "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
- * "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates.
- targets (`List[dict]`):
- A list of targets (len(targets) = batch_size), where each target is a dict containing:
- * "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of
- ground-truth
- objects in the target) containing the class labels
- * "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates.
-
- Returns:
- `List[Tuple]`: A list of size `batch_size`, containing tuples of (index_i, index_j) where:
- - index_i is the indices of the selected predictions (in order)
- - index_j is the indices of the corresponding selected targets (in order)
- For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
- """
- batch_size, num_queries = outputs["logits"].shape[:2]
-
- # We flatten to compute the cost matrices in a batch
- out_prob = outputs["logits"].flatten(0, 1).sigmoid() # [batch_size * num_queries, num_classes]
- out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4]
-
- # Also concat the target labels and boxes
- target_ids = torch.cat([v["class_labels"] for v in targets])
- target_bbox = torch.cat([v["boxes"] for v in targets])
-
- # Compute the classification cost.
- alpha = 0.25
- gamma = 2.0
- neg_cost_class = (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log())
- pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log())
- class_cost = pos_cost_class[:, target_ids] - neg_cost_class[:, target_ids]
-
- # Compute the L1 cost between boxes
- bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)
-
- # Compute the giou cost between boxes
- giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))
-
- # Final cost matrix
- cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
- cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()
-
- sizes = [len(v["boxes"]) for v in targets]
- indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]
- return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
-
-
-# Copied from transformers.models.detr.modeling_detr._upcast
-def _upcast(t: Tensor) -> Tensor:
- # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
- if t.is_floating_point():
- return t if t.dtype in (torch.float32, torch.float64) else t.float()
- else:
- return t if t.dtype in (torch.int32, torch.int64) else t.int()
-
-
-# Copied from transformers.models.detr.modeling_detr.box_area
-def box_area(boxes: Tensor) -> Tensor:
- """
- Computes the area of a set of bounding boxes, which are specified by its (x1, y1, x2, y2) coordinates.
-
- Args:
- boxes (`torch.FloatTensor` of shape `(number_of_boxes, 4)`):
- Boxes for which the area will be computed. They are expected to be in (x1, y1, x2, y2) format with `0 <= x1
- < x2` and `0 <= y1 < y2`.
-
- Returns:
- `torch.FloatTensor`: a tensor containing the area for each box.
- """
- boxes = _upcast(boxes)
- return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
-
-
-# Copied from transformers.models.detr.modeling_detr.box_iou
-def box_iou(boxes1, boxes2):
- area1 = box_area(boxes1)
- area2 = box_area(boxes2)
-
- left_top = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
- right_bottom = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
-
- width_height = (right_bottom - left_top).clamp(min=0) # [N,M,2]
- inter = width_height[:, :, 0] * width_height[:, :, 1] # [N,M]
-
- union = area1[:, None] + area2 - inter
-
- iou = inter / union
- return iou, union
-
-
-# Copied from transformers.models.detr.modeling_detr.generalized_box_iou
-def generalized_box_iou(boxes1, boxes2):
- """
- Generalized IoU from https://giou.stanford.edu/. The boxes should be in [x0, y0, x1, y1] (corner) format.
-
- Returns:
- `torch.FloatTensor`: a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2)
- """
- # degenerate boxes gives inf / nan results
- # so do an early check
- if not (boxes1[:, 2:] >= boxes1[:, :2]).all():
- raise ValueError(f"boxes1 must be in [x0, y0, x1, y1] (corner) format, but got {boxes1}")
- if not (boxes2[:, 2:] >= boxes2[:, :2]).all():
- raise ValueError(f"boxes2 must be in [x0, y0, x1, y1] (corner) format, but got {boxes2}")
- iou, union = box_iou(boxes1, boxes2)
-
- top_left = torch.min(boxes1[:, None, :2], boxes2[:, :2])
- bottom_right = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
-
- width_height = (bottom_right - top_left).clamp(min=0) # [N,M,2]
- area = width_height[:, :, 0] * width_height[:, :, 1]
-
- return iou - (area - union) / area
-
-
-# Copied from transformers.models.detr.modeling_detr._max_by_axis
-def _max_by_axis(the_list):
- # type: (List[List[int]]) -> List[int]
- maxes = the_list[0]
- for sublist in the_list[1:]:
- for index, item in enumerate(sublist):
- maxes[index] = max(maxes[index], item)
- return maxes
-
-
-# Copied from transformers.models.detr.modeling_detr.NestedTensor
-class NestedTensor(object):
- def __init__(self, tensors, mask: Optional[Tensor]):
- self.tensors = tensors
- self.mask = mask
-
- def to(self, device):
- cast_tensor = self.tensors.to(device)
- mask = self.mask
- if mask is not None:
- cast_mask = mask.to(device)
- else:
- cast_mask = None
- return NestedTensor(cast_tensor, cast_mask)
-
- def decompose(self):
- return self.tensors, self.mask
-
- def __repr__(self):
- return str(self.tensors)
-
-
-# Copied from transformers.models.detr.modeling_detr.nested_tensor_from_tensor_list
-def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
- if tensor_list[0].ndim == 3:
- max_size = _max_by_axis([list(img.shape) for img in tensor_list])
- batch_shape = [len(tensor_list)] + max_size
- batch_size, num_channels, height, width = batch_shape
- dtype = tensor_list[0].dtype
- device = tensor_list[0].device
- tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
- mask = torch.ones((batch_size, height, width), dtype=torch.bool, device=device)
- for img, pad_img, m in zip(tensor_list, tensor, mask):
- pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
- m[: img.shape[1], : img.shape[2]] = False
- else:
- raise ValueError("Only 3-dimensional tensors are supported")
- return NestedTensor(tensor, mask)
diff --git a/transformers/models/convbert/__init__.py b/transformers/models/convbert/__init__.py
deleted file mode 100644
index f1b19a949abbef25ed52f7e0d0d1efd6c2410d12..0000000000000000000000000000000000000000
--- a/transformers/models/convbert/__init__.py
+++ /dev/null
@@ -1,130 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_tf_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_convbert": ["CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "ConvBertConfig", "ConvBertOnnxConfig"],
- "tokenization_convbert": ["ConvBertTokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_convbert_fast"] = ["ConvBertTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_convbert"] = [
- "CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "ConvBertForMaskedLM",
- "ConvBertForMultipleChoice",
- "ConvBertForQuestionAnswering",
- "ConvBertForSequenceClassification",
- "ConvBertForTokenClassification",
- "ConvBertLayer",
- "ConvBertModel",
- "ConvBertPreTrainedModel",
- "load_tf_weights_in_convbert",
- ]
-
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_convbert"] = [
- "TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFConvBertForMaskedLM",
- "TFConvBertForMultipleChoice",
- "TFConvBertForQuestionAnswering",
- "TFConvBertForSequenceClassification",
- "TFConvBertForTokenClassification",
- "TFConvBertLayer",
- "TFConvBertModel",
- "TFConvBertPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_convbert import CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, ConvBertConfig, ConvBertOnnxConfig
- from .tokenization_convbert import ConvBertTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_convbert_fast import ConvBertTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_convbert import (
- CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
- ConvBertForMaskedLM,
- ConvBertForMultipleChoice,
- ConvBertForQuestionAnswering,
- ConvBertForSequenceClassification,
- ConvBertForTokenClassification,
- ConvBertLayer,
- ConvBertModel,
- ConvBertPreTrainedModel,
- load_tf_weights_in_convbert,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_convbert import (
- TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFConvBertForMaskedLM,
- TFConvBertForMultipleChoice,
- TFConvBertForQuestionAnswering,
- TFConvBertForSequenceClassification,
- TFConvBertForTokenClassification,
- TFConvBertLayer,
- TFConvBertModel,
- TFConvBertPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/convbert/__pycache__/__init__.cpython-310.pyc b/transformers/models/convbert/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 74c0ee722d391b79a8be74d28dd5d4eb72914b5a..0000000000000000000000000000000000000000
Binary files a/transformers/models/convbert/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convbert/__pycache__/configuration_convbert.cpython-310.pyc b/transformers/models/convbert/__pycache__/configuration_convbert.cpython-310.pyc
deleted file mode 100644
index 1d33da44ed1b52495e0b293986a2e6ca4ed1b1d1..0000000000000000000000000000000000000000
Binary files a/transformers/models/convbert/__pycache__/configuration_convbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convbert/__pycache__/convert_convbert_original_tf1_checkpoint_to_pytorch_and_tf2.cpython-310.pyc b/transformers/models/convbert/__pycache__/convert_convbert_original_tf1_checkpoint_to_pytorch_and_tf2.cpython-310.pyc
deleted file mode 100644
index 82375b5df01eea5376e7c6974d202aefbe69de34..0000000000000000000000000000000000000000
Binary files a/transformers/models/convbert/__pycache__/convert_convbert_original_tf1_checkpoint_to_pytorch_and_tf2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convbert/__pycache__/modeling_convbert.cpython-310.pyc b/transformers/models/convbert/__pycache__/modeling_convbert.cpython-310.pyc
deleted file mode 100644
index d11988c715ecc4fc3ac0c2b3b21cb87152cd4304..0000000000000000000000000000000000000000
Binary files a/transformers/models/convbert/__pycache__/modeling_convbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convbert/__pycache__/modeling_tf_convbert.cpython-310.pyc b/transformers/models/convbert/__pycache__/modeling_tf_convbert.cpython-310.pyc
deleted file mode 100644
index a63e72236643ae6f16c681180f0aa25f48f2c3d7..0000000000000000000000000000000000000000
Binary files a/transformers/models/convbert/__pycache__/modeling_tf_convbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convbert/__pycache__/tokenization_convbert.cpython-310.pyc b/transformers/models/convbert/__pycache__/tokenization_convbert.cpython-310.pyc
deleted file mode 100644
index 7d5aa4f49fc6d9149f7249dc6ef72f24371bfb86..0000000000000000000000000000000000000000
Binary files a/transformers/models/convbert/__pycache__/tokenization_convbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convbert/__pycache__/tokenization_convbert_fast.cpython-310.pyc b/transformers/models/convbert/__pycache__/tokenization_convbert_fast.cpython-310.pyc
deleted file mode 100644
index 2e3680a11bfbc97be65e56e61cb317c1bc06d3de..0000000000000000000000000000000000000000
Binary files a/transformers/models/convbert/__pycache__/tokenization_convbert_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convbert/configuration_convbert.py b/transformers/models/convbert/configuration_convbert.py
deleted file mode 100644
index d309ca396baffcfa707ecd599f7b4d280dc348a9..0000000000000000000000000000000000000000
--- a/transformers/models/convbert/configuration_convbert.py
+++ /dev/null
@@ -1,160 +0,0 @@
-# coding=utf-8
-# Copyright The HuggingFace team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" ConvBERT model configuration"""
-
-from collections import OrderedDict
-from typing import Mapping
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class ConvBertConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`ConvBertModel`]. It is used to instantiate an
- ConvBERT model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the ConvBERT
- [YituTech/conv-bert-base](https://huggingface.co/YituTech/conv-bert-base) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the ConvBERT model. Defines the number of different tokens that can be represented by
- the `inputs_ids` passed when calling [`ConvBertModel`] or [`TFConvBertModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 2):
- The vocabulary size of the `token_type_ids` passed when calling [`ConvBertModel`] or [`TFConvBertModel`].
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- head_ratio (`int`, *optional*, defaults to 2):
- Ratio gamma to reduce the number of attention heads.
- num_groups (`int`, *optional*, defaults to 1):
- The number of groups for grouped linear layers for ConvBert model
- conv_kernel_size (`int`, *optional*, defaults to 9):
- The size of the convolutional kernel.
- classifier_dropout (`float`, *optional*):
- The dropout ratio for the classification head.
-
- Example:
-
- ```python
- >>> from transformers import ConvBertConfig, ConvBertModel
-
- >>> # Initializing a ConvBERT convbert-base-uncased style configuration
- >>> configuration = ConvBertConfig()
-
- >>> # Initializing a model (with random weights) from the convbert-base-uncased style configuration
- >>> model = ConvBertModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "convbert"
-
- def __init__(
- self,
- vocab_size=30522,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=2,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- pad_token_id=1,
- bos_token_id=0,
- eos_token_id=2,
- embedding_size=768,
- head_ratio=2,
- conv_kernel_size=9,
- num_groups=1,
- classifier_dropout=None,
- **kwargs,
- ):
- super().__init__(
- pad_token_id=pad_token_id,
- bos_token_id=bos_token_id,
- eos_token_id=eos_token_id,
- **kwargs,
- )
-
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.embedding_size = embedding_size
- self.head_ratio = head_ratio
- self.conv_kernel_size = conv_kernel_size
- self.num_groups = num_groups
- self.classifier_dropout = classifier_dropout
-
-
-# Copied from transformers.models.bert.configuration_bert.BertOnnxConfig
-class ConvBertOnnxConfig(OnnxConfig):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- if self.task == "multiple-choice":
- dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
- else:
- dynamic_axis = {0: "batch", 1: "sequence"}
- return OrderedDict(
- [
- ("input_ids", dynamic_axis),
- ("attention_mask", dynamic_axis),
- ("token_type_ids", dynamic_axis),
- ]
- )
diff --git a/transformers/models/convbert/convert_convbert_original_tf1_checkpoint_to_pytorch_and_tf2.py b/transformers/models/convbert/convert_convbert_original_tf1_checkpoint_to_pytorch_and_tf2.py
deleted file mode 100644
index 3d4ff779874b30b0c094c596cedaca597e03ed36..0000000000000000000000000000000000000000
--- a/transformers/models/convbert/convert_convbert_original_tf1_checkpoint_to_pytorch_and_tf2.py
+++ /dev/null
@@ -1,57 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert ConvBERT checkpoint."""
-
-import argparse
-
-from transformers import ConvBertConfig, ConvBertModel, TFConvBertModel, load_tf_weights_in_convbert
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-
-
-def convert_orig_tf1_checkpoint_to_pytorch(tf_checkpoint_path, convbert_config_file, pytorch_dump_path):
- conf = ConvBertConfig.from_json_file(convbert_config_file)
- model = ConvBertModel(conf)
-
- model = load_tf_weights_in_convbert(model, conf, tf_checkpoint_path)
- model.save_pretrained(pytorch_dump_path)
-
- tf_model = TFConvBertModel.from_pretrained(pytorch_dump_path, from_pt=True)
- tf_model.save_pretrained(pytorch_dump_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
- )
- parser.add_argument(
- "--convbert_config_file",
- default=None,
- type=str,
- required=True,
- help=(
- "The config json file corresponding to the pre-trained ConvBERT model. \n"
- "This specifies the model architecture."
- ),
- )
- parser.add_argument(
- "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
- args = parser.parse_args()
- convert_orig_tf1_checkpoint_to_pytorch(args.tf_checkpoint_path, args.convbert_config_file, args.pytorch_dump_path)
diff --git a/transformers/models/convbert/modeling_convbert.py b/transformers/models/convbert/modeling_convbert.py
deleted file mode 100644
index d88add4e1390ef790c670c2407f280a8b4ab743a..0000000000000000000000000000000000000000
--- a/transformers/models/convbert/modeling_convbert.py
+++ /dev/null
@@ -1,1337 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch ConvBERT model."""
-
-
-import math
-import os
-from operator import attrgetter
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN, get_activation
-from ...modeling_outputs import (
- BaseModelOutputWithCrossAttentions,
- MaskedLMOutput,
- MultipleChoiceModelOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel, SequenceSummary
-from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_convbert import ConvBertConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "YituTech/conv-bert-base"
-_CONFIG_FOR_DOC = "ConvBertConfig"
-
-
-from ..deprecated._archive_maps import CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-def load_tf_weights_in_convbert(model, config, tf_checkpoint_path):
- """Load tf checkpoints in a pytorch model."""
- try:
- import tensorflow as tf
- except ImportError:
- logger.error(
- "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
- "https://www.tensorflow.org/install/ for installation instructions."
- )
- raise
- tf_path = os.path.abspath(tf_checkpoint_path)
- logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
- # Load weights from TF model
- init_vars = tf.train.list_variables(tf_path)
- tf_data = {}
- for name, shape in init_vars:
- logger.info(f"Loading TF weight {name} with shape {shape}")
- array = tf.train.load_variable(tf_path, name)
- tf_data[name] = array
-
- param_mapping = {
- "embeddings.word_embeddings.weight": "electra/embeddings/word_embeddings",
- "embeddings.position_embeddings.weight": "electra/embeddings/position_embeddings",
- "embeddings.token_type_embeddings.weight": "electra/embeddings/token_type_embeddings",
- "embeddings.LayerNorm.weight": "electra/embeddings/LayerNorm/gamma",
- "embeddings.LayerNorm.bias": "electra/embeddings/LayerNorm/beta",
- "embeddings_project.weight": "electra/embeddings_project/kernel",
- "embeddings_project.bias": "electra/embeddings_project/bias",
- }
- if config.num_groups > 1:
- group_dense_name = "g_dense"
- else:
- group_dense_name = "dense"
-
- for j in range(config.num_hidden_layers):
- param_mapping[
- f"encoder.layer.{j}.attention.self.query.weight"
- ] = f"electra/encoder/layer_{j}/attention/self/query/kernel"
- param_mapping[
- f"encoder.layer.{j}.attention.self.query.bias"
- ] = f"electra/encoder/layer_{j}/attention/self/query/bias"
- param_mapping[
- f"encoder.layer.{j}.attention.self.key.weight"
- ] = f"electra/encoder/layer_{j}/attention/self/key/kernel"
- param_mapping[
- f"encoder.layer.{j}.attention.self.key.bias"
- ] = f"electra/encoder/layer_{j}/attention/self/key/bias"
- param_mapping[
- f"encoder.layer.{j}.attention.self.value.weight"
- ] = f"electra/encoder/layer_{j}/attention/self/value/kernel"
- param_mapping[
- f"encoder.layer.{j}.attention.self.value.bias"
- ] = f"electra/encoder/layer_{j}/attention/self/value/bias"
- param_mapping[
- f"encoder.layer.{j}.attention.self.key_conv_attn_layer.depthwise.weight"
- ] = f"electra/encoder/layer_{j}/attention/self/conv_attn_key/depthwise_kernel"
- param_mapping[
- f"encoder.layer.{j}.attention.self.key_conv_attn_layer.pointwise.weight"
- ] = f"electra/encoder/layer_{j}/attention/self/conv_attn_key/pointwise_kernel"
- param_mapping[
- f"encoder.layer.{j}.attention.self.key_conv_attn_layer.bias"
- ] = f"electra/encoder/layer_{j}/attention/self/conv_attn_key/bias"
- param_mapping[
- f"encoder.layer.{j}.attention.self.conv_kernel_layer.weight"
- ] = f"electra/encoder/layer_{j}/attention/self/conv_attn_kernel/kernel"
- param_mapping[
- f"encoder.layer.{j}.attention.self.conv_kernel_layer.bias"
- ] = f"electra/encoder/layer_{j}/attention/self/conv_attn_kernel/bias"
- param_mapping[
- f"encoder.layer.{j}.attention.self.conv_out_layer.weight"
- ] = f"electra/encoder/layer_{j}/attention/self/conv_attn_point/kernel"
- param_mapping[
- f"encoder.layer.{j}.attention.self.conv_out_layer.bias"
- ] = f"electra/encoder/layer_{j}/attention/self/conv_attn_point/bias"
- param_mapping[
- f"encoder.layer.{j}.attention.output.dense.weight"
- ] = f"electra/encoder/layer_{j}/attention/output/dense/kernel"
- param_mapping[
- f"encoder.layer.{j}.attention.output.LayerNorm.weight"
- ] = f"electra/encoder/layer_{j}/attention/output/LayerNorm/gamma"
- param_mapping[
- f"encoder.layer.{j}.attention.output.dense.bias"
- ] = f"electra/encoder/layer_{j}/attention/output/dense/bias"
- param_mapping[
- f"encoder.layer.{j}.attention.output.LayerNorm.bias"
- ] = f"electra/encoder/layer_{j}/attention/output/LayerNorm/beta"
- param_mapping[
- f"encoder.layer.{j}.intermediate.dense.weight"
- ] = f"electra/encoder/layer_{j}/intermediate/{group_dense_name}/kernel"
- param_mapping[
- f"encoder.layer.{j}.intermediate.dense.bias"
- ] = f"electra/encoder/layer_{j}/intermediate/{group_dense_name}/bias"
- param_mapping[
- f"encoder.layer.{j}.output.dense.weight"
- ] = f"electra/encoder/layer_{j}/output/{group_dense_name}/kernel"
- param_mapping[
- f"encoder.layer.{j}.output.dense.bias"
- ] = f"electra/encoder/layer_{j}/output/{group_dense_name}/bias"
- param_mapping[
- f"encoder.layer.{j}.output.LayerNorm.weight"
- ] = f"electra/encoder/layer_{j}/output/LayerNorm/gamma"
- param_mapping[f"encoder.layer.{j}.output.LayerNorm.bias"] = f"electra/encoder/layer_{j}/output/LayerNorm/beta"
-
- for param in model.named_parameters():
- param_name = param[0]
- retriever = attrgetter(param_name)
- result = retriever(model)
- tf_name = param_mapping[param_name]
- value = torch.from_numpy(tf_data[tf_name])
- logger.info(f"TF: {tf_name}, PT: {param_name} ")
- if tf_name.endswith("/kernel"):
- if not tf_name.endswith("/intermediate/g_dense/kernel"):
- if not tf_name.endswith("/output/g_dense/kernel"):
- value = value.T
- if tf_name.endswith("/depthwise_kernel"):
- value = value.permute(1, 2, 0) # 2, 0, 1
- if tf_name.endswith("/pointwise_kernel"):
- value = value.permute(2, 1, 0) # 2, 1, 0
- if tf_name.endswith("/conv_attn_key/bias"):
- value = value.unsqueeze(-1)
- result.data = value
- return model
-
-
-class ConvBertEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.embedding_size)
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.embedding_size)
-
- # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
- # any TensorFlow checkpoint file
- self.LayerNorm = nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
- self.register_buffer(
- "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
- )
-
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- ) -> torch.LongTensor:
- if input_ids is not None:
- input_shape = input_ids.size()
- else:
- input_shape = inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, :seq_length]
-
- # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
- # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
- # issue #5664
- if token_type_ids is None:
- if hasattr(self, "token_type_ids"):
- buffered_token_type_ids = self.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
- position_embeddings = self.position_embeddings(position_ids)
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
- embeddings = inputs_embeds + position_embeddings + token_type_embeddings
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-class ConvBertPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ConvBertConfig
- load_tf_weights = load_tf_weights_in_convbert
- base_model_prefix = "convbert"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-class SeparableConv1D(nn.Module):
- """This class implements separable convolution, i.e. a depthwise and a pointwise layer"""
-
- def __init__(self, config, input_filters, output_filters, kernel_size, **kwargs):
- super().__init__()
- self.depthwise = nn.Conv1d(
- input_filters,
- input_filters,
- kernel_size=kernel_size,
- groups=input_filters,
- padding=kernel_size // 2,
- bias=False,
- )
- self.pointwise = nn.Conv1d(input_filters, output_filters, kernel_size=1, bias=False)
- self.bias = nn.Parameter(torch.zeros(output_filters, 1))
-
- self.depthwise.weight.data.normal_(mean=0.0, std=config.initializer_range)
- self.pointwise.weight.data.normal_(mean=0.0, std=config.initializer_range)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- x = self.depthwise(hidden_states)
- x = self.pointwise(x)
- x += self.bias
- return x
-
-
-class ConvBertSelfAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- new_num_attention_heads = config.num_attention_heads // config.head_ratio
- if new_num_attention_heads < 1:
- self.head_ratio = config.num_attention_heads
- self.num_attention_heads = 1
- else:
- self.num_attention_heads = new_num_attention_heads
- self.head_ratio = config.head_ratio
-
- self.conv_kernel_size = config.conv_kernel_size
- if config.hidden_size % self.num_attention_heads != 0:
- raise ValueError("hidden_size should be divisible by num_attention_heads")
-
- self.attention_head_size = (config.hidden_size // self.num_attention_heads) // 2
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size)
- self.key = nn.Linear(config.hidden_size, self.all_head_size)
- self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.key_conv_attn_layer = SeparableConv1D(
- config, config.hidden_size, self.all_head_size, self.conv_kernel_size
- )
- self.conv_kernel_layer = nn.Linear(self.all_head_size, self.num_attention_heads * self.conv_kernel_size)
- self.conv_out_layer = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.unfold = nn.Unfold(
- kernel_size=[self.conv_kernel_size, 1], padding=[int((self.conv_kernel_size - 1) / 2), 0]
- )
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- def transpose_for_scores(self, x):
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(*new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- mixed_query_layer = self.query(hidden_states)
- batch_size = hidden_states.size(0)
- # If this is instantiated as a cross-attention module, the keys
- # and values come from an encoder; the attention mask needs to be
- # such that the encoder's padding tokens are not attended to.
- if encoder_hidden_states is not None:
- mixed_key_layer = self.key(encoder_hidden_states)
- mixed_value_layer = self.value(encoder_hidden_states)
- else:
- mixed_key_layer = self.key(hidden_states)
- mixed_value_layer = self.value(hidden_states)
-
- mixed_key_conv_attn_layer = self.key_conv_attn_layer(hidden_states.transpose(1, 2))
- mixed_key_conv_attn_layer = mixed_key_conv_attn_layer.transpose(1, 2)
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
- key_layer = self.transpose_for_scores(mixed_key_layer)
- value_layer = self.transpose_for_scores(mixed_value_layer)
- conv_attn_layer = torch.multiply(mixed_key_conv_attn_layer, mixed_query_layer)
-
- conv_kernel_layer = self.conv_kernel_layer(conv_attn_layer)
- conv_kernel_layer = torch.reshape(conv_kernel_layer, [-1, self.conv_kernel_size, 1])
- conv_kernel_layer = torch.softmax(conv_kernel_layer, dim=1)
-
- conv_out_layer = self.conv_out_layer(hidden_states)
- conv_out_layer = torch.reshape(conv_out_layer, [batch_size, -1, self.all_head_size])
- conv_out_layer = conv_out_layer.transpose(1, 2).contiguous().unsqueeze(-1)
- conv_out_layer = nn.functional.unfold(
- conv_out_layer,
- kernel_size=[self.conv_kernel_size, 1],
- dilation=1,
- padding=[(self.conv_kernel_size - 1) // 2, 0],
- stride=1,
- )
- conv_out_layer = conv_out_layer.transpose(1, 2).reshape(
- batch_size, -1, self.all_head_size, self.conv_kernel_size
- )
- conv_out_layer = torch.reshape(conv_out_layer, [-1, self.attention_head_size, self.conv_kernel_size])
- conv_out_layer = torch.matmul(conv_out_layer, conv_kernel_layer)
- conv_out_layer = torch.reshape(conv_out_layer, [-1, self.all_head_size])
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in ConvBertModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
-
- conv_out = torch.reshape(conv_out_layer, [batch_size, -1, self.num_attention_heads, self.attention_head_size])
- context_layer = torch.cat([context_layer, conv_out], 2)
-
- # conv and context
- new_context_layer_shape = context_layer.size()[:-2] + (
- self.num_attention_heads * self.attention_head_size * 2,
- )
- context_layer = context_layer.view(*new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
- return outputs
-
-
-class ConvBertSelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class ConvBertAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.self = ConvBertSelfAttention(config)
- self.output = ConvBertSelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor, Optional[torch.FloatTensor]]:
- self_outputs = self.self(
- hidden_states,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- output_attentions,
- )
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-class GroupedLinearLayer(nn.Module):
- def __init__(self, input_size, output_size, num_groups):
- super().__init__()
- self.input_size = input_size
- self.output_size = output_size
- self.num_groups = num_groups
- self.group_in_dim = self.input_size // self.num_groups
- self.group_out_dim = self.output_size // self.num_groups
- self.weight = nn.Parameter(torch.empty(self.num_groups, self.group_in_dim, self.group_out_dim))
- self.bias = nn.Parameter(torch.empty(output_size))
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- batch_size = list(hidden_states.size())[0]
- x = torch.reshape(hidden_states, [-1, self.num_groups, self.group_in_dim])
- x = x.permute(1, 0, 2)
- x = torch.matmul(x, self.weight)
- x = x.permute(1, 0, 2)
- x = torch.reshape(x, [batch_size, -1, self.output_size])
- x = x + self.bias
- return x
-
-
-class ConvBertIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- if config.num_groups == 1:
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- else:
- self.dense = GroupedLinearLayer(
- input_size=config.hidden_size, output_size=config.intermediate_size, num_groups=config.num_groups
- )
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-class ConvBertOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- if config.num_groups == 1:
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- else:
- self.dense = GroupedLinearLayer(
- input_size=config.intermediate_size, output_size=config.hidden_size, num_groups=config.num_groups
- )
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class ConvBertLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = ConvBertAttention(config)
- self.is_decoder = config.is_decoder
- self.add_cross_attention = config.add_cross_attention
- if self.add_cross_attention:
- if not self.is_decoder:
- raise TypeError(f"{self} should be used as a decoder model if cross attention is added")
- self.crossattention = ConvBertAttention(config)
- self.intermediate = ConvBertIntermediate(config)
- self.output = ConvBertOutput(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor, Optional[torch.FloatTensor]]:
- self_attention_outputs = self.attention(
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions=output_attentions,
- )
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- if self.is_decoder and encoder_hidden_states is not None:
- if not hasattr(self, "crossattention"):
- raise AttributeError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
- " by setting `config.add_cross_attention=True`"
- )
- cross_attention_outputs = self.crossattention(
- attention_output,
- encoder_attention_mask,
- head_mask,
- encoder_hidden_states,
- output_attentions,
- )
- attention_output = cross_attention_outputs[0]
- outputs = outputs + cross_attention_outputs[1:] # add cross attentions if we output attention weights
-
- layer_output = apply_chunking_to_forward(
- self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
- )
- outputs = (layer_output,) + outputs
- return outputs
-
- def feed_forward_chunk(self, attention_output):
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.output(intermediate_output, attention_output)
- return layer_output
-
-
-class ConvBertEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([ConvBertLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple, BaseModelOutputWithCrossAttentions]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- output_attentions,
- )
- hidden_states = layer_outputs[0]
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, all_hidden_states, all_self_attentions, all_cross_attentions]
- if v is not None
- )
- return BaseModelOutputWithCrossAttentions(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-class ConvBertPredictionHeadTransform(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = ACT2FN[config.hidden_act]
- else:
- self.transform_act_fn = config.hidden_act
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
- return hidden_states
-
-
-CONVBERT_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`ConvBertConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CONVBERT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare ConvBERT Model transformer outputting raw hidden-states without any specific head on top.",
- CONVBERT_START_DOCSTRING,
-)
-class ConvBertModel(ConvBertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.embeddings = ConvBertEmbeddings(config)
-
- if config.embedding_size != config.hidden_size:
- self.embeddings_project = nn.Linear(config.embedding_size, config.hidden_size)
-
- self.encoder = ConvBertEncoder(config)
- self.config = config
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithCrossAttentions]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device)
- if token_type_ids is None:
- if hasattr(self.embeddings, "token_type_ids"):
- buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- hidden_states = self.embeddings(
- input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
- )
-
- if hasattr(self, "embeddings_project"):
- hidden_states = self.embeddings_project(hidden_states)
-
- hidden_states = self.encoder(
- hidden_states,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- return hidden_states
-
-
-class ConvBertGeneratorPredictions(nn.Module):
- """Prediction module for the generator, made up of two dense layers."""
-
- def __init__(self, config):
- super().__init__()
-
- self.activation = get_activation("gelu")
- self.LayerNorm = nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
- self.dense = nn.Linear(config.hidden_size, config.embedding_size)
-
- def forward(self, generator_hidden_states: torch.FloatTensor) -> torch.FloatTensor:
- hidden_states = self.dense(generator_hidden_states)
- hidden_states = self.activation(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
-
- return hidden_states
-
-
-@add_start_docstrings("""ConvBERT Model with a `language modeling` head on top.""", CONVBERT_START_DOCSTRING)
-class ConvBertForMaskedLM(ConvBertPreTrainedModel):
- _tied_weights_keys = ["generator.lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.convbert = ConvBertModel(config)
- self.generator_predictions = ConvBertGeneratorPredictions(config)
-
- self.generator_lm_head = nn.Linear(config.embedding_size, config.vocab_size)
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.generator_lm_head
-
- def set_output_embeddings(self, word_embeddings):
- self.generator_lm_head = word_embeddings
-
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- generator_hidden_states = self.convbert(
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict,
- )
- generator_sequence_output = generator_hidden_states[0]
-
- prediction_scores = self.generator_predictions(generator_sequence_output)
- prediction_scores = self.generator_lm_head(prediction_scores)
-
- loss = None
- # Masked language modeling softmax layer
- if labels is not None:
- loss_fct = nn.CrossEntropyLoss() # -100 index = padding token
- loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + generator_hidden_states[1:]
- return ((loss,) + output) if loss is not None else output
-
- return MaskedLMOutput(
- loss=loss,
- logits=prediction_scores,
- hidden_states=generator_hidden_states.hidden_states,
- attentions=generator_hidden_states.attentions,
- )
-
-
-class ConvBertClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
-
- self.config = config
-
- def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
- x = hidden_states[:, 0, :] # take token (equiv. to [CLS])
- x = self.dropout(x)
- x = self.dense(x)
- x = ACT2FN[self.config.hidden_act](x)
- x = self.dropout(x)
- x = self.out_proj(x)
- return x
-
-
-@add_start_docstrings(
- """
- ConvBERT Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- CONVBERT_START_DOCSTRING,
-)
-class ConvBertForSequenceClassification(ConvBertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
- self.convbert = ConvBertModel(config)
- self.classifier = ConvBertClassificationHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.convbert(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- ConvBERT Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- CONVBERT_START_DOCSTRING,
-)
-class ConvBertForMultipleChoice(ConvBertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.convbert = ConvBertModel(config)
- self.sequence_summary = SequenceSummary(config)
- self.classifier = nn.Linear(config.hidden_size, 1)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(
- CONVBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
- )
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
- position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
- inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.convbert(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- pooled_output = self.sequence_summary(sequence_output)
- logits = self.classifier(pooled_output)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- ConvBERT Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- CONVBERT_START_DOCSTRING,
-)
-class ConvBertForTokenClassification(ConvBertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.convbert = ConvBertModel(config)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.convbert(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- ConvBERT Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- CONVBERT_START_DOCSTRING,
-)
-class ConvBertForQuestionAnswering(ConvBertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.convbert = ConvBertModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- start_positions: Optional[torch.LongTensor] = None,
- end_positions: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.convbert(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[1:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/convbert/modeling_tf_convbert.py b/transformers/models/convbert/modeling_tf_convbert.py
deleted file mode 100644
index 7206b3558ace8a26994d9608d3963ee6f34f1e91..0000000000000000000000000000000000000000
--- a/transformers/models/convbert/modeling_tf_convbert.py
+++ /dev/null
@@ -1,1468 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 ConvBERT model."""
-
-
-from __future__ import annotations
-
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFMaskedLMOutput,
- TFMultipleChoiceModelOutput,
- TFQuestionAnsweringModelOutput,
- TFSequenceClassifierOutput,
- TFTokenClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFMaskedLanguageModelingLoss,
- TFModelInputType,
- TFMultipleChoiceLoss,
- TFPreTrainedModel,
- TFQuestionAnsweringLoss,
- TFSequenceClassificationLoss,
- TFSequenceSummary,
- TFTokenClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_convbert import ConvBertConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "YituTech/conv-bert-base"
-_CONFIG_FOR_DOC = "ConvBertConfig"
-
-
-from ..deprecated._archive_maps import TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.albert.modeling_tf_albert.TFAlbertEmbeddings with Albert->ConvBert
-class TFConvBertEmbeddings(keras.layers.Layer):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config: ConvBertConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embedding_size = config.embedding_size
- self.max_position_embeddings = config.max_position_embeddings
- self.initializer_range = config.initializer_range
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
-
- def build(self, input_shape=None):
- with tf.name_scope("word_embeddings"):
- self.weight = self.add_weight(
- name="weight",
- shape=[self.config.vocab_size, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
-
- with tf.name_scope("token_type_embeddings"):
- self.token_type_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.config.type_vocab_size, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
-
- with tf.name_scope("position_embeddings"):
- self.position_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.max_position_embeddings, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.embedding_size])
-
- # Copied from transformers.models.bert.modeling_tf_bert.TFBertEmbeddings.call
- def call(
- self,
- input_ids: tf.Tensor = None,
- position_ids: tf.Tensor = None,
- token_type_ids: tf.Tensor = None,
- inputs_embeds: tf.Tensor = None,
- past_key_values_length=0,
- training: bool = False,
- ) -> tf.Tensor:
- """
- Applies embedding based on inputs tensor.
-
- Returns:
- final_embeddings (`tf.Tensor`): output embedding tensor.
- """
- if input_ids is None and inputs_embeds is None:
- raise ValueError("Need to provide either `input_ids` or `input_embeds`.")
-
- if input_ids is not None:
- check_embeddings_within_bounds(input_ids, self.config.vocab_size)
- inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
-
- input_shape = shape_list(inputs_embeds)[:-1]
-
- if token_type_ids is None:
- token_type_ids = tf.fill(dims=input_shape, value=0)
-
- if position_ids is None:
- position_ids = tf.expand_dims(
- tf.range(start=past_key_values_length, limit=input_shape[1] + past_key_values_length), axis=0
- )
-
- position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
- token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)
- final_embeddings = inputs_embeds + position_embeds + token_type_embeds
- final_embeddings = self.LayerNorm(inputs=final_embeddings)
- final_embeddings = self.dropout(inputs=final_embeddings, training=training)
-
- return final_embeddings
-
-
-class TFConvBertSelfAttention(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- if config.hidden_size % config.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- new_num_attention_heads = int(config.num_attention_heads / config.head_ratio)
- if new_num_attention_heads < 1:
- self.head_ratio = config.num_attention_heads
- num_attention_heads = 1
- else:
- num_attention_heads = new_num_attention_heads
- self.head_ratio = config.head_ratio
-
- self.num_attention_heads = num_attention_heads
- self.conv_kernel_size = config.conv_kernel_size
-
- if config.hidden_size % self.num_attention_heads != 0:
- raise ValueError("hidden_size should be divisible by num_attention_heads")
-
- self.attention_head_size = config.hidden_size // config.num_attention_heads
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.query = keras.layers.Dense(
- self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"
- )
- self.key = keras.layers.Dense(
- self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key"
- )
- self.value = keras.layers.Dense(
- self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"
- )
-
- self.key_conv_attn_layer = keras.layers.SeparableConv1D(
- self.all_head_size,
- self.conv_kernel_size,
- padding="same",
- activation=None,
- depthwise_initializer=get_initializer(1 / self.conv_kernel_size),
- pointwise_initializer=get_initializer(config.initializer_range),
- name="key_conv_attn_layer",
- )
-
- self.conv_kernel_layer = keras.layers.Dense(
- self.num_attention_heads * self.conv_kernel_size,
- activation=None,
- name="conv_kernel_layer",
- kernel_initializer=get_initializer(config.initializer_range),
- )
-
- self.conv_out_layer = keras.layers.Dense(
- self.all_head_size,
- activation=None,
- name="conv_out_layer",
- kernel_initializer=get_initializer(config.initializer_range),
- )
-
- self.dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)
- self.config = config
-
- def transpose_for_scores(self, x, batch_size):
- # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]
- x = tf.reshape(x, (batch_size, -1, self.num_attention_heads, self.attention_head_size))
- return tf.transpose(x, perm=[0, 2, 1, 3])
-
- def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False):
- batch_size = shape_list(hidden_states)[0]
- mixed_query_layer = self.query(hidden_states)
- mixed_key_layer = self.key(hidden_states)
- mixed_value_layer = self.value(hidden_states)
-
- mixed_key_conv_attn_layer = self.key_conv_attn_layer(hidden_states)
-
- query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
- key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
- conv_attn_layer = tf.multiply(mixed_key_conv_attn_layer, mixed_query_layer)
-
- conv_kernel_layer = self.conv_kernel_layer(conv_attn_layer)
- conv_kernel_layer = tf.reshape(conv_kernel_layer, [-1, self.conv_kernel_size, 1])
- conv_kernel_layer = stable_softmax(conv_kernel_layer, axis=1)
-
- paddings = tf.constant(
- [
- [
- 0,
- 0,
- ],
- [int((self.conv_kernel_size - 1) / 2), int((self.conv_kernel_size - 1) / 2)],
- [0, 0],
- ]
- )
-
- conv_out_layer = self.conv_out_layer(hidden_states)
- conv_out_layer = tf.reshape(conv_out_layer, [batch_size, -1, self.all_head_size])
- conv_out_layer = tf.pad(conv_out_layer, paddings, "CONSTANT")
-
- unfold_conv_out_layer = tf.stack(
- [
- tf.slice(conv_out_layer, [0, i, 0], [batch_size, shape_list(mixed_query_layer)[1], self.all_head_size])
- for i in range(self.conv_kernel_size)
- ],
- axis=-1,
- )
-
- conv_out_layer = tf.reshape(unfold_conv_out_layer, [-1, self.attention_head_size, self.conv_kernel_size])
-
- conv_out_layer = tf.matmul(conv_out_layer, conv_kernel_layer)
- conv_out_layer = tf.reshape(conv_out_layer, [-1, self.all_head_size])
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = tf.matmul(
- query_layer, key_layer, transpose_b=True
- ) # (batch size, num_heads, seq_len_q, seq_len_k)
- dk = tf.cast(shape_list(key_layer)[-1], attention_scores.dtype) # scale attention_scores
- attention_scores = attention_scores / tf.math.sqrt(dk)
-
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = stable_softmax(attention_scores, axis=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs, training=training)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- value_layer = tf.reshape(
- mixed_value_layer, [batch_size, -1, self.num_attention_heads, self.attention_head_size]
- )
- value_layer = tf.transpose(value_layer, [0, 2, 1, 3])
-
- context_layer = tf.matmul(attention_probs, value_layer)
- context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
-
- conv_out = tf.reshape(conv_out_layer, [batch_size, -1, self.num_attention_heads, self.attention_head_size])
- context_layer = tf.concat([context_layer, conv_out], 2)
- context_layer = tf.reshape(
- context_layer, (batch_size, -1, self.head_ratio * self.all_head_size)
- ) # (batch_size, seq_len_q, all_head_size)
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "query", None) is not None:
- with tf.name_scope(self.query.name):
- self.query.build([None, None, self.config.hidden_size])
- if getattr(self, "key", None) is not None:
- with tf.name_scope(self.key.name):
- self.key.build([None, None, self.config.hidden_size])
- if getattr(self, "value", None) is not None:
- with tf.name_scope(self.value.name):
- self.value.build([None, None, self.config.hidden_size])
- if getattr(self, "key_conv_attn_layer", None) is not None:
- with tf.name_scope(self.key_conv_attn_layer.name):
- self.key_conv_attn_layer.build([None, None, self.config.hidden_size])
- if getattr(self, "conv_kernel_layer", None) is not None:
- with tf.name_scope(self.conv_kernel_layer.name):
- self.conv_kernel_layer.build([None, None, self.all_head_size])
- if getattr(self, "conv_out_layer", None) is not None:
- with tf.name_scope(self.conv_out_layer.name):
- self.conv_out_layer.build([None, None, self.config.hidden_size])
-
-
-class TFConvBertSelfOutput(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states, input_tensor, training=False):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states, training=training)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
-
-
-class TFConvBertAttention(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.self_attention = TFConvBertSelfAttention(config, name="self")
- self.dense_output = TFConvBertSelfOutput(config, name="output")
-
- def prune_heads(self, heads):
- raise NotImplementedError
-
- def call(self, input_tensor, attention_mask, head_mask, output_attentions, training=False):
- self_outputs = self.self_attention(
- input_tensor, attention_mask, head_mask, output_attentions, training=training
- )
- attention_output = self.dense_output(self_outputs[0], input_tensor, training=training)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "self_attention", None) is not None:
- with tf.name_scope(self.self_attention.name):
- self.self_attention.build(None)
- if getattr(self, "dense_output", None) is not None:
- with tf.name_scope(self.dense_output.name):
- self.dense_output.build(None)
-
-
-class GroupedLinearLayer(keras.layers.Layer):
- def __init__(self, input_size, output_size, num_groups, kernel_initializer, **kwargs):
- super().__init__(**kwargs)
- self.input_size = input_size
- self.output_size = output_size
- self.num_groups = num_groups
- self.kernel_initializer = kernel_initializer
- self.group_in_dim = self.input_size // self.num_groups
- self.group_out_dim = self.output_size // self.num_groups
-
- def build(self, input_shape=None):
- self.kernel = self.add_weight(
- "kernel",
- shape=[self.group_out_dim, self.group_in_dim, self.num_groups],
- initializer=self.kernel_initializer,
- trainable=True,
- )
-
- self.bias = self.add_weight(
- "bias", shape=[self.output_size], initializer=self.kernel_initializer, dtype=self.dtype, trainable=True
- )
- super().build(input_shape)
-
- def call(self, hidden_states):
- batch_size = shape_list(hidden_states)[0]
- x = tf.transpose(tf.reshape(hidden_states, [-1, self.num_groups, self.group_in_dim]), [1, 0, 2])
- x = tf.matmul(x, tf.transpose(self.kernel, [2, 1, 0]))
- x = tf.transpose(x, [1, 0, 2])
- x = tf.reshape(x, [batch_size, -1, self.output_size])
- x = tf.nn.bias_add(value=x, bias=self.bias)
- return x
-
-
-class TFConvBertIntermediate(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- if config.num_groups == 1:
- self.dense = keras.layers.Dense(
- config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- else:
- self.dense = GroupedLinearLayer(
- config.hidden_size,
- config.intermediate_size,
- num_groups=config.num_groups,
- kernel_initializer=get_initializer(config.initializer_range),
- name="dense",
- )
-
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.intermediate_act_fn = config.hidden_act
- self.config = config
-
- def call(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFConvBertOutput(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- if config.num_groups == 1:
- self.dense = keras.layers.Dense(
- config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- else:
- self.dense = GroupedLinearLayer(
- config.intermediate_size,
- config.hidden_size,
- num_groups=config.num_groups,
- kernel_initializer=get_initializer(config.initializer_range),
- name="dense",
- )
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states, input_tensor, training=False):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states, training=training)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.intermediate_size])
-
-
-class TFConvBertLayer(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.attention = TFConvBertAttention(config, name="attention")
- self.intermediate = TFConvBertIntermediate(config, name="intermediate")
- self.bert_output = TFConvBertOutput(config, name="output")
-
- def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False):
- attention_outputs = self.attention(
- hidden_states, attention_mask, head_mask, output_attentions, training=training
- )
- attention_output = attention_outputs[0]
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.bert_output(intermediate_output, attention_output, training=training)
- outputs = (layer_output,) + attention_outputs[1:] # add attentions if we output them
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "intermediate", None) is not None:
- with tf.name_scope(self.intermediate.name):
- self.intermediate.build(None)
- if getattr(self, "bert_output", None) is not None:
- with tf.name_scope(self.bert_output.name):
- self.bert_output.build(None)
-
-
-class TFConvBertEncoder(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.layer = [TFConvBertLayer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]
-
- def call(
- self,
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions,
- output_hidden_states,
- return_dict,
- training=False,
- ):
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_outputs = layer_module(
- hidden_states, attention_mask, head_mask[i], output_attentions, training=training
- )
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- # Add last layer
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)
-
- return TFBaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "layer", None) is not None:
- for layer in self.layer:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-class TFConvBertPredictionHeadTransform(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- config.embedding_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
-
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.transform_act_fn = config.hidden_act
-
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.config = config
-
- def call(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
-
-
-@keras_serializable
-class TFConvBertMainLayer(keras.layers.Layer):
- config_class = ConvBertConfig
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.embeddings = TFConvBertEmbeddings(config, name="embeddings")
-
- if config.embedding_size != config.hidden_size:
- self.embeddings_project = keras.layers.Dense(config.hidden_size, name="embeddings_project")
-
- self.encoder = TFConvBertEncoder(config, name="encoder")
- self.config = config
-
- def get_input_embeddings(self):
- return self.embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.weight = value
- self.embeddings.vocab_size = value.shape[0]
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- raise NotImplementedError
-
- def get_extended_attention_mask(self, attention_mask, input_shape, dtype):
- if attention_mask is None:
- attention_mask = tf.fill(input_shape, 1)
-
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- extended_attention_mask = tf.reshape(attention_mask, (input_shape[0], 1, 1, input_shape[1]))
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and -10000.0 for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- extended_attention_mask = tf.cast(extended_attention_mask, dtype)
- extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
-
- return extended_attention_mask
-
- def get_head_mask(self, head_mask):
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.config.num_hidden_layers
-
- return head_mask
-
- @unpack_inputs
- def call(
- self,
- input_ids=None,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- inputs_embeds=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- training=False,
- ):
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = tf.fill(input_shape, 1)
-
- if token_type_ids is None:
- token_type_ids = tf.fill(input_shape, 0)
-
- hidden_states = self.embeddings(input_ids, position_ids, token_type_ids, inputs_embeds, training=training)
- extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, hidden_states.dtype)
- head_mask = self.get_head_mask(head_mask)
-
- if hasattr(self, "embeddings_project"):
- hidden_states = self.embeddings_project(hidden_states, training=training)
-
- hidden_states = self.encoder(
- hidden_states,
- extended_attention_mask,
- head_mask,
- output_attentions,
- output_hidden_states,
- return_dict,
- training=training,
- )
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "embeddings_project", None) is not None:
- with tf.name_scope(self.embeddings_project.name):
- self.embeddings_project.build([None, None, self.config.embedding_size])
-
-
-class TFConvBertPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ConvBertConfig
- base_model_prefix = "convbert"
-
-
-CONVBERT_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Args:
- config ([`ConvBertConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CONVBERT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`Numpy array` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare ConvBERT Model transformer outputting raw hidden-states without any specific head on top.",
- CONVBERT_START_DOCSTRING,
-)
-class TFConvBertModel(TFConvBertPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.convbert = TFConvBertMainLayer(config, name="convbert")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: Optional[Union[np.array, tf.Tensor]] = None,
- token_type_ids: Optional[Union[np.array, tf.Tensor]] = None,
- position_ids: Optional[Union[np.array, tf.Tensor]] = None,
- head_mask: Optional[Union[np.array, tf.Tensor]] = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- outputs = self.convbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convbert", None) is not None:
- with tf.name_scope(self.convbert.name):
- self.convbert.build(None)
-
-
-class TFConvBertMaskedLMHead(keras.layers.Layer):
- def __init__(self, config, input_embeddings, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embedding_size = config.embedding_size
- self.input_embeddings = input_embeddings
-
- def build(self, input_shape):
- self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")
-
- super().build(input_shape)
-
- def get_output_embeddings(self):
- return self.input_embeddings
-
- def set_output_embeddings(self, value):
- self.input_embeddings.weight = value
- self.input_embeddings.vocab_size = shape_list(value)[0]
-
- def get_bias(self):
- return {"bias": self.bias}
-
- def set_bias(self, value):
- self.bias = value["bias"]
- self.config.vocab_size = shape_list(value["bias"])[0]
-
- def call(self, hidden_states):
- seq_length = shape_list(tensor=hidden_states)[1]
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.embedding_size])
- hidden_states = tf.matmul(a=hidden_states, b=self.input_embeddings.weight, transpose_b=True)
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])
- hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)
-
- return hidden_states
-
-
-class TFConvBertGeneratorPredictions(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dense = keras.layers.Dense(config.embedding_size, name="dense")
- self.config = config
-
- def call(self, generator_hidden_states, training=False):
- hidden_states = self.dense(generator_hidden_states)
- hidden_states = get_tf_activation("gelu")(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.embedding_size])
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings("""ConvBERT Model with a `language modeling` head on top.""", CONVBERT_START_DOCSTRING)
-class TFConvBertForMaskedLM(TFConvBertPreTrainedModel, TFMaskedLanguageModelingLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, **kwargs)
-
- self.config = config
- self.convbert = TFConvBertMainLayer(config, name="convbert")
- self.generator_predictions = TFConvBertGeneratorPredictions(config, name="generator_predictions")
-
- if isinstance(config.hidden_act, str):
- self.activation = get_tf_activation(config.hidden_act)
- else:
- self.activation = config.hidden_act
-
- self.generator_lm_head = TFConvBertMaskedLMHead(config, self.convbert.embeddings, name="generator_lm_head")
-
- def get_lm_head(self):
- return self.generator_lm_head
-
- def get_prefix_bias_name(self):
- return self.name + "/" + self.generator_lm_head.name
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFMaskedLMOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- generator_hidden_states = self.convbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- generator_sequence_output = generator_hidden_states[0]
- prediction_scores = self.generator_predictions(generator_sequence_output, training=training)
- prediction_scores = self.generator_lm_head(prediction_scores, training=training)
- loss = None if labels is None else self.hf_compute_loss(labels, prediction_scores)
-
- if not return_dict:
- output = (prediction_scores,) + generator_hidden_states[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFMaskedLMOutput(
- loss=loss,
- logits=prediction_scores,
- hidden_states=generator_hidden_states.hidden_states,
- attentions=generator_hidden_states.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convbert", None) is not None:
- with tf.name_scope(self.convbert.name):
- self.convbert.build(None)
- if getattr(self, "generator_predictions", None) is not None:
- with tf.name_scope(self.generator_predictions.name):
- self.generator_predictions.build(None)
- if getattr(self, "generator_lm_head", None) is not None:
- with tf.name_scope(self.generator_lm_head.name):
- self.generator_lm_head.build(None)
-
-
-class TFConvBertClassificationHead(keras.layers.Layer):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = keras.layers.Dropout(classifier_dropout)
- self.out_proj = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="out_proj"
- )
-
- self.config = config
-
- def call(self, hidden_states, **kwargs):
- x = hidden_states[:, 0, :] # take token (equiv. to [CLS])
- x = self.dropout(x)
- x = self.dense(x)
- x = get_tf_activation(self.config.hidden_act)(x)
- x = self.dropout(x)
- x = self.out_proj(x)
-
- return x
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "out_proj", None) is not None:
- with tf.name_scope(self.out_proj.name):
- self.out_proj.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- ConvBERT Model transformer with a sequence classification/regression head on top e.g., for GLUE tasks.
- """,
- CONVBERT_START_DOCSTRING,
-)
-class TFConvBertForSequenceClassification(TFConvBertPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
- self.convbert = TFConvBertMainLayer(config, name="convbert")
- self.classifier = TFConvBertClassificationHead(config, name="classifier")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFSequenceClassifierOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- outputs = self.convbert(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- logits = self.classifier(outputs[0], training=training)
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convbert", None) is not None:
- with tf.name_scope(self.convbert.name):
- self.convbert.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build(None)
-
-
-@add_start_docstrings(
- """
- ConvBERT Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- CONVBERT_START_DOCSTRING,
-)
-class TFConvBertForMultipleChoice(TFConvBertPreTrainedModel, TFMultipleChoiceLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.convbert = TFConvBertMainLayer(config, name="convbert")
- self.sequence_summary = TFSequenceSummary(
- config, initializer_range=config.initializer_range, name="sequence_summary"
- )
- self.classifier = keras.layers.Dense(
- 1, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(
- CONVBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
- )
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFMultipleChoiceModelOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]`
- where `num_choices` is the size of the second dimension of the input tensors. (See `input_ids` above)
- """
- if input_ids is not None:
- num_choices = shape_list(input_ids)[1]
- seq_length = shape_list(input_ids)[2]
- else:
- num_choices = shape_list(inputs_embeds)[1]
- seq_length = shape_list(inputs_embeds)[2]
-
- flat_input_ids = tf.reshape(input_ids, (-1, seq_length)) if input_ids is not None else None
- flat_attention_mask = tf.reshape(attention_mask, (-1, seq_length)) if attention_mask is not None else None
- flat_token_type_ids = tf.reshape(token_type_ids, (-1, seq_length)) if token_type_ids is not None else None
- flat_position_ids = tf.reshape(position_ids, (-1, seq_length)) if position_ids is not None else None
- flat_inputs_embeds = (
- tf.reshape(inputs_embeds, (-1, seq_length, shape_list(inputs_embeds)[3]))
- if inputs_embeds is not None
- else None
- )
- outputs = self.convbert(
- flat_input_ids,
- flat_attention_mask,
- flat_token_type_ids,
- flat_position_ids,
- head_mask,
- flat_inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- logits = self.sequence_summary(outputs[0], training=training)
- logits = self.classifier(logits)
- reshaped_logits = tf.reshape(logits, (-1, num_choices))
- loss = None if labels is None else self.hf_compute_loss(labels, reshaped_logits)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFMultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convbert", None) is not None:
- with tf.name_scope(self.convbert.name):
- self.convbert.build(None)
- if getattr(self, "sequence_summary", None) is not None:
- with tf.name_scope(self.sequence_summary.name):
- self.sequence_summary.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- ConvBERT Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- CONVBERT_START_DOCSTRING,
-)
-class TFConvBertForTokenClassification(TFConvBertPreTrainedModel, TFTokenClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
- self.convbert = TFConvBertMainLayer(config, name="convbert")
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = keras.layers.Dropout(classifier_dropout)
- self.classifier = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFTokenClassifierOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- outputs = self.convbert(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- sequence_output = self.dropout(sequence_output, training=training)
- logits = self.classifier(sequence_output)
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFTokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convbert", None) is not None:
- with tf.name_scope(self.convbert.name):
- self.convbert.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- ConvBERT Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- CONVBERT_START_DOCSTRING,
-)
-class TFConvBertForQuestionAnswering(TFConvBertPreTrainedModel, TFQuestionAnsweringLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
- self.convbert = TFConvBertMainLayer(config, name="convbert")
- self.qa_outputs = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFQuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- start_positions: tf.Tensor | None = None,
- end_positions: tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFQuestionAnsweringModelOutput]:
- r"""
- start_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- outputs = self.convbert(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = tf.split(logits, 2, axis=-1)
- start_logits = tf.squeeze(start_logits, axis=-1)
- end_logits = tf.squeeze(end_logits, axis=-1)
- loss = None
-
- if start_positions is not None and end_positions is not None:
- labels = {"start_position": start_positions}
- labels["end_position"] = end_positions
- loss = self.hf_compute_loss(labels, (start_logits, end_logits))
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFQuestionAnsweringModelOutput(
- loss=loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convbert", None) is not None:
- with tf.name_scope(self.convbert.name):
- self.convbert.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.config.hidden_size])
diff --git a/transformers/models/convbert/tokenization_convbert.py b/transformers/models/convbert/tokenization_convbert.py
deleted file mode 100644
index c0fe2c018341c55b2446a6f12052a3fa36bd9246..0000000000000000000000000000000000000000
--- a/transformers/models/convbert/tokenization_convbert.py
+++ /dev/null
@@ -1,503 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for ConvBERT."""
-import collections
-import os
-import unicodedata
-from typing import List, Optional, Tuple
-
-from ...tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
-
-
-# Copied from transformers.models.bert.tokenization_bert.load_vocab
-def load_vocab(vocab_file):
- """Loads a vocabulary file into a dictionary."""
- vocab = collections.OrderedDict()
- with open(vocab_file, "r", encoding="utf-8") as reader:
- tokens = reader.readlines()
- for index, token in enumerate(tokens):
- token = token.rstrip("\n")
- vocab[token] = index
- return vocab
-
-
-# Copied from transformers.models.bert.tokenization_bert.whitespace_tokenize
-def whitespace_tokenize(text):
- """Runs basic whitespace cleaning and splitting on a piece of text."""
- text = text.strip()
- if not text:
- return []
- tokens = text.split()
- return tokens
-
-
-# Copied from transformers.models.bert.tokenization_bert.BertTokenizer with bert-base-cased->YituTech/conv-bert-base, ConvBertTokenizer->BertTokenizer, BERT->ConvBERT
-class ConvBertTokenizer(PreTrainedTokenizer):
- r"""
- Construct a ConvBERT tokenizer. Based on WordPiece.
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- do_basic_tokenize (`bool`, *optional*, defaults to `True`):
- Whether or not to do basic tokenization before WordPiece.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original ConvBERT).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
-
- def __init__(
- self,
- vocab_file,
- do_lower_case=True,
- do_basic_tokenize=True,
- never_split=None,
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- if not os.path.isfile(vocab_file):
- raise ValueError(
- f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
- " model use `tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- self.vocab = load_vocab(vocab_file)
- self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
- self.do_basic_tokenize = do_basic_tokenize
- if do_basic_tokenize:
- self.basic_tokenizer = BasicTokenizer(
- do_lower_case=do_lower_case,
- never_split=never_split,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- )
-
- self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token))
-
- super().__init__(
- do_lower_case=do_lower_case,
- do_basic_tokenize=do_basic_tokenize,
- never_split=never_split,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- @property
- def do_lower_case(self):
- return self.basic_tokenizer.do_lower_case
-
- @property
- def vocab_size(self):
- return len(self.vocab)
-
- def get_vocab(self):
- return dict(self.vocab, **self.added_tokens_encoder)
-
- def _tokenize(self, text, split_special_tokens=False):
- split_tokens = []
- if self.do_basic_tokenize:
- for token in self.basic_tokenizer.tokenize(
- text, never_split=self.all_special_tokens if not split_special_tokens else None
- ):
- # If the token is part of the never_split set
- if token in self.basic_tokenizer.never_split:
- split_tokens.append(token)
- else:
- split_tokens += self.wordpiece_tokenizer.tokenize(token)
- else:
- split_tokens = self.wordpiece_tokenizer.tokenize(text)
- return split_tokens
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.vocab.get(token, self.vocab.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.ids_to_tokens.get(index, self.unk_token)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- out_string = " ".join(tokens).replace(" ##", "").strip()
- return out_string
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A ConvBERT sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A ConvBERT sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- index = 0
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- else:
- vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(token + "\n")
- index += 1
- return (vocab_file,)
-
-
-# Copied from transformers.models.bert.tokenization_bert.BasicTokenizer
-class BasicTokenizer(object):
- """
- Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
-
- Args:
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- do_split_on_punc (`bool`, *optional*, defaults to `True`):
- In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
- the full context of the words, such as contractions.
- """
-
- def __init__(
- self,
- do_lower_case=True,
- never_split=None,
- tokenize_chinese_chars=True,
- strip_accents=None,
- do_split_on_punc=True,
- ):
- if never_split is None:
- never_split = []
- self.do_lower_case = do_lower_case
- self.never_split = set(never_split)
- self.tokenize_chinese_chars = tokenize_chinese_chars
- self.strip_accents = strip_accents
- self.do_split_on_punc = do_split_on_punc
-
- def tokenize(self, text, never_split=None):
- """
- Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
-
- Args:
- never_split (`List[str]`, *optional*)
- Kept for backward compatibility purposes. Now implemented directly at the base class level (see
- [`PreTrainedTokenizer.tokenize`]) List of token not to split.
- """
- # union() returns a new set by concatenating the two sets.
- never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
- text = self._clean_text(text)
-
- # This was added on November 1st, 2018 for the multilingual and Chinese
- # models. This is also applied to the English models now, but it doesn't
- # matter since the English models were not trained on any Chinese data
- # and generally don't have any Chinese data in them (there are Chinese
- # characters in the vocabulary because Wikipedia does have some Chinese
- # words in the English Wikipedia.).
- if self.tokenize_chinese_chars:
- text = self._tokenize_chinese_chars(text)
- # prevents treating the same character with different unicode codepoints as different characters
- unicode_normalized_text = unicodedata.normalize("NFC", text)
- orig_tokens = whitespace_tokenize(unicode_normalized_text)
- split_tokens = []
- for token in orig_tokens:
- if token not in never_split:
- if self.do_lower_case:
- token = token.lower()
- if self.strip_accents is not False:
- token = self._run_strip_accents(token)
- elif self.strip_accents:
- token = self._run_strip_accents(token)
- split_tokens.extend(self._run_split_on_punc(token, never_split))
-
- output_tokens = whitespace_tokenize(" ".join(split_tokens))
- return output_tokens
-
- def _run_strip_accents(self, text):
- """Strips accents from a piece of text."""
- text = unicodedata.normalize("NFD", text)
- output = []
- for char in text:
- cat = unicodedata.category(char)
- if cat == "Mn":
- continue
- output.append(char)
- return "".join(output)
-
- def _run_split_on_punc(self, text, never_split=None):
- """Splits punctuation on a piece of text."""
- if not self.do_split_on_punc or (never_split is not None and text in never_split):
- return [text]
- chars = list(text)
- i = 0
- start_new_word = True
- output = []
- while i < len(chars):
- char = chars[i]
- if _is_punctuation(char):
- output.append([char])
- start_new_word = True
- else:
- if start_new_word:
- output.append([])
- start_new_word = False
- output[-1].append(char)
- i += 1
-
- return ["".join(x) for x in output]
-
- def _tokenize_chinese_chars(self, text):
- """Adds whitespace around any CJK character."""
- output = []
- for char in text:
- cp = ord(char)
- if self._is_chinese_char(cp):
- output.append(" ")
- output.append(char)
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
- def _is_chinese_char(self, cp):
- """Checks whether CP is the codepoint of a CJK character."""
- # This defines a "chinese character" as anything in the CJK Unicode block:
- # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
- #
- # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
- # despite its name. The modern Korean Hangul alphabet is a different block,
- # as is Japanese Hiragana and Katakana. Those alphabets are used to write
- # space-separated words, so they are not treated specially and handled
- # like the all of the other languages.
- if (
- (cp >= 0x4E00 and cp <= 0x9FFF)
- or (cp >= 0x3400 and cp <= 0x4DBF) #
- or (cp >= 0x20000 and cp <= 0x2A6DF) #
- or (cp >= 0x2A700 and cp <= 0x2B73F) #
- or (cp >= 0x2B740 and cp <= 0x2B81F) #
- or (cp >= 0x2B820 and cp <= 0x2CEAF) #
- or (cp >= 0xF900 and cp <= 0xFAFF)
- or (cp >= 0x2F800 and cp <= 0x2FA1F) #
- ): #
- return True
-
- return False
-
- def _clean_text(self, text):
- """Performs invalid character removal and whitespace cleanup on text."""
- output = []
- for char in text:
- cp = ord(char)
- if cp == 0 or cp == 0xFFFD or _is_control(char):
- continue
- if _is_whitespace(char):
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
-
-# Copied from transformers.models.bert.tokenization_bert.WordpieceTokenizer
-class WordpieceTokenizer(object):
- """Runs WordPiece tokenization."""
-
- def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
- self.vocab = vocab
- self.unk_token = unk_token
- self.max_input_chars_per_word = max_input_chars_per_word
-
- def tokenize(self, text):
- """
- Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
- tokenization using the given vocabulary.
-
- For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
-
- Args:
- text: A single token or whitespace separated tokens. This should have
- already been passed through *BasicTokenizer*.
-
- Returns:
- A list of wordpiece tokens.
- """
-
- output_tokens = []
- for token in whitespace_tokenize(text):
- chars = list(token)
- if len(chars) > self.max_input_chars_per_word:
- output_tokens.append(self.unk_token)
- continue
-
- is_bad = False
- start = 0
- sub_tokens = []
- while start < len(chars):
- end = len(chars)
- cur_substr = None
- while start < end:
- substr = "".join(chars[start:end])
- if start > 0:
- substr = "##" + substr
- if substr in self.vocab:
- cur_substr = substr
- break
- end -= 1
- if cur_substr is None:
- is_bad = True
- break
- sub_tokens.append(cur_substr)
- start = end
-
- if is_bad:
- output_tokens.append(self.unk_token)
- else:
- output_tokens.extend(sub_tokens)
- return output_tokens
diff --git a/transformers/models/convbert/tokenization_convbert_fast.py b/transformers/models/convbert/tokenization_convbert_fast.py
deleted file mode 100644
index 65bedb73fe9171f0473c7d5a35c08dee7432eb78..0000000000000000000000000000000000000000
--- a/transformers/models/convbert/tokenization_convbert_fast.py
+++ /dev/null
@@ -1,172 +0,0 @@
-# coding=utf-8
-# Copyright The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for ConvBERT."""
-import json
-from typing import List, Optional, Tuple
-
-from tokenizers import normalizers
-
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import logging
-from .tokenization_convbert import ConvBertTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
-
-
-# Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast with bert-base-cased->YituTech/conv-bert-base, Bert->ConvBert, BERT->ConvBERT
-class ConvBertTokenizerFast(PreTrainedTokenizerFast):
- r"""
- Construct a "fast" ConvBERT tokenizer (backed by HuggingFace's *tokenizers* library). Based on WordPiece.
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- clean_text (`bool`, *optional*, defaults to `True`):
- Whether or not to clean the text before tokenization by removing any control characters and replacing all
- whitespaces by the classic one.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see [this
- issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original ConvBERT).
- wordpieces_prefix (`str`, *optional*, defaults to `"##"`):
- The prefix for subwords.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = ConvBertTokenizer
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- do_lower_case=True,
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- super().__init__(
- vocab_file,
- tokenizer_file=tokenizer_file,
- do_lower_case=do_lower_case,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- normalizer_state = json.loads(self.backend_tokenizer.normalizer.__getstate__())
- if (
- normalizer_state.get("lowercase", do_lower_case) != do_lower_case
- or normalizer_state.get("strip_accents", strip_accents) != strip_accents
- or normalizer_state.get("handle_chinese_chars", tokenize_chinese_chars) != tokenize_chinese_chars
- ):
- normalizer_class = getattr(normalizers, normalizer_state.pop("type"))
- normalizer_state["lowercase"] = do_lower_case
- normalizer_state["strip_accents"] = strip_accents
- normalizer_state["handle_chinese_chars"] = tokenize_chinese_chars
- self.backend_tokenizer.normalizer = normalizer_class(**normalizer_state)
-
- self.do_lower_case = do_lower_case
-
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A ConvBERT sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- output = [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
-
- if token_ids_1 is not None:
- output += token_ids_1 + [self.sep_token_id]
-
- return output
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A ConvBERT sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
diff --git a/transformers/models/convnext/__init__.py b/transformers/models/convnext/__init__.py
deleted file mode 100644
index 099a7fc9d63da4ef2cbe0308371d7b26d586e447..0000000000000000000000000000000000000000
--- a/transformers/models/convnext/__init__.py
+++ /dev/null
@@ -1,102 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_tf_available,
- is_torch_available,
- is_vision_available,
-)
-
-
-_import_structure = {
- "configuration_convnext": ["CONVNEXT_PRETRAINED_CONFIG_ARCHIVE_MAP", "ConvNextConfig", "ConvNextOnnxConfig"]
-}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_convnext"] = ["ConvNextFeatureExtractor"]
- _import_structure["image_processing_convnext"] = ["ConvNextImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_convnext"] = [
- "CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "ConvNextForImageClassification",
- "ConvNextModel",
- "ConvNextPreTrainedModel",
- "ConvNextBackbone",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_convnext"] = [
- "TFConvNextForImageClassification",
- "TFConvNextModel",
- "TFConvNextPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_convnext import CONVNEXT_PRETRAINED_CONFIG_ARCHIVE_MAP, ConvNextConfig, ConvNextOnnxConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_convnext import ConvNextFeatureExtractor
- from .image_processing_convnext import ConvNextImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_convnext import (
- CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
- ConvNextBackbone,
- ConvNextForImageClassification,
- ConvNextModel,
- ConvNextPreTrainedModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_convnext import TFConvNextForImageClassification, TFConvNextModel, TFConvNextPreTrainedModel
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
diff --git a/transformers/models/convnext/__pycache__/__init__.cpython-310.pyc b/transformers/models/convnext/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 81618be669f68545a7b09e4b406d131c13b10f03..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnext/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnext/__pycache__/configuration_convnext.cpython-310.pyc b/transformers/models/convnext/__pycache__/configuration_convnext.cpython-310.pyc
deleted file mode 100644
index 0a3f9cf299ea3352671b6fc76b6efaef5ea550d5..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnext/__pycache__/configuration_convnext.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnext/__pycache__/convert_convnext_to_pytorch.cpython-310.pyc b/transformers/models/convnext/__pycache__/convert_convnext_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 2c2f6a9ee5d539cf16ad30163303f84c3f3fbfeb..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnext/__pycache__/convert_convnext_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnext/__pycache__/feature_extraction_convnext.cpython-310.pyc b/transformers/models/convnext/__pycache__/feature_extraction_convnext.cpython-310.pyc
deleted file mode 100644
index a22079ca66e3f73de70d30a94988ec4065ec4566..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnext/__pycache__/feature_extraction_convnext.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnext/__pycache__/image_processing_convnext.cpython-310.pyc b/transformers/models/convnext/__pycache__/image_processing_convnext.cpython-310.pyc
deleted file mode 100644
index c2773cecaf6fb739cdae7e514feaf9cfd7e2ccc7..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnext/__pycache__/image_processing_convnext.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnext/__pycache__/modeling_convnext.cpython-310.pyc b/transformers/models/convnext/__pycache__/modeling_convnext.cpython-310.pyc
deleted file mode 100644
index 989c82513062944c145fdb263220797f03ac3268..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnext/__pycache__/modeling_convnext.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnext/__pycache__/modeling_tf_convnext.cpython-310.pyc b/transformers/models/convnext/__pycache__/modeling_tf_convnext.cpython-310.pyc
deleted file mode 100644
index 8aad0d11320992a4807fe385c2944309ae452064..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnext/__pycache__/modeling_tf_convnext.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnext/configuration_convnext.py b/transformers/models/convnext/configuration_convnext.py
deleted file mode 100644
index f84c31079ea34eddcc6b02e2b7fe51c647b8ab1d..0000000000000000000000000000000000000000
--- a/transformers/models/convnext/configuration_convnext.py
+++ /dev/null
@@ -1,142 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" ConvNeXT model configuration"""
-
-from collections import OrderedDict
-from typing import Mapping
-
-from packaging import version
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import CONVNEXT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class ConvNextConfig(BackboneConfigMixin, PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`ConvNextModel`]. It is used to instantiate an
- ConvNeXT model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the ConvNeXT
- [facebook/convnext-tiny-224](https://huggingface.co/facebook/convnext-tiny-224) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- patch_size (`int`, optional, defaults to 4):
- Patch size to use in the patch embedding layer.
- num_stages (`int`, optional, defaults to 4):
- The number of stages in the model.
- hidden_sizes (`List[int]`, *optional*, defaults to [96, 192, 384, 768]):
- Dimensionality (hidden size) at each stage.
- depths (`List[int]`, *optional*, defaults to [3, 3, 9, 3]):
- Depth (number of blocks) for each stage.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in each block. If string, `"gelu"`, `"relu"`,
- `"selu"` and `"gelu_new"` are supported.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- layer_scale_init_value (`float`, *optional*, defaults to 1e-6):
- The initial value for the layer scale.
- drop_path_rate (`float`, *optional*, defaults to 0.0):
- The drop rate for stochastic depth.
- out_features (`List[str]`, *optional*):
- If used as backbone, list of features to output. Can be any of `"stem"`, `"stage1"`, `"stage2"`, etc.
- (depending on how many stages the model has). If unset and `out_indices` is set, will default to the
- corresponding stages. If unset and `out_indices` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
- out_indices (`List[int]`, *optional*):
- If used as backbone, list of indices of features to output. Can be any of 0, 1, 2, etc. (depending on how
- many stages the model has). If unset and `out_features` is set, will default to the corresponding stages.
- If unset and `out_features` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
-
- Example:
- ```python
- >>> from transformers import ConvNextConfig, ConvNextModel
-
- >>> # Initializing a ConvNext convnext-tiny-224 style configuration
- >>> configuration = ConvNextConfig()
-
- >>> # Initializing a model (with random weights) from the convnext-tiny-224 style configuration
- >>> model = ConvNextModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "convnext"
-
- def __init__(
- self,
- num_channels=3,
- patch_size=4,
- num_stages=4,
- hidden_sizes=None,
- depths=None,
- hidden_act="gelu",
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- layer_scale_init_value=1e-6,
- drop_path_rate=0.0,
- image_size=224,
- out_features=None,
- out_indices=None,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.num_channels = num_channels
- self.patch_size = patch_size
- self.num_stages = num_stages
- self.hidden_sizes = [96, 192, 384, 768] if hidden_sizes is None else hidden_sizes
- self.depths = [3, 3, 9, 3] if depths is None else depths
- self.hidden_act = hidden_act
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.layer_scale_init_value = layer_scale_init_value
- self.drop_path_rate = drop_path_rate
- self.image_size = image_size
- self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
- self._out_features, self._out_indices = get_aligned_output_features_output_indices(
- out_features=out_features, out_indices=out_indices, stage_names=self.stage_names
- )
-
-
-class ConvNextOnnxConfig(OnnxConfig):
- torch_onnx_minimum_version = version.parse("1.11")
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- return OrderedDict(
- [
- ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
- ]
- )
-
- @property
- def atol_for_validation(self) -> float:
- return 1e-5
diff --git a/transformers/models/convnext/convert_convnext_to_pytorch.py b/transformers/models/convnext/convert_convnext_to_pytorch.py
deleted file mode 100644
index cdcbf24d552389ba34f55c8fa1af717aa26dd60f..0000000000000000000000000000000000000000
--- a/transformers/models/convnext/convert_convnext_to_pytorch.py
+++ /dev/null
@@ -1,243 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert ConvNext checkpoints from the original repository.
-
-URL: https://github.com/facebookresearch/ConvNeXt"""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_convnext_config(checkpoint_url):
- config = ConvNextConfig()
-
- if "tiny" in checkpoint_url:
- depths = [3, 3, 9, 3]
- hidden_sizes = [96, 192, 384, 768]
- if "small" in checkpoint_url:
- depths = [3, 3, 27, 3]
- hidden_sizes = [96, 192, 384, 768]
- if "base" in checkpoint_url:
- depths = [3, 3, 27, 3]
- hidden_sizes = [128, 256, 512, 1024]
- if "large" in checkpoint_url:
- depths = [3, 3, 27, 3]
- hidden_sizes = [192, 384, 768, 1536]
- if "xlarge" in checkpoint_url:
- depths = [3, 3, 27, 3]
- hidden_sizes = [256, 512, 1024, 2048]
-
- if "1k" in checkpoint_url:
- num_labels = 1000
- filename = "imagenet-1k-id2label.json"
- expected_shape = (1, 1000)
- else:
- num_labels = 21841
- filename = "imagenet-22k-id2label.json"
- expected_shape = (1, 21841)
-
- repo_id = "huggingface/label-files"
- config.num_labels = num_labels
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- if "1k" not in checkpoint_url:
- # this dataset contains 21843 labels but the model only has 21841
- # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
- del id2label[9205]
- del id2label[15027]
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
- config.hidden_sizes = hidden_sizes
- config.depths = depths
-
- return config, expected_shape
-
-
-def rename_key(name):
- if "downsample_layers.0.0" in name:
- name = name.replace("downsample_layers.0.0", "embeddings.patch_embeddings")
- if "downsample_layers.0.1" in name:
- name = name.replace("downsample_layers.0.1", "embeddings.norm") # we rename to layernorm later on
- if "downsample_layers.1.0" in name:
- name = name.replace("downsample_layers.1.0", "stages.1.downsampling_layer.0")
- if "downsample_layers.1.1" in name:
- name = name.replace("downsample_layers.1.1", "stages.1.downsampling_layer.1")
- if "downsample_layers.2.0" in name:
- name = name.replace("downsample_layers.2.0", "stages.2.downsampling_layer.0")
- if "downsample_layers.2.1" in name:
- name = name.replace("downsample_layers.2.1", "stages.2.downsampling_layer.1")
- if "downsample_layers.3.0" in name:
- name = name.replace("downsample_layers.3.0", "stages.3.downsampling_layer.0")
- if "downsample_layers.3.1" in name:
- name = name.replace("downsample_layers.3.1", "stages.3.downsampling_layer.1")
- if "stages" in name and "downsampling_layer" not in name:
- # stages.0.0. for instance should be renamed to stages.0.layers.0.
- name = name[: len("stages.0")] + ".layers" + name[len("stages.0") :]
- if "stages" in name:
- name = name.replace("stages", "encoder.stages")
- if "norm" in name:
- name = name.replace("norm", "layernorm")
- if "gamma" in name:
- name = name.replace("gamma", "layer_scale_parameter")
- if "head" in name:
- name = name.replace("head", "classifier")
-
- return name
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-@torch.no_grad()
-def convert_convnext_checkpoint(checkpoint_url, pytorch_dump_folder_path):
- """
- Copy/paste/tweak model's weights to our ConvNext structure.
- """
-
- # define ConvNext configuration based on URL
- config, expected_shape = get_convnext_config(checkpoint_url)
- # load original state_dict from URL
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url)["model"]
- # rename keys
- for key in state_dict.copy().keys():
- val = state_dict.pop(key)
- state_dict[rename_key(key)] = val
- # add prefix to all keys expect classifier head
- for key in state_dict.copy().keys():
- val = state_dict.pop(key)
- if not key.startswith("classifier"):
- key = "convnext." + key
- state_dict[key] = val
-
- # load HuggingFace model
- model = ConvNextForImageClassification(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- # Check outputs on an image, prepared by ConvNextImageProcessor
- size = 224 if "224" in checkpoint_url else 384
- image_processor = ConvNextImageProcessor(size=size)
- pixel_values = image_processor(images=prepare_img(), return_tensors="pt").pixel_values
-
- logits = model(pixel_values).logits
-
- # note: the logits below were obtained without center cropping
- if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
- expected_logits = torch.tensor([-0.1210, -0.6605, 0.1918])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
- expected_logits = torch.tensor([-0.4473, -0.1847, -0.6365])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
- expected_logits = torch.tensor([0.4525, 0.7539, 0.0308])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
- expected_logits = torch.tensor([0.3561, 0.6350, -0.0384])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
- expected_logits = torch.tensor([0.4174, -0.0989, 0.1489])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
- expected_logits = torch.tensor([0.2513, -0.1349, -0.1613])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
- expected_logits = torch.tensor([1.2980, 0.3631, -0.1198])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
- expected_logits = torch.tensor([1.2963, 0.1227, 0.1723])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
- expected_logits = torch.tensor([1.7956, 0.8390, 0.2820])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
- expected_logits = torch.tensor([-0.2822, -0.0502, -0.0878])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
- expected_logits = torch.tensor([-0.5672, -0.0730, -0.4348])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
- expected_logits = torch.tensor([0.2681, 0.2365, 0.6246])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
- expected_logits = torch.tensor([-0.2642, 0.3931, 0.5116])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
- expected_logits = torch.tensor([-0.6677, -0.1873, -0.8379])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
- expected_logits = torch.tensor([-0.7749, -0.2967, -0.6444])
- else:
- raise ValueError(f"Unknown URL: {checkpoint_url}")
-
- assert torch.allclose(logits[0, :3], expected_logits, atol=1e-3)
- assert logits.shape == expected_shape
-
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- print(f"Saving image processor to {pytorch_dump_folder_path}")
- image_processor.save_pretrained(pytorch_dump_folder_path)
-
- print("Pushing model to the hub...")
- model_name = "convnext"
- if "tiny" in checkpoint_url:
- model_name += "-tiny"
- elif "small" in checkpoint_url:
- model_name += "-small"
- elif "base" in checkpoint_url:
- model_name += "-base"
- elif "xlarge" in checkpoint_url:
- model_name += "-xlarge"
- elif "large" in checkpoint_url:
- model_name += "-large"
- if "224" in checkpoint_url:
- model_name += "-224"
- elif "384" in checkpoint_url:
- model_name += "-384"
- if "22k" in checkpoint_url and "1k" not in checkpoint_url:
- model_name += "-22k"
- if "22k" in checkpoint_url and "1k" in checkpoint_url:
- model_name += "-22k-1k"
-
- model.push_to_hub(
- repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
- organization="nielsr",
- commit_message="Add model",
- )
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--checkpoint_url",
- default="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
- type=str,
- help="URL of the original ConvNeXT checkpoint you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- required=True,
- help="Path to the output PyTorch model directory.",
- )
-
- args = parser.parse_args()
- convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)
diff --git a/transformers/models/convnext/feature_extraction_convnext.py b/transformers/models/convnext/feature_extraction_convnext.py
deleted file mode 100644
index 92b8a8f4fba82fb72b83384d2cbcb6abfe773ea2..0000000000000000000000000000000000000000
--- a/transformers/models/convnext/feature_extraction_convnext.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for ConvNeXT."""
-
-import warnings
-
-from ...utils import logging
-from .image_processing_convnext import ConvNextImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-class ConvNextFeatureExtractor(ConvNextImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class ConvNextFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
- " Please use ConvNextImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/convnext/image_processing_convnext.py b/transformers/models/convnext/image_processing_convnext.py
deleted file mode 100644
index 54060105f59eb264af6d2ee5c58c8308e0a8fa49..0000000000000000000000000000000000000000
--- a/transformers/models/convnext/image_processing_convnext.py
+++ /dev/null
@@ -1,338 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for ConvNeXT."""
-
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
-from ...image_transforms import (
- center_crop,
- get_resize_output_image_size,
- resize,
- to_channel_dimension_format,
-)
-from ...image_utils import (
- IMAGENET_STANDARD_MEAN,
- IMAGENET_STANDARD_STD,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import TensorType, is_vision_available, logging
-
-
-if is_vision_available():
- import PIL
-
-
-logger = logging.get_logger(__name__)
-
-
-class ConvNextImageProcessor(BaseImageProcessor):
- r"""
- Constructs a ConvNeXT image processor.
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Controls whether to resize the image's (height, width) dimensions to the specified `size`. Can be overriden
- by `do_resize` in the `preprocess` method.
- size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 384}`):
- Resolution of the output image after `resize` is applied. If `size["shortest_edge"]` >= 384, the image is
- resized to `(size["shortest_edge"], size["shortest_edge"])`. Otherwise, the smaller edge of the image will
- be matched to `int(size["shortest_edge"]/crop_pct)`, after which the image is cropped to
- `(size["shortest_edge"], size["shortest_edge"])`. Only has an effect if `do_resize` is set to `True`. Can
- be overriden by `size` in the `preprocess` method.
- crop_pct (`float` *optional*, defaults to 224 / 256):
- Percentage of the image to crop. Only has an effect if `do_resize` is `True` and size < 384. Can be
- overriden by `crop_pct` in the `preprocess` method.
- resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
- Resampling filter to use if resizing the image. Can be overriden by `resample` in the `preprocess` method.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image by the specified scale `rescale_factor`. Can be overriden by `do_rescale` in
- the `preprocess` method.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overriden by `rescale_factor` in the `preprocess`
- method.
- do_normalize (`bool`, *optional*, defaults to `True`):
- Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
- method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
- Mean to use if normalizing the image. This is a float or list of floats the length of the number of
- channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
- Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
- number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
- """
-
- model_input_names = ["pixel_values"]
-
- def __init__(
- self,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- crop_pct: float = None,
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- do_normalize: bool = True,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- size = size if size is not None else {"shortest_edge": 384}
- size = get_size_dict(size, default_to_square=False)
-
- self.do_resize = do_resize
- self.size = size
- # Default value set here for backwards compatibility where the value in config is None
- self.crop_pct = crop_pct if crop_pct is not None else 224 / 256
- self.resample = resample
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_normalize = do_normalize
- self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size",
- "crop_pct",
- "resample",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "image_mean",
- "image_std",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- crop_pct: float,
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize an image.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary of the form `{"shortest_edge": int}`, specifying the size of the output image. If
- `size["shortest_edge"]` >= 384 image is resized to `(size["shortest_edge"], size["shortest_edge"])`.
- Otherwise, the smaller edge of the image will be matched to `int(size["shortest_edge"] / crop_pct)`,
- after which the image is cropped to `(size["shortest_edge"], size["shortest_edge"])`.
- crop_pct (`float`):
- Percentage of the image to crop. Only has an effect if size < 384.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
- Resampling filter to use when resizing the image.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the image. If not provided, it will be the same as the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred from the input
- image.
- """
- size = get_size_dict(size, default_to_square=False)
- if "shortest_edge" not in size:
- raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {size.keys()}")
- shortest_edge = size["shortest_edge"]
-
- if shortest_edge < 384:
- # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
- resize_shortest_edge = int(shortest_edge / crop_pct)
- resize_size = get_resize_output_image_size(
- image, size=resize_shortest_edge, default_to_square=False, input_data_format=input_data_format
- )
- image = resize(
- image=image,
- size=resize_size,
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
- # then crop to (shortest_edge, shortest_edge)
- return center_crop(
- image=image,
- size=(shortest_edge, shortest_edge),
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
- else:
- # warping (no cropping) when evaluated at 384 or larger
- return resize(
- image,
- size=(shortest_edge, shortest_edge),
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
-
- def preprocess(
- self,
- images: ImageInput,
- do_resize: bool = None,
- size: Dict[str, int] = None,
- crop_pct: float = None,
- resample: PILImageResampling = None,
- do_rescale: bool = None,
- rescale_factor: float = None,
- do_normalize: bool = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: ChannelDimension = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> PIL.Image.Image:
- """
- Preprocess an image or batch of images.
-
- Args:
- images (`ImageInput`):
- Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
- passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
- Size of the output image after `resize` has been applied. If `size["shortest_edge"]` >= 384, the image
- is resized to `(size["shortest_edge"], size["shortest_edge"])`. Otherwise, the smaller edge of the
- image will be matched to `int(size["shortest_edge"]/ crop_pct)`, after which the image is cropped to
- `(size["shortest_edge"], size["shortest_edge"])`. Only has an effect if `do_resize` is set to `True`.
- crop_pct (`float`, *optional*, defaults to `self.crop_pct`):
- Percentage of the image to crop if size < 384.
- resample (`int`, *optional*, defaults to `self.resample`):
- Resampling filter to use if resizing the image. This can be one of `PILImageResampling`, filters. Only
- has an effect if `do_resize` is set to `True`.
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether to rescale the image values between [0 - 1].
- rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
- Rescale factor to rescale the image by if `do_rescale` is set to `True`.
- do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
- Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
- Image mean.
- image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
- Image standard deviation.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - Unset: Use the channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- do_resize = do_resize if do_resize is not None else self.do_resize
- crop_pct = crop_pct if crop_pct is not None else self.crop_pct
- resample = resample if resample is not None else self.resample
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- do_normalize = do_normalize if do_normalize is not None else self.do_normalize
- image_mean = image_mean if image_mean is not None else self.image_mean
- image_std = image_std if image_std is not None else self.image_std
-
- size = size if size is not None else self.size
- size = get_size_dict(size, default_to_square=False)
-
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- images = make_list_of_images(images)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
-
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
-
- # All transformations expect numpy arrays.
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- if do_resize:
- images = [
- self.resize(
- image=image, size=size, crop_pct=crop_pct, resample=resample, input_data_format=input_data_format
- )
- for image in images
- ]
-
- if do_rescale:
- images = [
- self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_normalize:
- images = [
- self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
- for image in images
- ]
-
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
- ]
-
- data = {"pixel_values": images}
- return BatchFeature(data=data, tensor_type=return_tensors)
diff --git a/transformers/models/convnext/modeling_convnext.py b/transformers/models/convnext/modeling_convnext.py
deleted file mode 100644
index 147d2ac22dac4525be7ed1924d9aad50d5b09cc1..0000000000000000000000000000000000000000
--- a/transformers/models/convnext/modeling_convnext.py
+++ /dev/null
@@ -1,551 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch ConvNext model."""
-
-
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BackboneOutput,
- BaseModelOutputWithNoAttention,
- BaseModelOutputWithPoolingAndNoAttention,
- ImageClassifierOutputWithNoAttention,
-)
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from ...utils.backbone_utils import BackboneMixin
-from .configuration_convnext import ConvNextConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "ConvNextConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/convnext-tiny-224"
-_EXPECTED_OUTPUT_SHAPE = [1, 768, 7, 7]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "facebook/convnext-tiny-224"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->ConvNext
-class ConvNextDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-class ConvNextLayerNorm(nn.Module):
- r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
- The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
- width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).
- """
-
- def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
- super().__init__()
- self.weight = nn.Parameter(torch.ones(normalized_shape))
- self.bias = nn.Parameter(torch.zeros(normalized_shape))
- self.eps = eps
- self.data_format = data_format
- if self.data_format not in ["channels_last", "channels_first"]:
- raise NotImplementedError(f"Unsupported data format: {self.data_format}")
- self.normalized_shape = (normalized_shape,)
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- if self.data_format == "channels_last":
- x = torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
- elif self.data_format == "channels_first":
- input_dtype = x.dtype
- x = x.float()
- u = x.mean(1, keepdim=True)
- s = (x - u).pow(2).mean(1, keepdim=True)
- x = (x - u) / torch.sqrt(s + self.eps)
- x = x.to(dtype=input_dtype)
- x = self.weight[:, None, None] * x + self.bias[:, None, None]
- return x
-
-
-class ConvNextEmbeddings(nn.Module):
- """This class is comparable to (and inspired by) the SwinEmbeddings class
- found in src/transformers/models/swin/modeling_swin.py.
- """
-
- def __init__(self, config):
- super().__init__()
- self.patch_embeddings = nn.Conv2d(
- config.num_channels, config.hidden_sizes[0], kernel_size=config.patch_size, stride=config.patch_size
- )
- self.layernorm = ConvNextLayerNorm(config.hidden_sizes[0], eps=1e-6, data_format="channels_first")
- self.num_channels = config.num_channels
-
- def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
- num_channels = pixel_values.shape[1]
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- embeddings = self.patch_embeddings(pixel_values)
- embeddings = self.layernorm(embeddings)
- return embeddings
-
-
-class ConvNextLayer(nn.Module):
- """This corresponds to the `Block` class in the original implementation.
-
- There are two equivalent implementations: [DwConv, LayerNorm (channels_first), Conv, GELU,1x1 Conv]; all in (N, C,
- H, W) (2) [DwConv, Permute to (N, H, W, C), LayerNorm (channels_last), Linear, GELU, Linear]; Permute back
-
- The authors used (2) as they find it slightly faster in PyTorch.
-
- Args:
- config ([`ConvNextConfig`]): Model configuration class.
- dim (`int`): Number of input channels.
- drop_path (`float`): Stochastic depth rate. Default: 0.0.
- """
-
- def __init__(self, config, dim, drop_path=0):
- super().__init__()
- self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
- self.layernorm = ConvNextLayerNorm(dim, eps=1e-6)
- self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
- self.act = ACT2FN[config.hidden_act]
- self.pwconv2 = nn.Linear(4 * dim, dim)
- self.layer_scale_parameter = (
- nn.Parameter(config.layer_scale_init_value * torch.ones((dim)), requires_grad=True)
- if config.layer_scale_init_value > 0
- else None
- )
- self.drop_path = ConvNextDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
- input = hidden_states
- x = self.dwconv(hidden_states)
- x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
- x = self.layernorm(x)
- x = self.pwconv1(x)
- x = self.act(x)
- x = self.pwconv2(x)
- if self.layer_scale_parameter is not None:
- x = self.layer_scale_parameter * x
- x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
-
- x = input + self.drop_path(x)
- return x
-
-
-class ConvNextStage(nn.Module):
- """ConvNeXT stage, consisting of an optional downsampling layer + multiple residual blocks.
-
- Args:
- config ([`ConvNextConfig`]): Model configuration class.
- in_channels (`int`): Number of input channels.
- out_channels (`int`): Number of output channels.
- depth (`int`): Number of residual blocks.
- drop_path_rates(`List[float]`): Stochastic depth rates for each layer.
- """
-
- def __init__(self, config, in_channels, out_channels, kernel_size=2, stride=2, depth=2, drop_path_rates=None):
- super().__init__()
-
- if in_channels != out_channels or stride > 1:
- self.downsampling_layer = nn.Sequential(
- ConvNextLayerNorm(in_channels, eps=1e-6, data_format="channels_first"),
- nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride),
- )
- else:
- self.downsampling_layer = nn.Identity()
- drop_path_rates = drop_path_rates or [0.0] * depth
- self.layers = nn.Sequential(
- *[ConvNextLayer(config, dim=out_channels, drop_path=drop_path_rates[j]) for j in range(depth)]
- )
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
- hidden_states = self.downsampling_layer(hidden_states)
- hidden_states = self.layers(hidden_states)
- return hidden_states
-
-
-class ConvNextEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.stages = nn.ModuleList()
- drop_path_rates = [
- x.tolist() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths)).split(config.depths)
- ]
- prev_chs = config.hidden_sizes[0]
- for i in range(config.num_stages):
- out_chs = config.hidden_sizes[i]
- stage = ConvNextStage(
- config,
- in_channels=prev_chs,
- out_channels=out_chs,
- stride=2 if i > 0 else 1,
- depth=config.depths[i],
- drop_path_rates=drop_path_rates[i],
- )
- self.stages.append(stage)
- prev_chs = out_chs
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple, BaseModelOutputWithNoAttention]:
- all_hidden_states = () if output_hidden_states else None
-
- for i, layer_module in enumerate(self.stages):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- hidden_states = layer_module(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
-
- return BaseModelOutputWithNoAttention(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- )
-
-
-class ConvNextPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ConvNextConfig
- base_model_prefix = "convnext"
- main_input_name = "pixel_values"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-CONVNEXT_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`ConvNextConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CONVNEXT_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`ConvNextImageProcessor.__call__`] for details.
-
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare ConvNext model outputting raw features without any specific head on top.",
- CONVNEXT_START_DOCSTRING,
-)
-class ConvNextModel(ConvNextPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.config = config
-
- self.embeddings = ConvNextEmbeddings(config)
- self.encoder = ConvNextEncoder(config)
-
- # final layernorm layer
- self.layernorm = nn.LayerNorm(config.hidden_sizes[-1], eps=config.layer_norm_eps)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVNEXT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPoolingAndNoAttention,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: torch.FloatTensor = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPoolingAndNoAttention]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- embedding_output = self.embeddings(pixel_values)
-
- encoder_outputs = self.encoder(
- embedding_output,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden_state = encoder_outputs[0]
-
- # global average pooling, (N, C, H, W) -> (N, C)
- pooled_output = self.layernorm(last_hidden_state.mean([-2, -1]))
-
- if not return_dict:
- return (last_hidden_state, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPoolingAndNoAttention(
- last_hidden_state=last_hidden_state,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- ConvNext Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
- ImageNet.
- """,
- CONVNEXT_START_DOCSTRING,
-)
-class ConvNextForImageClassification(ConvNextPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.convnext = ConvNextModel(config)
-
- # Classifier head
- self.classifier = (
- nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVNEXT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=ImageClassifierOutputWithNoAttention,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: torch.FloatTensor = None,
- labels: Optional[torch.LongTensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.convnext(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
-
- pooled_output = outputs.pooler_output if return_dict else outputs[1]
-
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutputWithNoAttention(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- ConvNeXt backbone, to be used with frameworks like DETR and MaskFormer.
- """,
- CONVNEXT_START_DOCSTRING,
-)
-class ConvNextBackbone(ConvNextPreTrainedModel, BackboneMixin):
- def __init__(self, config):
- super().__init__(config)
- super()._init_backbone(config)
-
- self.embeddings = ConvNextEmbeddings(config)
- self.encoder = ConvNextEncoder(config)
- self.num_features = [config.hidden_sizes[0]] + config.hidden_sizes
-
- # Add layer norms to hidden states of out_features
- hidden_states_norms = {}
- for stage, num_channels in zip(self._out_features, self.channels):
- hidden_states_norms[stage] = ConvNextLayerNorm(num_channels, data_format="channels_first")
- self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
-
- # initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVNEXT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.Tensor,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> BackboneOutput:
- """
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, AutoBackbone
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> processor = AutoImageProcessor.from_pretrained("facebook/convnext-tiny-224")
- >>> model = AutoBackbone.from_pretrained("facebook/convnext-tiny-224")
-
- >>> inputs = processor(image, return_tensors="pt")
- >>> outputs = model(**inputs)
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- embedding_output = self.embeddings(pixel_values)
-
- outputs = self.encoder(
- embedding_output,
- output_hidden_states=True,
- return_dict=return_dict,
- )
-
- hidden_states = outputs.hidden_states if return_dict else outputs[1]
-
- feature_maps = ()
- for stage, hidden_state in zip(self.stage_names, hidden_states):
- if stage in self.out_features:
- hidden_state = self.hidden_states_norms[stage](hidden_state)
- feature_maps += (hidden_state,)
-
- if not return_dict:
- output = (feature_maps,)
- if output_hidden_states:
- output += (hidden_states,)
- return output
-
- return BackboneOutput(
- feature_maps=feature_maps,
- hidden_states=hidden_states if output_hidden_states else None,
- attentions=None,
- )
diff --git a/transformers/models/convnext/modeling_tf_convnext.py b/transformers/models/convnext/modeling_tf_convnext.py
deleted file mode 100644
index b92ac446d94f21c988f431150b663b919e52e975..0000000000000000000000000000000000000000
--- a/transformers/models/convnext/modeling_tf_convnext.py
+++ /dev/null
@@ -1,667 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 ConvNext model."""
-
-
-from __future__ import annotations
-
-from typing import List, Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import TFBaseModelOutput, TFBaseModelOutputWithPooling, TFSequenceClassifierOutput
-from ...modeling_tf_utils import (
- TFModelInputType,
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import shape_list
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from .configuration_convnext import ConvNextConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-_CONFIG_FOR_DOC = "ConvNextConfig"
-_CHECKPOINT_FOR_DOC = "facebook/convnext-tiny-224"
-
-
-class TFConvNextDropPath(keras.layers.Layer):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
- References:
- (1) github.com:rwightman/pytorch-image-models
- """
-
- def __init__(self, drop_path: float, **kwargs):
- super().__init__(**kwargs)
- self.drop_path = drop_path
-
- def call(self, x: tf.Tensor, training=None):
- if training:
- keep_prob = 1 - self.drop_path
- shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)
- random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)
- random_tensor = tf.floor(random_tensor)
- return (x / keep_prob) * random_tensor
- return x
-
-
-class TFConvNextEmbeddings(keras.layers.Layer):
- """This class is comparable to (and inspired by) the SwinEmbeddings class
- found in src/transformers/models/swin/modeling_swin.py.
- """
-
- def __init__(self, config: ConvNextConfig, **kwargs):
- super().__init__(**kwargs)
- self.patch_embeddings = keras.layers.Conv2D(
- filters=config.hidden_sizes[0],
- kernel_size=config.patch_size,
- strides=config.patch_size,
- name="patch_embeddings",
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer=keras.initializers.Zeros(),
- )
- self.layernorm = keras.layers.LayerNormalization(epsilon=1e-6, name="layernorm")
- self.num_channels = config.num_channels
- self.config = config
-
- def call(self, pixel_values):
- if isinstance(pixel_values, dict):
- pixel_values = pixel_values["pixel_values"]
-
- tf.debugging.assert_equal(
- shape_list(pixel_values)[1],
- self.num_channels,
- message="Make sure that the channel dimension of the pixel values match with the one set in the configuration.",
- )
-
- # When running on CPU, `keras.layers.Conv2D` doesn't support `NCHW` format.
- # So change the input format from `NCHW` to `NHWC`.
- # shape = (batch_size, in_height, in_width, in_channels)
- pixel_values = tf.transpose(pixel_values, perm=(0, 2, 3, 1))
-
- embeddings = self.patch_embeddings(pixel_values)
- embeddings = self.layernorm(embeddings)
- return embeddings
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "patch_embeddings", None) is not None:
- with tf.name_scope(self.patch_embeddings.name):
- self.patch_embeddings.build([None, None, None, self.config.num_channels])
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, None, None, self.config.hidden_sizes[0]])
-
-
-class TFConvNextLayer(keras.layers.Layer):
- """This corresponds to the `Block` class in the original implementation.
-
- There are two equivalent implementations: [DwConv, LayerNorm (channels_first), Conv, GELU,1x1 Conv]; all in (N, C,
- H, W) (2) [DwConv, Permute to (N, H, W, C), LayerNorm (channels_last), Linear, GELU, Linear]; Permute back
-
- The authors used (2) as they find it slightly faster in PyTorch. Since we already permuted the inputs to follow
- NHWC ordering, we can just apply the operations straight-away without the permutation.
-
- Args:
- config ([`ConvNextConfig`]): Model configuration class.
- dim (`int`): Number of input channels.
- drop_path (`float`): Stochastic depth rate. Default: 0.0.
- """
-
- def __init__(self, config, dim, drop_path=0.0, **kwargs):
- super().__init__(**kwargs)
- self.dim = dim
- self.config = config
- self.dwconv = keras.layers.Conv2D(
- filters=dim,
- kernel_size=7,
- padding="same",
- groups=dim,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer="zeros",
- name="dwconv",
- ) # depthwise conv
- self.layernorm = keras.layers.LayerNormalization(
- epsilon=1e-6,
- name="layernorm",
- )
- self.pwconv1 = keras.layers.Dense(
- units=4 * dim,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer="zeros",
- name="pwconv1",
- ) # pointwise/1x1 convs, implemented with linear layers
- self.act = get_tf_activation(config.hidden_act)
- self.pwconv2 = keras.layers.Dense(
- units=dim,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer="zeros",
- name="pwconv2",
- )
- # Using `layers.Activation` instead of `tf.identity` to better control `training`
- # behaviour.
- self.drop_path = (
- TFConvNextDropPath(drop_path, name="drop_path")
- if drop_path > 0.0
- else keras.layers.Activation("linear", name="drop_path")
- )
-
- def build(self, input_shape: tf.TensorShape = None):
- # PT's `nn.Parameters` must be mapped to a TF layer weight to inherit the same name hierarchy (and vice-versa)
- self.layer_scale_parameter = (
- self.add_weight(
- shape=(self.dim,),
- initializer=keras.initializers.Constant(value=self.config.layer_scale_init_value),
- trainable=True,
- name="layer_scale_parameter",
- )
- if self.config.layer_scale_init_value > 0
- else None
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "dwconv", None) is not None:
- with tf.name_scope(self.dwconv.name):
- self.dwconv.build([None, None, None, self.dim])
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, None, None, self.dim])
- if getattr(self, "pwconv1", None) is not None:
- with tf.name_scope(self.pwconv1.name):
- self.pwconv1.build([None, None, self.dim])
- if getattr(self, "pwconv2", None) is not None:
- with tf.name_scope(self.pwconv2.name):
- self.pwconv2.build([None, None, 4 * self.dim])
- if getattr(self, "drop_path", None) is not None:
- with tf.name_scope(self.drop_path.name):
- self.drop_path.build(None)
-
- def call(self, hidden_states, training=False):
- input = hidden_states
- x = self.dwconv(hidden_states)
- x = self.layernorm(x)
- x = self.pwconv1(x)
- x = self.act(x)
- x = self.pwconv2(x)
-
- if self.layer_scale_parameter is not None:
- x = self.layer_scale_parameter * x
-
- x = input + self.drop_path(x, training=training)
- return x
-
-
-class TFConvNextStage(keras.layers.Layer):
- """ConvNext stage, consisting of an optional downsampling layer + multiple residual blocks.
-
- Args:
- config (`ConvNextV2Config`):
- Model configuration class.
- in_channels (`int`):
- Number of input channels.
- out_channels (`int`):
- Number of output channels.
- depth (`int`):
- Number of residual blocks.
- drop_path_rates(`List[float]`):
- Stochastic depth rates for each layer.
- """
-
- def __init__(
- self,
- config: ConvNextConfig,
- in_channels: int,
- out_channels: int,
- kernel_size: int = 2,
- stride: int = 2,
- depth: int = 2,
- drop_path_rates: Optional[List[float]] = None,
- **kwargs,
- ):
- super().__init__(**kwargs)
- if in_channels != out_channels or stride > 1:
- self.downsampling_layer = [
- keras.layers.LayerNormalization(
- epsilon=1e-6,
- name="downsampling_layer.0",
- ),
- # Inputs to this layer will follow NHWC format since we
- # transposed the inputs from NCHW to NHWC in the `TFConvNextEmbeddings`
- # layer. All the outputs throughout the model will be in NHWC
- # from this point on until the output where we again change to
- # NCHW.
- keras.layers.Conv2D(
- filters=out_channels,
- kernel_size=kernel_size,
- strides=stride,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer=keras.initializers.Zeros(),
- name="downsampling_layer.1",
- ),
- ]
- else:
- self.downsampling_layer = [tf.identity]
-
- drop_path_rates = drop_path_rates or [0.0] * depth
- self.layers = [
- TFConvNextLayer(
- config,
- dim=out_channels,
- drop_path=drop_path_rates[j],
- name=f"layers.{j}",
- )
- for j in range(depth)
- ]
- self.in_channels = in_channels
- self.out_channels = out_channels
- self.stride = stride
-
- def call(self, hidden_states):
- for layer in self.downsampling_layer:
- hidden_states = layer(hidden_states)
- for layer in self.layers:
- hidden_states = layer(hidden_states)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "layers", None) is not None:
- for layer in self.layers:
- with tf.name_scope(layer.name):
- layer.build(None)
- if self.in_channels != self.out_channels or self.stride > 1:
- with tf.name_scope(self.downsampling_layer[0].name):
- self.downsampling_layer[0].build([None, None, None, self.in_channels])
- with tf.name_scope(self.downsampling_layer[1].name):
- self.downsampling_layer[1].build([None, None, None, self.in_channels])
-
-
-class TFConvNextEncoder(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- self.stages = []
- drop_path_rates = tf.linspace(0.0, config.drop_path_rate, sum(config.depths))
- drop_path_rates = tf.split(drop_path_rates, config.depths)
- drop_path_rates = [x.numpy().tolist() for x in drop_path_rates]
- prev_chs = config.hidden_sizes[0]
- for i in range(config.num_stages):
- out_chs = config.hidden_sizes[i]
- stage = TFConvNextStage(
- config,
- in_channels=prev_chs,
- out_channels=out_chs,
- stride=2 if i > 0 else 1,
- depth=config.depths[i],
- drop_path_rates=drop_path_rates[i],
- name=f"stages.{i}",
- )
- self.stages.append(stage)
- prev_chs = out_chs
-
- def call(self, hidden_states, output_hidden_states=False, return_dict=True):
- all_hidden_states = () if output_hidden_states else None
-
- for i, layer_module in enumerate(self.stages):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- hidden_states = layer_module(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
-
- return TFBaseModelOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
-
- def build(self, input_shape=None):
- for stage in self.stages:
- with tf.name_scope(stage.name):
- stage.build(None)
-
-
-@keras_serializable
-class TFConvNextMainLayer(keras.layers.Layer):
- config_class = ConvNextConfig
-
- def __init__(self, config: ConvNextConfig, add_pooling_layer: bool = True, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embeddings = TFConvNextEmbeddings(config, name="embeddings")
- self.encoder = TFConvNextEncoder(config, name="encoder")
- self.layernorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
- # We are setting the `data_format` like so because from here on we will revert to the
- # NCHW output format
- self.pooler = keras.layers.GlobalAvgPool2D(data_format="channels_first") if add_pooling_layer else None
-
- @unpack_inputs
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- embedding_output = self.embeddings(pixel_values, training=training)
-
- encoder_outputs = self.encoder(
- embedding_output,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- last_hidden_state = encoder_outputs[0]
- # Change to NCHW output format have uniformity in the modules
- last_hidden_state = tf.transpose(last_hidden_state, perm=(0, 3, 1, 2))
- pooled_output = self.layernorm(self.pooler(last_hidden_state))
-
- # Change the other hidden state outputs to NCHW as well
- if output_hidden_states:
- hidden_states = tuple([tf.transpose(h, perm=(0, 3, 1, 2)) for h in encoder_outputs[1]])
-
- if not return_dict:
- hidden_states = hidden_states if output_hidden_states else ()
- return (last_hidden_state, pooled_output) + hidden_states
-
- return TFBaseModelOutputWithPooling(
- last_hidden_state=last_hidden_state,
- pooler_output=pooled_output,
- hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, self.config.hidden_sizes[-1]])
-
-
-class TFConvNextPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ConvNextConfig
- base_model_prefix = "convnext"
- main_input_name = "pixel_values"
-
-
-CONVNEXT_START_DOCSTRING = r"""
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `pixel_values` only and nothing else: `model(pixel_values)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([pixel_values, attention_mask])` or `model([pixel_values, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"pixel_values": pixel_values, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`ConvNextConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CONVNEXT_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`np.ndarray`, `tf.Tensor`, `List[tf.Tensor]` ``Dict[str, tf.Tensor]` or `Dict[str, np.ndarray]` and each example must have the shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`ConvNextImageProcessor.__call__`] for details.
-
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
-"""
-
-
-@add_start_docstrings(
- "The bare ConvNext model outputting raw features without any specific head on top.",
- CONVNEXT_START_DOCSTRING,
-)
-class TFConvNextModel(TFConvNextPreTrainedModel):
- def __init__(self, config, *inputs, add_pooling_layer=True, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.convnext = TFConvNextMainLayer(config, add_pooling_layer=add_pooling_layer, name="convnext")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVNEXT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFBaseModelOutputWithPooling, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, TFConvNextModel
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/convnext-tiny-224")
- >>> model = TFConvNextModel.from_pretrained("facebook/convnext-tiny-224")
-
- >>> inputs = image_processor(images=image, return_tensors="tf")
- >>> outputs = model(**inputs)
- >>> last_hidden_states = outputs.last_hidden_state
- ```"""
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- outputs = self.convnext(
- pixel_values=pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- if not return_dict:
- return (outputs[0],) + outputs[1:]
-
- return TFBaseModelOutputWithPooling(
- last_hidden_state=outputs.last_hidden_state,
- pooler_output=outputs.pooler_output,
- hidden_states=outputs.hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convnext", None) is not None:
- with tf.name_scope(self.convnext.name):
- self.convnext.build(None)
-
-
-@add_start_docstrings(
- """
- ConvNext Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
- ImageNet.
- """,
- CONVNEXT_START_DOCSTRING,
-)
-class TFConvNextForImageClassification(TFConvNextPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: ConvNextConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
- self.convnext = TFConvNextMainLayer(config, name="convnext")
-
- # Classifier head
- self.classifier = keras.layers.Dense(
- units=config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer="zeros",
- name="classifier",
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVNEXT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFSequenceClassifierOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, TFConvNextForImageClassification
- >>> import tensorflow as tf
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/convnext-tiny-224")
- >>> model = TFConvNextForImageClassification.from_pretrained("facebook/convnext-tiny-224")
-
- >>> inputs = image_processor(images=image, return_tensors="tf")
- >>> outputs = model(**inputs)
- >>> logits = outputs.logits
- >>> # model predicts one of the 1000 ImageNet classes
- >>> predicted_class_idx = tf.math.argmax(logits, axis=-1)[0]
- >>> print("Predicted class:", model.config.id2label[int(predicted_class_idx)])
- ```"""
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- outputs = self.convnext(
- pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- pooled_output = outputs.pooler_output if return_dict else outputs[1]
-
- logits = self.classifier(pooled_output)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convnext", None) is not None:
- with tf.name_scope(self.convnext.name):
- self.convnext.build(None)
- if getattr(self, "classifier", None) is not None:
- if hasattr(self.classifier, "name"):
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_sizes[-1]])
diff --git a/transformers/models/convnextv2/__init__.py b/transformers/models/convnextv2/__init__.py
deleted file mode 100644
index d2a484b9b8285083ba958772743207d64a8403bc..0000000000000000000000000000000000000000
--- a/transformers/models/convnextv2/__init__.py
+++ /dev/null
@@ -1,97 +0,0 @@
-# flake8: noqa
-# There's no way to ignore "F401 '...' imported but unused" warnings in this
-# module, but to preserve other warnings. So, don't check this module at all.
-
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-# rely on isort to merge the imports
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_torch_available,
- is_tf_available,
-)
-
-
-_import_structure = {
- "configuration_convnextv2": [
- "CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "ConvNextV2Config",
- ]
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_convnextv2"] = [
- "CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST",
- "ConvNextV2ForImageClassification",
- "ConvNextV2Model",
- "ConvNextV2PreTrainedModel",
- "ConvNextV2Backbone",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_convnextv2"] = [
- "TFConvNextV2ForImageClassification",
- "TFConvNextV2Model",
- "TFConvNextV2PreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_convnextv2 import (
- CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP,
- ConvNextV2Config,
- )
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_convnextv2 import (
- CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST,
- ConvNextV2Backbone,
- ConvNextV2ForImageClassification,
- ConvNextV2Model,
- ConvNextV2PreTrainedModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_convnextv2 import (
- TFConvNextV2ForImageClassification,
- TFConvNextV2Model,
- TFConvNextV2PreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
diff --git a/transformers/models/convnextv2/__pycache__/__init__.cpython-310.pyc b/transformers/models/convnextv2/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 9ffe0f697dfca53282daaa7e8506b7347bd5c163..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnextv2/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnextv2/__pycache__/configuration_convnextv2.cpython-310.pyc b/transformers/models/convnextv2/__pycache__/configuration_convnextv2.cpython-310.pyc
deleted file mode 100644
index 5af2fd8e32d8994aa9b0682dcc8562468e7381ea..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnextv2/__pycache__/configuration_convnextv2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnextv2/__pycache__/convert_convnextv2_to_pytorch.cpython-310.pyc b/transformers/models/convnextv2/__pycache__/convert_convnextv2_to_pytorch.cpython-310.pyc
deleted file mode 100644
index e96141b566b98926fab233e5f91809fae5c83953..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnextv2/__pycache__/convert_convnextv2_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnextv2/__pycache__/modeling_convnextv2.cpython-310.pyc b/transformers/models/convnextv2/__pycache__/modeling_convnextv2.cpython-310.pyc
deleted file mode 100644
index e1fe58b1d1bff09a8bc5c611161da56c6267bc31..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnextv2/__pycache__/modeling_convnextv2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnextv2/__pycache__/modeling_tf_convnextv2.cpython-310.pyc b/transformers/models/convnextv2/__pycache__/modeling_tf_convnextv2.cpython-310.pyc
deleted file mode 100644
index 0ad613d5b1be5ec0f6308429b23c8efbd376f7e0..0000000000000000000000000000000000000000
Binary files a/transformers/models/convnextv2/__pycache__/modeling_tf_convnextv2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/convnextv2/configuration_convnextv2.py b/transformers/models/convnextv2/configuration_convnextv2.py
deleted file mode 100644
index ccee03eef6a492332c09e523994d8bcedbcfca5f..0000000000000000000000000000000000000000
--- a/transformers/models/convnextv2/configuration_convnextv2.py
+++ /dev/null
@@ -1,117 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" ConvNeXTV2 model configuration"""
-
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class ConvNextV2Config(BackboneConfigMixin, PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`ConvNextV2Model`]. It is used to instantiate an
- ConvNeXTV2 model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the ConvNeXTV2
- [facebook/convnextv2-tiny-1k-224](https://huggingface.co/facebook/convnextv2-tiny-1k-224) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- patch_size (`int`, optional, defaults to 4):
- Patch size to use in the patch embedding layer.
- num_stages (`int`, optional, defaults to 4):
- The number of stages in the model.
- hidden_sizes (`List[int]`, *optional*, defaults to `[96, 192, 384, 768]`):
- Dimensionality (hidden size) at each stage.
- depths (`List[int]`, *optional*, defaults to `[3, 3, 9, 3]`):
- Depth (number of blocks) for each stage.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in each block. If string, `"gelu"`, `"relu"`,
- `"selu"` and `"gelu_new"` are supported.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- drop_path_rate (`float`, *optional*, defaults to 0.0):
- The drop rate for stochastic depth.
- out_features (`List[str]`, *optional*):
- If used as backbone, list of features to output. Can be any of `"stem"`, `"stage1"`, `"stage2"`, etc.
- (depending on how many stages the model has). If unset and `out_indices` is set, will default to the
- corresponding stages. If unset and `out_indices` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
- out_indices (`List[int]`, *optional*):
- If used as backbone, list of indices of features to output. Can be any of 0, 1, 2, etc. (depending on how
- many stages the model has). If unset and `out_features` is set, will default to the corresponding stages.
- If unset and `out_features` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
-
- Example:
- ```python
- >>> from transformers import ConvNeXTV2Config, ConvNextV2Model
-
- >>> # Initializing a ConvNeXTV2 convnextv2-tiny-1k-224 style configuration
- >>> configuration = ConvNeXTV2Config()
-
- >>> # Initializing a model (with random weights) from the convnextv2-tiny-1k-224 style configuration
- >>> model = ConvNextV2Model(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "convnextv2"
-
- def __init__(
- self,
- num_channels=3,
- patch_size=4,
- num_stages=4,
- hidden_sizes=None,
- depths=None,
- hidden_act="gelu",
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- drop_path_rate=0.0,
- image_size=224,
- out_features=None,
- out_indices=None,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.num_channels = num_channels
- self.patch_size = patch_size
- self.num_stages = num_stages
- self.hidden_sizes = [96, 192, 384, 768] if hidden_sizes is None else hidden_sizes
- self.depths = [3, 3, 9, 3] if depths is None else depths
- self.hidden_act = hidden_act
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.drop_path_rate = drop_path_rate
- self.image_size = image_size
- self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
- self._out_features, self._out_indices = get_aligned_output_features_output_indices(
- out_features=out_features, out_indices=out_indices, stage_names=self.stage_names
- )
diff --git a/transformers/models/convnextv2/convert_convnextv2_to_pytorch.py b/transformers/models/convnextv2/convert_convnextv2_to_pytorch.py
deleted file mode 100644
index 8094ecf0d6157a1bb2343817f7e9303f622d9102..0000000000000000000000000000000000000000
--- a/transformers/models/convnextv2/convert_convnextv2_to_pytorch.py
+++ /dev/null
@@ -1,286 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert ConvNeXTV2 checkpoints from the original repository.
-
-URL: https://github.com/facebookresearch/ConvNeXt"""
-
-import argparse
-import json
-import os
-
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from transformers import ConvNextImageProcessor, ConvNextV2Config, ConvNextV2ForImageClassification
-from transformers.image_utils import PILImageResampling
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_convnextv2_config(checkpoint_url):
- config = ConvNextV2Config()
-
- if "atto" in checkpoint_url:
- depths = [2, 2, 6, 2]
- hidden_sizes = [40, 80, 160, 320]
- if "femto" in checkpoint_url:
- depths = [2, 2, 6, 2]
- hidden_sizes = [48, 96, 192, 384]
- if "pico" in checkpoint_url:
- depths = [2, 2, 6, 2]
- hidden_sizes = [64, 128, 256, 512]
- if "nano" in checkpoint_url:
- depths = [2, 2, 8, 2]
- hidden_sizes = [80, 160, 320, 640]
- if "tiny" in checkpoint_url:
- depths = [3, 3, 9, 3]
- hidden_sizes = [96, 192, 384, 768]
- if "base" in checkpoint_url:
- depths = [3, 3, 27, 3]
- hidden_sizes = [128, 256, 512, 1024]
- if "large" in checkpoint_url:
- depths = [3, 3, 27, 3]
- hidden_sizes = [192, 384, 768, 1536]
- if "huge" in checkpoint_url:
- depths = [3, 3, 27, 3]
- hidden_sizes = [352, 704, 1408, 2816]
-
- num_labels = 1000
- filename = "imagenet-1k-id2label.json"
- expected_shape = (1, 1000)
-
- repo_id = "huggingface/label-files"
- config.num_labels = num_labels
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
-
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
- config.hidden_sizes = hidden_sizes
- config.depths = depths
-
- return config, expected_shape
-
-
-def rename_key(name):
- if "downsample_layers.0.0" in name:
- name = name.replace("downsample_layers.0.0", "embeddings.patch_embeddings")
- if "downsample_layers.0.1" in name:
- name = name.replace("downsample_layers.0.1", "embeddings.norm") # we rename to layernorm later on
- if "downsample_layers.1.0" in name:
- name = name.replace("downsample_layers.1.0", "stages.1.downsampling_layer.0")
- if "downsample_layers.1.1" in name:
- name = name.replace("downsample_layers.1.1", "stages.1.downsampling_layer.1")
- if "downsample_layers.2.0" in name:
- name = name.replace("downsample_layers.2.0", "stages.2.downsampling_layer.0")
- if "downsample_layers.2.1" in name:
- name = name.replace("downsample_layers.2.1", "stages.2.downsampling_layer.1")
- if "downsample_layers.3.0" in name:
- name = name.replace("downsample_layers.3.0", "stages.3.downsampling_layer.0")
- if "downsample_layers.3.1" in name:
- name = name.replace("downsample_layers.3.1", "stages.3.downsampling_layer.1")
- if "stages" in name and "downsampling_layer" not in name:
- # stages.0.0. for instance should be renamed to stages.0.layers.0.
- name = name[: len("stages.0")] + ".layers" + name[len("stages.0") :]
- if "gamma" in name:
- name = name.replace("gamma", "weight")
- if "beta" in name:
- name = name.replace("beta", "bias")
- if "stages" in name:
- name = name.replace("stages", "encoder.stages")
- if "norm" in name:
- name = name.replace("norm", "layernorm")
- if "head" in name:
- name = name.replace("head", "classifier")
-
- return name
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-def convert_preprocessor(checkpoint_url):
- if "224" in checkpoint_url:
- size = 224
- crop_pct = 224 / 256
- elif "384" in checkpoint_url:
- size = 384
- crop_pct = None
- else:
- size = 512
- crop_pct = None
-
- return ConvNextImageProcessor(
- size=size,
- crop_pct=crop_pct,
- image_mean=[0.485, 0.456, 0.406],
- image_std=[0.229, 0.224, 0.225],
- resample=PILImageResampling.BICUBIC,
- )
-
-
-@torch.no_grad()
-def convert_convnextv2_checkpoint(checkpoint_url, pytorch_dump_folder_path, save_model, push_to_hub):
- """
- Copy/paste/tweak model's weights to our ConvNeXTV2 structure.
- """
- print("Downloading original model from checkpoint...")
- # define ConvNeXTV2 configuration based on URL
- config, expected_shape = get_convnextv2_config(checkpoint_url)
- # load original state_dict from URL
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url)["model"]
-
- print("Converting model parameters...")
- # rename keys
- for key in state_dict.copy().keys():
- val = state_dict.pop(key)
- state_dict[rename_key(key)] = val
- # add prefix to all keys expect classifier head
- for key in state_dict.copy().keys():
- val = state_dict.pop(key)
- if not key.startswith("classifier"):
- key = "convnextv2." + key
- state_dict[key] = val
-
- # load HuggingFace model
- model = ConvNextV2ForImageClassification(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- # Check outputs on an image, prepared by ConvNextImageProcessor
- preprocessor = convert_preprocessor(checkpoint_url)
- inputs = preprocessor(images=prepare_img(), return_tensors="pt")
- logits = model(**inputs).logits
-
- # note: the logits below were obtained without center cropping
- if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
- expected_logits = torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
- expected_logits = torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
- expected_logits = torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
- expected_logits = torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
- expected_logits = torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
- expected_logits = torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
- expected_logits = torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
- expected_logits = torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
- expected_logits = torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
- expected_logits = torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
- expected_logits = torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
- expected_logits = torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
- expected_logits = torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
- expected_logits = torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
- expected_logits = torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
- expected_logits = torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
- expected_logits = torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309])
- elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
- expected_logits = torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826])
- else:
- raise ValueError(f"Unknown URL: {checkpoint_url}")
-
- assert torch.allclose(logits[0, :5], expected_logits, atol=1e-3)
- assert logits.shape == expected_shape
- print("Model outputs match the original results!")
-
- if save_model:
- print("Saving model to local...")
- # Create folder to save model
- if not os.path.isdir(pytorch_dump_folder_path):
- os.mkdir(pytorch_dump_folder_path)
-
- model.save_pretrained(pytorch_dump_folder_path)
- preprocessor.save_pretrained(pytorch_dump_folder_path)
-
- model_name = "convnextv2"
- if "atto" in checkpoint_url:
- model_name += "-atto"
- if "femto" in checkpoint_url:
- model_name += "-femto"
- if "pico" in checkpoint_url:
- model_name += "-pico"
- if "nano" in checkpoint_url:
- model_name += "-nano"
- elif "tiny" in checkpoint_url:
- model_name += "-tiny"
- elif "base" in checkpoint_url:
- model_name += "-base"
- elif "large" in checkpoint_url:
- model_name += "-large"
- elif "huge" in checkpoint_url:
- model_name += "-huge"
- if "22k" in checkpoint_url and "1k" not in checkpoint_url:
- model_name += "-22k"
- elif "22k" in checkpoint_url and "1k" in checkpoint_url:
- model_name += "-22k-1k"
- elif "1k" in checkpoint_url:
- model_name += "-1k"
- if "224" in checkpoint_url:
- model_name += "-224"
- elif "384" in checkpoint_url:
- model_name += "-384"
- elif "512" in checkpoint_url:
- model_name += "-512"
-
- if push_to_hub:
- print(f"Pushing {model_name} to the hub...")
- model.push_to_hub(model_name)
- preprocessor.push_to_hub(model_name)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--checkpoint_url",
- default="https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt",
- type=str,
- help="URL of the original ConvNeXTV2 checkpoint you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default="model",
- type=str,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument("--save_model", action="store_true", help="Save model to local")
- parser.add_argument("--push_to_hub", action="store_true", help="Push model and image preprocessor to the hub")
-
- args = parser.parse_args()
- convert_convnextv2_checkpoint(
- args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
- )
diff --git a/transformers/models/convnextv2/modeling_convnextv2.py b/transformers/models/convnextv2/modeling_convnextv2.py
deleted file mode 100644
index 7439f212971ec11c21215cbecb665f21405145f0..0000000000000000000000000000000000000000
--- a/transformers/models/convnextv2/modeling_convnextv2.py
+++ /dev/null
@@ -1,574 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch ConvNextV2 model."""
-
-
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BackboneOutput,
- BaseModelOutputWithNoAttention,
- BaseModelOutputWithPoolingAndNoAttention,
- ImageClassifierOutputWithNoAttention,
-)
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from ...utils.backbone_utils import BackboneMixin
-from .configuration_convnextv2 import ConvNextV2Config
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "ConvNextV2Config"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/convnextv2-tiny-1k-224"
-_EXPECTED_OUTPUT_SHAPE = [1, 768, 7, 7]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "facebook/convnextv2-tiny-1k-224"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->ConvNextV2
-class ConvNextV2DropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-class ConvNextV2GRN(nn.Module):
- """GRN (Global Response Normalization) layer"""
-
- def __init__(self, dim: int):
- super().__init__()
- self.weight = nn.Parameter(torch.zeros(1, 1, 1, dim))
- self.bias = nn.Parameter(torch.zeros(1, 1, 1, dim))
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
- # Compute and normalize global spatial feature maps
- global_features = torch.norm(hidden_states, p=2, dim=(1, 2), keepdim=True)
- norm_features = global_features / (global_features.mean(dim=-1, keepdim=True) + 1e-6)
- hidden_states = self.weight * (hidden_states * norm_features) + self.bias + hidden_states
-
- return hidden_states
-
-
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextLayerNorm with ConvNext->ConvNextV2
-class ConvNextV2LayerNorm(nn.Module):
- r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
- The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
- width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).
- """
-
- def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
- super().__init__()
- self.weight = nn.Parameter(torch.ones(normalized_shape))
- self.bias = nn.Parameter(torch.zeros(normalized_shape))
- self.eps = eps
- self.data_format = data_format
- if self.data_format not in ["channels_last", "channels_first"]:
- raise NotImplementedError(f"Unsupported data format: {self.data_format}")
- self.normalized_shape = (normalized_shape,)
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- if self.data_format == "channels_last":
- x = torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
- elif self.data_format == "channels_first":
- input_dtype = x.dtype
- x = x.float()
- u = x.mean(1, keepdim=True)
- s = (x - u).pow(2).mean(1, keepdim=True)
- x = (x - u) / torch.sqrt(s + self.eps)
- x = x.to(dtype=input_dtype)
- x = self.weight[:, None, None] * x + self.bias[:, None, None]
- return x
-
-
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextEmbeddings with ConvNext->ConvNextV2
-class ConvNextV2Embeddings(nn.Module):
- """This class is comparable to (and inspired by) the SwinEmbeddings class
- found in src/transformers/models/swin/modeling_swin.py.
- """
-
- def __init__(self, config):
- super().__init__()
- self.patch_embeddings = nn.Conv2d(
- config.num_channels, config.hidden_sizes[0], kernel_size=config.patch_size, stride=config.patch_size
- )
- self.layernorm = ConvNextV2LayerNorm(config.hidden_sizes[0], eps=1e-6, data_format="channels_first")
- self.num_channels = config.num_channels
-
- def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
- num_channels = pixel_values.shape[1]
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- embeddings = self.patch_embeddings(pixel_values)
- embeddings = self.layernorm(embeddings)
- return embeddings
-
-
-class ConvNextV2Layer(nn.Module):
- """This corresponds to the `Block` class in the original implementation.
-
- There are two equivalent implementations: [DwConv, LayerNorm (channels_first), Conv, GELU,1x1 Conv]; all in (N, C,
- H, W) (2) [DwConv, Permute to (N, H, W, C), LayerNorm (channels_last), Linear, GELU, Linear]; Permute back
-
- The authors used (2) as they find it slightly faster in PyTorch.
-
- Args:
- config ([`ConvNextV2Config`]): Model configuration class.
- dim (`int`): Number of input channels.
- drop_path (`float`): Stochastic depth rate. Default: 0.0.
- """
-
- def __init__(self, config, dim, drop_path=0):
- super().__init__()
- # depthwise conv
- self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
- self.layernorm = ConvNextV2LayerNorm(dim, eps=1e-6)
- # pointwise/1x1 convs, implemented with linear layers
- self.pwconv1 = nn.Linear(dim, 4 * dim)
- self.act = ACT2FN[config.hidden_act]
- self.grn = ConvNextV2GRN(4 * dim)
- self.pwconv2 = nn.Linear(4 * dim, dim)
- self.drop_path = ConvNextV2DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
- input = hidden_states
- x = self.dwconv(hidden_states)
- # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
- x = x.permute(0, 2, 3, 1)
- x = self.layernorm(x)
- x = self.pwconv1(x)
- x = self.act(x)
- x = self.grn(x)
- x = self.pwconv2(x)
- # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
- x = x.permute(0, 3, 1, 2)
-
- x = input + self.drop_path(x)
- return x
-
-
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextStage with ConvNeXT->ConvNeXTV2, ConvNext->ConvNextV2
-class ConvNextV2Stage(nn.Module):
- """ConvNeXTV2 stage, consisting of an optional downsampling layer + multiple residual blocks.
-
- Args:
- config ([`ConvNextV2Config`]): Model configuration class.
- in_channels (`int`): Number of input channels.
- out_channels (`int`): Number of output channels.
- depth (`int`): Number of residual blocks.
- drop_path_rates(`List[float]`): Stochastic depth rates for each layer.
- """
-
- def __init__(self, config, in_channels, out_channels, kernel_size=2, stride=2, depth=2, drop_path_rates=None):
- super().__init__()
-
- if in_channels != out_channels or stride > 1:
- self.downsampling_layer = nn.Sequential(
- ConvNextV2LayerNorm(in_channels, eps=1e-6, data_format="channels_first"),
- nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride),
- )
- else:
- self.downsampling_layer = nn.Identity()
- drop_path_rates = drop_path_rates or [0.0] * depth
- self.layers = nn.Sequential(
- *[ConvNextV2Layer(config, dim=out_channels, drop_path=drop_path_rates[j]) for j in range(depth)]
- )
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
- hidden_states = self.downsampling_layer(hidden_states)
- hidden_states = self.layers(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextEncoder with ConvNext->ConvNextV2
-class ConvNextV2Encoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.stages = nn.ModuleList()
- drop_path_rates = [
- x.tolist() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths)).split(config.depths)
- ]
- prev_chs = config.hidden_sizes[0]
- for i in range(config.num_stages):
- out_chs = config.hidden_sizes[i]
- stage = ConvNextV2Stage(
- config,
- in_channels=prev_chs,
- out_channels=out_chs,
- stride=2 if i > 0 else 1,
- depth=config.depths[i],
- drop_path_rates=drop_path_rates[i],
- )
- self.stages.append(stage)
- prev_chs = out_chs
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple, BaseModelOutputWithNoAttention]:
- all_hidden_states = () if output_hidden_states else None
-
- for i, layer_module in enumerate(self.stages):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- hidden_states = layer_module(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
-
- return BaseModelOutputWithNoAttention(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- )
-
-
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextPreTrainedModel with ConvNext->ConvNextV2, convnext->convnextv2
-class ConvNextV2PreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ConvNextV2Config
- base_model_prefix = "convnextv2"
- main_input_name = "pixel_values"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-CONVNEXTV2_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`ConvNextV2Config`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CONVNEXTV2_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`ConvNextImageProcessor`]. See
- [`ConvNextImageProcessor.__call__`] for details.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare ConvNextV2 model outputting raw features without any specific head on top.",
- CONVNEXTV2_START_DOCSTRING,
-)
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextModel with CONVNEXT->CONVNEXTV2, ConvNext->ConvNextV2
-class ConvNextV2Model(ConvNextV2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.config = config
-
- self.embeddings = ConvNextV2Embeddings(config)
- self.encoder = ConvNextV2Encoder(config)
-
- # final layernorm layer
- self.layernorm = nn.LayerNorm(config.hidden_sizes[-1], eps=config.layer_norm_eps)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVNEXTV2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPoolingAndNoAttention,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: torch.FloatTensor = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPoolingAndNoAttention]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- embedding_output = self.embeddings(pixel_values)
-
- encoder_outputs = self.encoder(
- embedding_output,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden_state = encoder_outputs[0]
-
- # global average pooling, (N, C, H, W) -> (N, C)
- pooled_output = self.layernorm(last_hidden_state.mean([-2, -1]))
-
- if not return_dict:
- return (last_hidden_state, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPoolingAndNoAttention(
- last_hidden_state=last_hidden_state,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- ConvNextV2 Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
- ImageNet.
- """,
- CONVNEXTV2_START_DOCSTRING,
-)
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextForImageClassification with CONVNEXT->CONVNEXTV2,ConvNext->ConvNextV2,convnext->convnextv2
-class ConvNextV2ForImageClassification(ConvNextV2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.convnextv2 = ConvNextV2Model(config)
-
- # Classifier head
- self.classifier = (
- nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVNEXTV2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=ImageClassifierOutputWithNoAttention,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: torch.FloatTensor = None,
- labels: Optional[torch.LongTensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.convnextv2(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
-
- pooled_output = outputs.pooler_output if return_dict else outputs[1]
-
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutputWithNoAttention(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- ConvNeXT V2 backbone, to be used with frameworks like DETR and MaskFormer.
- """,
- CONVNEXTV2_START_DOCSTRING,
-)
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextBackbone with CONVNEXT->CONVNEXTV2,ConvNext->ConvNextV2,facebook/convnext-tiny-224->facebook/convnextv2-tiny-1k-224
-class ConvNextV2Backbone(ConvNextV2PreTrainedModel, BackboneMixin):
- def __init__(self, config):
- super().__init__(config)
- super()._init_backbone(config)
-
- self.embeddings = ConvNextV2Embeddings(config)
- self.encoder = ConvNextV2Encoder(config)
- self.num_features = [config.hidden_sizes[0]] + config.hidden_sizes
-
- # Add layer norms to hidden states of out_features
- hidden_states_norms = {}
- for stage, num_channels in zip(self._out_features, self.channels):
- hidden_states_norms[stage] = ConvNextV2LayerNorm(num_channels, data_format="channels_first")
- self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
-
- # initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CONVNEXTV2_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.Tensor,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> BackboneOutput:
- """
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, AutoBackbone
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> processor = AutoImageProcessor.from_pretrained("facebook/convnextv2-tiny-1k-224")
- >>> model = AutoBackbone.from_pretrained("facebook/convnextv2-tiny-1k-224")
-
- >>> inputs = processor(image, return_tensors="pt")
- >>> outputs = model(**inputs)
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- embedding_output = self.embeddings(pixel_values)
-
- outputs = self.encoder(
- embedding_output,
- output_hidden_states=True,
- return_dict=return_dict,
- )
-
- hidden_states = outputs.hidden_states if return_dict else outputs[1]
-
- feature_maps = ()
- for stage, hidden_state in zip(self.stage_names, hidden_states):
- if stage in self.out_features:
- hidden_state = self.hidden_states_norms[stage](hidden_state)
- feature_maps += (hidden_state,)
-
- if not return_dict:
- output = (feature_maps,)
- if output_hidden_states:
- output += (hidden_states,)
- return output
-
- return BackboneOutput(
- feature_maps=feature_maps,
- hidden_states=hidden_states if output_hidden_states else None,
- attentions=None,
- )
diff --git a/transformers/models/convnextv2/modeling_tf_convnextv2.py b/transformers/models/convnextv2/modeling_tf_convnextv2.py
deleted file mode 100644
index 0debe6fd0c54d6387bcb471370585472dc9d6921..0000000000000000000000000000000000000000
--- a/transformers/models/convnextv2/modeling_tf_convnextv2.py
+++ /dev/null
@@ -1,681 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Meta Platforms Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 ConvNextV2 model."""
-
-
-from __future__ import annotations
-
-from typing import List, Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutputWithNoAttention,
- TFBaseModelOutputWithPooling,
- TFBaseModelOutputWithPoolingAndNoAttention,
- TFImageClassifierOutputWithNoAttention,
-)
-from ...modeling_tf_utils import (
- TFModelInputType,
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import shape_list
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_convnextv2 import ConvNextV2Config
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "ConvNextV2Config"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/convnextv2-tiny-1k-224"
-_EXPECTED_OUTPUT_SHAPE = [1, 768, 7, 7]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "facebook/convnextv2-tiny-1k-224"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-# Copied from transformers.models.convnext.modeling_tf_convnext.TFConvNextDropPath with ConvNext->ConvNextV2
-class TFConvNextV2DropPath(keras.layers.Layer):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
- References:
- (1) github.com:rwightman/pytorch-image-models
- """
-
- def __init__(self, drop_path: float, **kwargs):
- super().__init__(**kwargs)
- self.drop_path = drop_path
-
- def call(self, x: tf.Tensor, training=None):
- if training:
- keep_prob = 1 - self.drop_path
- shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)
- random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)
- random_tensor = tf.floor(random_tensor)
- return (x / keep_prob) * random_tensor
- return x
-
-
-class TFConvNextV2GRN(keras.layers.Layer):
- """GRN (Global Response Normalization) layer"""
-
- def __init__(self, config: ConvNextV2Config, dim: int, **kwargs):
- super().__init__(**kwargs)
- self.dim = dim
-
- def build(self, input_shape: tf.TensorShape = None):
- # PT's `nn.Parameters` must be mapped to a TF layer weight to inherit the same name hierarchy (and vice-versa)
- self.weight = self.add_weight(
- name="weight",
- shape=(1, 1, 1, self.dim),
- initializer=keras.initializers.Zeros(),
- )
- self.bias = self.add_weight(
- name="bias",
- shape=(1, 1, 1, self.dim),
- initializer=keras.initializers.Zeros(),
- )
- return super().build(input_shape)
-
- def call(self, hidden_states: tf.Tensor):
- global_features = tf.norm(hidden_states, ord="euclidean", axis=(1, 2), keepdims=True)
- norm_features = global_features / (tf.reduce_mean(global_features, axis=-1, keepdims=True) + 1e-6)
- hidden_states = self.weight * (hidden_states * norm_features) + self.bias + hidden_states
- return hidden_states
-
-
-# Copied from transformers.models.convnext.modeling_tf_convnext.TFConvNextEmbeddings with ConvNext->ConvNextV2
-class TFConvNextV2Embeddings(keras.layers.Layer):
- """This class is comparable to (and inspired by) the SwinEmbeddings class
- found in src/transformers/models/swin/modeling_swin.py.
- """
-
- def __init__(self, config: ConvNextV2Config, **kwargs):
- super().__init__(**kwargs)
- self.patch_embeddings = keras.layers.Conv2D(
- filters=config.hidden_sizes[0],
- kernel_size=config.patch_size,
- strides=config.patch_size,
- name="patch_embeddings",
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer=keras.initializers.Zeros(),
- )
- self.layernorm = keras.layers.LayerNormalization(epsilon=1e-6, name="layernorm")
- self.num_channels = config.num_channels
- self.config = config
-
- def call(self, pixel_values):
- if isinstance(pixel_values, dict):
- pixel_values = pixel_values["pixel_values"]
-
- tf.debugging.assert_equal(
- shape_list(pixel_values)[1],
- self.num_channels,
- message="Make sure that the channel dimension of the pixel values match with the one set in the configuration.",
- )
-
- # When running on CPU, `keras.layers.Conv2D` doesn't support `NCHW` format.
- # So change the input format from `NCHW` to `NHWC`.
- # shape = (batch_size, in_height, in_width, in_channels)
- pixel_values = tf.transpose(pixel_values, perm=(0, 2, 3, 1))
-
- embeddings = self.patch_embeddings(pixel_values)
- embeddings = self.layernorm(embeddings)
- return embeddings
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "patch_embeddings", None) is not None:
- with tf.name_scope(self.patch_embeddings.name):
- self.patch_embeddings.build([None, None, None, self.config.num_channels])
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, None, None, self.config.hidden_sizes[0]])
-
-
-class TFConvNextV2Layer(keras.layers.Layer):
- """This corresponds to the `Block` class in the original implementation.
-
- There are two equivalent implementations: [DwConv, LayerNorm (channels_first), Conv, GELU,1x1 Conv]; all in (N, C,
- H, W) (2) [DwConv, Permute to (N, H, W, C), LayerNorm (channels_last), Linear, GELU, Linear]; Permute back
-
- The authors used (2) as they find it slightly faster in PyTorch. Since we already permuted the inputs to follow
- NHWC ordering, we can just apply the operations straight-away without the permutation.
-
- Args:
- config (`ConvNextV2Config`):
- Model configuration class.
- dim (`int`):
- Number of input channels.
- drop_path (`float`, defaults to 0.0):
- Stochastic depth rate.
- """
-
- def __init__(self, config: ConvNextV2Config, dim: int, drop_path: float = 0.0, **kwargs):
- super().__init__(**kwargs)
- self.dim = dim
- self.config = config
- self.dwconv = keras.layers.Conv2D(
- filters=dim,
- kernel_size=7,
- padding="same",
- groups=dim,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer=keras.initializers.Zeros(),
- name="dwconv",
- ) # depthwise conv
- self.layernorm = keras.layers.LayerNormalization(
- epsilon=1e-6,
- name="layernorm",
- )
- self.pwconv1 = keras.layers.Dense(
- units=4 * dim,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer=keras.initializers.Zeros(),
- name="pwconv1",
- ) # pointwise/1x1 convs, implemented with linear layers
- self.act = get_tf_activation(config.hidden_act)
- self.grn = TFConvNextV2GRN(config, 4 * dim, dtype=tf.float32, name="grn")
- self.pwconv2 = keras.layers.Dense(
- units=dim,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer=keras.initializers.Zeros(),
- name="pwconv2",
- )
- # Using `layers.Activation` instead of `tf.identity` to better control `training`
- # behaviour.
- self.drop_path = (
- TFConvNextV2DropPath(drop_path, name="drop_path")
- if drop_path > 0.0
- else keras.layers.Activation("linear", name="drop_path")
- )
-
- def call(self, hidden_states, training=False):
- input = hidden_states
- x = self.dwconv(hidden_states)
- x = self.layernorm(x)
- x = self.pwconv1(x)
- x = self.act(x)
- x = self.grn(x)
- x = self.pwconv2(x)
- x = self.drop_path(x, training=training)
- x = input + x
- return x
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dwconv", None) is not None:
- with tf.name_scope(self.dwconv.name):
- self.dwconv.build([None, None, None, self.dim])
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, None, None, self.dim])
- if getattr(self, "pwconv1", None) is not None:
- with tf.name_scope(self.pwconv1.name):
- self.pwconv1.build([None, None, self.dim])
- if getattr(self, "grn", None) is not None:
- with tf.name_scope(self.grn.name):
- self.grn.build(None)
- if getattr(self, "pwconv2", None) is not None:
- with tf.name_scope(self.pwconv2.name):
- self.pwconv2.build([None, None, 4 * self.dim])
- if getattr(self, "drop_path", None) is not None:
- with tf.name_scope(self.drop_path.name):
- self.drop_path.build(None)
-
-
-# Copied from transformers.models.convnext.modeling_tf_convnext.TFConvNextStage with ConvNext->ConvNextV2
-class TFConvNextV2Stage(keras.layers.Layer):
- """ConvNextV2 stage, consisting of an optional downsampling layer + multiple residual blocks.
-
- Args:
- config (`ConvNextV2V2Config`):
- Model configuration class.
- in_channels (`int`):
- Number of input channels.
- out_channels (`int`):
- Number of output channels.
- depth (`int`):
- Number of residual blocks.
- drop_path_rates(`List[float]`):
- Stochastic depth rates for each layer.
- """
-
- def __init__(
- self,
- config: ConvNextV2Config,
- in_channels: int,
- out_channels: int,
- kernel_size: int = 2,
- stride: int = 2,
- depth: int = 2,
- drop_path_rates: Optional[List[float]] = None,
- **kwargs,
- ):
- super().__init__(**kwargs)
- if in_channels != out_channels or stride > 1:
- self.downsampling_layer = [
- keras.layers.LayerNormalization(
- epsilon=1e-6,
- name="downsampling_layer.0",
- ),
- # Inputs to this layer will follow NHWC format since we
- # transposed the inputs from NCHW to NHWC in the `TFConvNextV2Embeddings`
- # layer. All the outputs throughout the model will be in NHWC
- # from this point on until the output where we again change to
- # NCHW.
- keras.layers.Conv2D(
- filters=out_channels,
- kernel_size=kernel_size,
- strides=stride,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer=keras.initializers.Zeros(),
- name="downsampling_layer.1",
- ),
- ]
- else:
- self.downsampling_layer = [tf.identity]
-
- drop_path_rates = drop_path_rates or [0.0] * depth
- self.layers = [
- TFConvNextV2Layer(
- config,
- dim=out_channels,
- drop_path=drop_path_rates[j],
- name=f"layers.{j}",
- )
- for j in range(depth)
- ]
- self.in_channels = in_channels
- self.out_channels = out_channels
- self.stride = stride
-
- def call(self, hidden_states):
- for layer in self.downsampling_layer:
- hidden_states = layer(hidden_states)
- for layer in self.layers:
- hidden_states = layer(hidden_states)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "layers", None) is not None:
- for layer in self.layers:
- with tf.name_scope(layer.name):
- layer.build(None)
- if self.in_channels != self.out_channels or self.stride > 1:
- with tf.name_scope(self.downsampling_layer[0].name):
- self.downsampling_layer[0].build([None, None, None, self.in_channels])
- with tf.name_scope(self.downsampling_layer[1].name):
- self.downsampling_layer[1].build([None, None, None, self.in_channels])
-
-
-class TFConvNextV2Encoder(keras.layers.Layer):
- def __init__(self, config: ConvNextV2Config, **kwargs):
- super().__init__(**kwargs)
- self.stages = []
- drop_path_rates = tf.linspace(0.0, config.drop_path_rate, sum(config.depths))
- drop_path_rates = tf.split(drop_path_rates, config.depths)
- drop_path_rates = [x.numpy().tolist() for x in drop_path_rates]
- prev_chs = config.hidden_sizes[0]
- for i in range(config.num_stages):
- out_chs = config.hidden_sizes[i]
- stage = TFConvNextV2Stage(
- config,
- in_channels=prev_chs,
- out_channels=out_chs,
- stride=2 if i > 0 else 1,
- depth=config.depths[i],
- drop_path_rates=drop_path_rates[i],
- name=f"stages.{i}",
- )
- self.stages.append(stage)
- prev_chs = out_chs
-
- def call(
- self,
- hidden_states: tf.Tensor,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple, TFBaseModelOutputWithNoAttention]:
- all_hidden_states = () if output_hidden_states else None
-
- for i, layer_module in enumerate(self.stages):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- hidden_states = layer_module(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
-
- return TFBaseModelOutputWithNoAttention(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
-
- def build(self, input_shape=None):
- for stage in self.stages:
- with tf.name_scope(stage.name):
- stage.build(None)
-
-
-@keras_serializable
-class TFConvNextV2MainLayer(keras.layers.Layer):
- config_class = ConvNextV2Config
-
- def __init__(self, config: ConvNextV2Config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embeddings = TFConvNextV2Embeddings(config, name="embeddings")
- self.encoder = TFConvNextV2Encoder(config, name="encoder")
- self.layernorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
- # We are setting the `data_format` like so because from here on we will revert to the
- # NCHW output format
- self.pooler = keras.layers.GlobalAvgPool2D(data_format="channels_last")
-
- @unpack_inputs
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor]]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- embedding_output = self.embeddings(pixel_values, training=training)
-
- encoder_outputs = self.encoder(
- embedding_output,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- last_hidden_state = encoder_outputs[0]
-
- # Change to NCHW output format have uniformity in the modules
- pooled_output = self.pooler(last_hidden_state)
- last_hidden_state = tf.transpose(last_hidden_state, perm=(0, 3, 1, 2))
- pooled_output = self.layernorm(pooled_output)
-
- # Change the other hidden state outputs to NCHW as well
- if output_hidden_states:
- hidden_states = tuple([tf.transpose(h, perm=(0, 3, 1, 2)) for h in encoder_outputs[1]])
-
- if not return_dict:
- hidden_states = hidden_states if output_hidden_states else ()
- return (last_hidden_state, pooled_output) + hidden_states
-
- return TFBaseModelOutputWithPoolingAndNoAttention(
- last_hidden_state=last_hidden_state,
- pooler_output=pooled_output,
- hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, self.config.hidden_sizes[-1]])
-
-
-class TFConvNextV2PreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ConvNextV2Config
- base_model_prefix = "convnextv2"
- main_input_name = "pixel_values"
-
-
-CONVNEXTV2_START_DOCSTRING = r"""
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `pixel_values` only and nothing else: `model(pixel_values)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([pixel_values, attention_mask])` or `model([pixel_values, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"pixel_values": pixel_values, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`ConvNextV2Config`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CONVNEXTV2_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`np.ndarray`, `tf.Tensor`, `List[tf.Tensor]`, `Dict[str, tf.Tensor]` or `Dict[str, np.ndarray]` and each example must have the shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`ConvNextImageProcessor.__call__`] for details.
-
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to `True`.
-"""
-
-
-@add_start_docstrings(
- "The bare ConvNextV2 model outputting raw features without any specific head on top.",
- CONVNEXTV2_START_DOCSTRING,
-)
-class TFConvNextV2Model(TFConvNextV2PreTrainedModel):
- def __init__(self, config: ConvNextV2Config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.convnextv2 = TFConvNextV2MainLayer(config, name="convnextv2")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVNEXTV2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutputWithPoolingAndNoAttention,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- outputs = self.convnextv2(
- pixel_values=pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- if not return_dict:
- return outputs[:]
-
- return TFBaseModelOutputWithPoolingAndNoAttention(
- last_hidden_state=outputs.last_hidden_state,
- pooler_output=outputs.pooler_output,
- hidden_states=outputs.hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convnextv2", None) is not None:
- with tf.name_scope(self.convnextv2.name):
- self.convnextv2.build(None)
-
-
-@add_start_docstrings(
- """
- ConvNextV2 Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
- ImageNet.
- """,
- CONVNEXTV2_START_DOCSTRING,
-)
-class TFConvNextV2ForImageClassification(TFConvNextV2PreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: ConvNextV2Config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
- self.convnextv2 = TFConvNextV2MainLayer(config, name="convnextv2")
-
- # Classifier head
- self.classifier = keras.layers.Dense(
- units=config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- bias_initializer=keras.initializers.Zeros(),
- name="classifier",
- )
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CONVNEXTV2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=TFImageClassifierOutputWithNoAttention,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFImageClassifierOutputWithNoAttention, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- outputs = self.convnextv2(
- pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- pooled_output = outputs.pooler_output if return_dict else outputs[1]
-
- logits = self.classifier(pooled_output)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFImageClassifierOutputWithNoAttention(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convnextv2", None) is not None:
- with tf.name_scope(self.convnextv2.name):
- self.convnextv2.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_sizes[-1]])
diff --git a/transformers/models/cpm/__init__.py b/transformers/models/cpm/__init__.py
deleted file mode 100644
index be6b0f66898ecbef786f311097f6a49c676762bd..0000000000000000000000000000000000000000
--- a/transformers/models/cpm/__init__.py
+++ /dev/null
@@ -1,59 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_tokenizers_available
-
-
-_import_structure = {}
-
-try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_cpm"] = ["CpmTokenizer"]
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_cpm_fast"] = ["CpmTokenizerFast"]
-
-
-if TYPE_CHECKING:
- try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_cpm import CpmTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_cpm_fast import CpmTokenizerFast
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/cpm/__pycache__/__init__.cpython-310.pyc b/transformers/models/cpm/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 6208e7439fece13f667fa04f00ab37c29842417f..0000000000000000000000000000000000000000
Binary files a/transformers/models/cpm/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cpm/__pycache__/tokenization_cpm.cpython-310.pyc b/transformers/models/cpm/__pycache__/tokenization_cpm.cpython-310.pyc
deleted file mode 100644
index 3ac0a99d2927dfb79021c6b1ceab5797fae737c7..0000000000000000000000000000000000000000
Binary files a/transformers/models/cpm/__pycache__/tokenization_cpm.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cpm/__pycache__/tokenization_cpm_fast.cpython-310.pyc b/transformers/models/cpm/__pycache__/tokenization_cpm_fast.cpython-310.pyc
deleted file mode 100644
index 77f1e49e176ecfb3ea777789c005909bf0ad2c3a..0000000000000000000000000000000000000000
Binary files a/transformers/models/cpm/__pycache__/tokenization_cpm_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cpm/tokenization_cpm.py b/transformers/models/cpm/tokenization_cpm.py
deleted file mode 100644
index ac454898b5572a1bd9f1143d25ba52558f5b5ac9..0000000000000000000000000000000000000000
--- a/transformers/models/cpm/tokenization_cpm.py
+++ /dev/null
@@ -1,344 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes."""
-import os
-import unicodedata
-from shutil import copyfile
-from typing import Any, Dict, List, Optional, Tuple
-
-import sentencepiece as spm
-
-from ...tokenization_utils import AddedToken, PreTrainedTokenizer
-from ...utils import SPIECE_UNDERLINE, logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}
-
-
-class CpmTokenizer(PreTrainedTokenizer):
- """Runs pre-tokenization with Jieba segmentation tool. It is used in CPM models."""
-
- vocab_files_names = VOCAB_FILES_NAMES
-
- def __init__(
- self,
- vocab_file,
- do_lower_case=False,
- remove_space=True,
- keep_accents=False,
- bos_token="",
- eos_token="",
- unk_token="",
- sep_token="",
- pad_token="",
- cls_token="",
- mask_token="",
- additional_special_tokens=["", ""],
- sp_model_kwargs: Optional[Dict[str, Any]] = None,
- **kwargs,
- ) -> None:
- """
- Construct a CPM tokenizer. Based on [Jieba](https://pypi.org/project/jieba/) and
- [SentencePiece](https://github.com/google/sentencepiece).
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .spm extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether to lowercase the input when tokenizing.
- remove_space (`bool`, *optional*, defaults to `True`):
- Whether to strip the text when tokenizing (removing excess spaces before and after the string).
- keep_accents (`bool`, *optional*, defaults to `False`):
- Whether to keep accents when tokenizing.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier
- token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the beginning of
- sequence. The token used is the `cls_token`.
-
-
-
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sequence token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the end of
- sequence. The token used is the `sep_token`.
-
-
-
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be
- this token instead.
- sep_token (`str`, *optional*, defaults to `""`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences
- for sequence classification or for a text and a question for question answering. It is also used as the
- last token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `""`):
- The classifier token which is used when doing sequence classification (classification of the whole
- sequence instead of per-token classification). It is the first token of the sequence when built with
- special tokens.
- mask_token (`str`, *optional*, defaults to `""`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- additional_special_tokens (`List[str]`, *optional*, defaults to `["", ""]`):
- Additional special tokens used by the tokenizer.
-
- Attributes:
- sp_model (`SentencePieceProcessor`):
- The *SentencePiece* processor that is used for every conversion (string, tokens and IDs).
- """
- # Mask token behave like a normal word, i.e. include the space before it
- mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token
-
- self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
-
- self.do_lower_case = do_lower_case
- self.remove_space = remove_space
- self.keep_accents = keep_accents
- self.vocab_file = vocab_file
-
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(vocab_file)
-
- try:
- import jieba
- except ModuleNotFoundError as error:
- raise error.__class__(
- "You need to install jieba to use CpmTokenizer or CpmTokenizerFast. "
- "See https://pypi.org/project/jieba/ for installation."
- )
- self.jieba = jieba
- self.translator = str.maketrans(" \n", "\u2582\u2583")
-
- super().__init__(
- do_lower_case=do_lower_case,
- remove_space=remove_space,
- keep_accents=keep_accents,
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- additional_special_tokens=additional_special_tokens,
- sp_model_kwargs=self.sp_model_kwargs,
- **kwargs,
- )
-
- self._pad_token_type_id = 3
-
- @property
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.vocab_size
- def vocab_size(self):
- return len(self.sp_model)
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.get_vocab
- def get_vocab(self):
- vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
- vocab.update(self.added_tokens_encoder)
- return vocab
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.__getstate__
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sp_model"] = None
- return state
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.__setstate__
- def __setstate__(self, d):
- self.__dict__ = d
-
- # for backward compatibility
- if not hasattr(self, "sp_model_kwargs"):
- self.sp_model_kwargs = {}
-
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(self.vocab_file)
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.preprocess_text
- def preprocess_text(self, inputs):
- if self.remove_space:
- outputs = " ".join(inputs.strip().split())
- else:
- outputs = inputs
- outputs = outputs.replace("``", '"').replace("''", '"')
-
- if not self.keep_accents:
- outputs = unicodedata.normalize("NFKD", outputs)
- outputs = "".join([c for c in outputs if not unicodedata.combining(c)])
- if self.do_lower_case:
- outputs = outputs.lower()
-
- return outputs
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer._tokenize
- def _tokenize(self, text: str) -> List[str]:
- """Tokenize a string."""
- text = self.preprocess_text(text)
- pieces = self.sp_model.encode(text, out_type=str)
- new_pieces = []
- for piece in pieces:
- if len(piece) > 1 and piece[-1] == str(",") and piece[-2].isdigit():
- cur_pieces = self.sp_model.EncodeAsPieces(piece[:-1].replace(SPIECE_UNDERLINE, ""))
- if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
- if len(cur_pieces[0]) == 1:
- cur_pieces = cur_pieces[1:]
- else:
- cur_pieces[0] = cur_pieces[0][1:]
- cur_pieces.append(piece[-1])
- new_pieces.extend(cur_pieces)
- else:
- new_pieces.append(piece)
-
- return new_pieces
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.sp_model.PieceToId(token)
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.sp_model.IdToPiece(index)
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.convert_tokens_to_string
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (strings for sub-words) in a single string."""
- out_string = "".join(tokens).replace(SPIECE_UNDERLINE, " ").strip()
- return out_string
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. An XLNet sequence has the following format:
-
- - single sequence: `X `
- - pair of sequences: `A B `
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return token_ids_0 + sep + cls
- return token_ids_0 + sep + token_ids_1 + sep + cls
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.get_special_tokens_mask
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1, 1]
- return ([0] * len(token_ids_0)) + [1, 1]
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. An XLNet
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls_segment_id = [2]
-
- if token_ids_1 is None:
- return len(token_ids_0 + sep) * [0] + cls_segment_id
- return len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1] + cls_segment_id
-
- # Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
- elif not os.path.isfile(self.vocab_file):
- with open(out_vocab_file, "wb") as fi:
- content_spiece_model = self.sp_model.serialized_model_proto()
- fi.write(content_spiece_model)
-
- return (out_vocab_file,)
-
- def _decode(self, *args, **kwargs):
- text = super()._decode(*args, **kwargs)
- text = text.replace(" ", "").replace("\u2582", " ").replace("\u2583", "\n")
- return text
diff --git a/transformers/models/cpm/tokenization_cpm_fast.py b/transformers/models/cpm/tokenization_cpm_fast.py
deleted file mode 100644
index 9b7b6da118ab4beb162f5e73e7a2cd2d12366997..0000000000000000000000000000000000000000
--- a/transformers/models/cpm/tokenization_cpm_fast.py
+++ /dev/null
@@ -1,237 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes."""
-import os
-from shutil import copyfile
-from typing import List, Optional, Tuple
-
-from ...tokenization_utils_fast import AddedToken, PreTrainedTokenizerFast
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}
-
-
-class CpmTokenizerFast(PreTrainedTokenizerFast):
- """Runs pre-tokenization with Jieba segmentation tool. It is used in CPM models."""
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- do_lower_case=False,
- remove_space=True,
- keep_accents=False,
- bos_token="",
- eos_token="",
- unk_token="",
- sep_token="",
- pad_token="",
- cls_token="",
- mask_token="",
- additional_special_tokens=["", ""],
- **kwargs,
- ):
- """
- Construct a CPM tokenizer. Based on [Jieba](https://pypi.org/project/jieba/) and
- [SentencePiece](https://github.com/google/sentencepiece).
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .spm extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether to lowercase the input when tokenizing.
- remove_space (`bool`, *optional*, defaults to `True`):
- Whether to strip the text when tokenizing (removing excess spaces before and after the string).
- keep_accents (`bool`, *optional*, defaults to `False`):
- Whether to keep accents when tokenizing.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier
- token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the beginning of
- sequence. The token used is the `cls_token`.
-
-
-
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sequence token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the end of
- sequence. The token used is the `sep_token`.
-
-
-
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be
- this token instead.
- sep_token (`str`, *optional*, defaults to `""`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences
- for sequence classification or for a text and a question for question answering. It is also used as the
- last token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `""`):
- The classifier token which is used when doing sequence classification (classification of the whole
- sequence instead of per-token classification). It is the first token of the sequence when built with
- special tokens.
- mask_token (`str`, *optional*, defaults to `""`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- additional_special_tokens (`List[str]`, *optional*, defaults to `["", ""]`):
- Additional special tokens used by the tokenizer.
-
- Attributes:
- sp_model (`SentencePieceProcessor`):
- The *SentencePiece* processor that is used for every conversion (string, tokens and IDs).
- """
- # Mask token behave like a normal word, i.e. include the space before it
- mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token
-
- super().__init__(
- vocab_file=vocab_file,
- tokenizer_file=tokenizer_file,
- do_lower_case=do_lower_case,
- remove_space=remove_space,
- keep_accents=keep_accents,
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- additional_special_tokens=additional_special_tokens,
- **kwargs,
- )
-
- self._pad_token_type_id = 3
- self.do_lower_case = do_lower_case
- self.remove_space = remove_space
- self.keep_accents = keep_accents
- self.vocab_file = vocab_file
-
- try:
- import jieba
- except ModuleNotFoundError as error:
- raise error.__class__(
- "You need to install jieba to use CpmTokenizer or CpmTokenizerFast. "
- "See https://pypi.org/project/jieba/ for installation."
- )
- self.jieba = jieba
- self.translator = str.maketrans(" \n", "\u2582\u2583")
-
- @property
- def can_save_slow_tokenizer(self) -> bool:
- return os.path.isfile(self.vocab_file) if self.vocab_file else False
-
- # Copied from transformers.models.xlnet.tokenization_xlnet_fast.XLNetTokenizerFast.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. An XLNet sequence has the following format:
-
- - single sequence: `X `
- - pair of sequences: `A B `
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return token_ids_0 + sep + cls
- return token_ids_0 + sep + token_ids_1 + sep + cls
-
- # Copied from transformers.models.xlnet.tokenization_xlnet_fast.XLNetTokenizerFast.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. An XLNet
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls_segment_id = [2]
-
- if token_ids_1 is None:
- return len(token_ids_0 + sep) * [0] + cls_segment_id
- return len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1] + cls_segment_id
-
- # Copied from transformers.models.xlnet.tokenization_xlnet_fast.XLNetTokenizerFast.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not self.can_save_slow_tokenizer:
- raise ValueError(
- "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
- "tokenizer."
- )
-
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
-
- return (out_vocab_file,)
-
- def _batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
- batch_text_or_text_pairs = [
- " ".join([x.translate(self.translator) for x in self.jieba.cut(text, cut_all=False)])
- for text in batch_text_or_text_pairs
- ]
- return super()._batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
-
- def _decode(self, *args, **kwargs):
- text = super()._decode(*args, **kwargs)
- text = text.replace(" ", "").replace("\u2582", " ").replace("\u2583", "\n")
- return text
diff --git a/transformers/models/cpmant/__init__.py b/transformers/models/cpmant/__init__.py
deleted file mode 100644
index 8140009b60f15680663fc61569f55675e6d71196..0000000000000000000000000000000000000000
--- a/transformers/models/cpmant/__init__.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# flake8: noqa
-# There's no way to ignore "F401 '...' imported but unused" warnings in this
-# module, but to preserve other warnings. So, don't check this module at all.
-
-# Copyright 2022 The HuggingFace Team and The OpenBMB Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-# rely on isort to merge the imports
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available
-
-
-_import_structure = {
- "configuration_cpmant": ["CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP", "CpmAntConfig"],
- "tokenization_cpmant": ["CpmAntTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_cpmant"] = [
- "CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "CpmAntForCausalLM",
- "CpmAntModel",
- "CpmAntPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_cpmant import CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP, CpmAntConfig
- from .tokenization_cpmant import CpmAntTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_cpmant import (
- CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST,
- CpmAntForCausalLM,
- CpmAntModel,
- CpmAntPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/cpmant/__pycache__/__init__.cpython-310.pyc b/transformers/models/cpmant/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index f83e39441e608dfb73245092477d79cf190c4632..0000000000000000000000000000000000000000
Binary files a/transformers/models/cpmant/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cpmant/__pycache__/configuration_cpmant.cpython-310.pyc b/transformers/models/cpmant/__pycache__/configuration_cpmant.cpython-310.pyc
deleted file mode 100644
index 589d30b21ca6eadb8d9406b0bd31c6690869051b..0000000000000000000000000000000000000000
Binary files a/transformers/models/cpmant/__pycache__/configuration_cpmant.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cpmant/__pycache__/modeling_cpmant.cpython-310.pyc b/transformers/models/cpmant/__pycache__/modeling_cpmant.cpython-310.pyc
deleted file mode 100644
index 5e7ca54d93f678d971ad7b576e827c9a18e90925..0000000000000000000000000000000000000000
Binary files a/transformers/models/cpmant/__pycache__/modeling_cpmant.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cpmant/__pycache__/tokenization_cpmant.cpython-310.pyc b/transformers/models/cpmant/__pycache__/tokenization_cpmant.cpython-310.pyc
deleted file mode 100644
index 8c9ea88fa7d0c826178b9e183bf7156bd30e8002..0000000000000000000000000000000000000000
Binary files a/transformers/models/cpmant/__pycache__/tokenization_cpmant.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cpmant/configuration_cpmant.py b/transformers/models/cpmant/configuration_cpmant.py
deleted file mode 100644
index 62bbce8ada50e11dccf82468fff1c2417f44ab79..0000000000000000000000000000000000000000
--- a/transformers/models/cpmant/configuration_cpmant.py
+++ /dev/null
@@ -1,122 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" CPMAnt model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class CpmAntConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`CpmAntModel`]. It is used to instantiate an
- CPMAnt model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the CPMAnt
- [openbmb/cpm-ant-10b](https://huggingface.co/openbmb/cpm-ant-10b) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30720):
- Vocabulary size of the CPMAnt model. Defines the number of different tokens that can be represented by the
- `input` passed when calling [`CpmAntModel`].
- hidden_size (`int`, *optional*, defaults to 4096):
- Dimension of the encoder layers.
- num_attention_heads (`int`, *optional*, defaults to 32):
- Number of attention heads in the Transformer encoder.
- dim_head (`int`, *optional*, defaults to 128):
- Dimension of attention heads for each attention layer in the Transformer encoder.
- dim_ff (`int`, *optional*, defaults to 10240):
- Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- num_hidden_layers (`int`, *optional*, defaults to 48):
- Number of layers of the Transformer encoder.
- dropout_p (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder.
- position_bias_num_buckets (`int`, *optional*, defaults to 512):
- The number of position_bias buckets.
- position_bias_max_distance (`int`, *optional*, defaults to 2048):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- eps (`float`, *optional*, defaults to 1e-06):
- The epsilon used by the layer normalization layers.
- init_std (`float`, *optional*, defaults to 1.0):
- Initialize parameters with std = init_std.
- prompt_types (`int`, *optional*, defaults to 32):
- The type of prompt.
- prompt_length (`int`, *optional*, defaults to 32):
- The length of prompt.
- segment_types (`int`, *optional*, defaults to 32):
- The type of segment.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether to use cache.
-
- Example:
-
- ```python
- >>> from transformers import CpmAntModel, CpmAntConfig
-
- >>> # Initializing a CPMAnt cpm-ant-10b style configuration
- >>> configuration = CpmAntConfig()
-
- >>> # Initializing a model from the cpm-ant-10b style configuration
- >>> model = CpmAntModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "cpmant"
-
- def __init__(
- self,
- vocab_size: int = 30720,
- hidden_size: int = 4096,
- num_attention_heads: int = 32,
- dim_head: int = 128,
- dim_ff: int = 10240,
- num_hidden_layers: int = 48,
- dropout_p: int = 0.0,
- position_bias_num_buckets: int = 512,
- position_bias_max_distance: int = 2048,
- eps: int = 1e-6,
- init_std: float = 1.0,
- prompt_types: int = 32,
- prompt_length: int = 32,
- segment_types: int = 32,
- use_cache: bool = True,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.prompt_types = prompt_types
- self.prompt_length = prompt_length
- self.segment_types = segment_types
- self.hidden_size = hidden_size
- self.num_attention_heads = num_attention_heads
- self.dim_head = dim_head
- self.dim_ff = dim_ff
- self.num_hidden_layers = num_hidden_layers
- self.position_bias_num_buckets = position_bias_num_buckets
- self.position_bias_max_distance = position_bias_max_distance
- self.dropout_p = dropout_p
- self.eps = eps
- self.use_cache = use_cache
- self.vocab_size = vocab_size
- self.init_std = init_std
diff --git a/transformers/models/cpmant/modeling_cpmant.py b/transformers/models/cpmant/modeling_cpmant.py
deleted file mode 100644
index 63bb467e64e3542e4385c39c14ef7664ee1cdd71..0000000000000000000000000000000000000000
--- a/transformers/models/cpmant/modeling_cpmant.py
+++ /dev/null
@@ -1,872 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch CPMAnt"""
-
-
-import math
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
-from ...modeling_utils import PreTrainedModel
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_cpmant import CpmAntConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "openbmb/cpm-ant-10b"
-_CONFIG_FOR_DOC = "CpmAntConfig"
-
-
-from ..deprecated._archive_maps import CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class CpmAntLayerNorm(nn.Module):
- """
- We use Root Mean Square (RMS) Layer Normalization, please see https://arxiv.org/abs/1910.07467 for details."
- """
-
- def __init__(self, config: CpmAntConfig):
- super().__init__()
-
- self.eps = config.eps
- self.dim_norm = config.hidden_size
- self.weight = nn.Parameter(torch.empty(config.hidden_size))
-
- def forward(self, hidden_states: torch.Tensor):
- """
- Args:
- hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
- """
- if hidden_states.size(-1) != self.dim_norm:
- raise AssertionError("hidden_states.size(-1) != self.dim_norm")
- old_dtype = hidden_states.dtype
- variance = hidden_states.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
- hidden_states = (hidden_states * torch.rsqrt(variance + self.eps)).to(old_dtype) * self.weight
- return hidden_states
-
-
-class CpmAntAttention(nn.Module):
- def __init__(self, config: CpmAntConfig):
- super().__init__()
- self.dim_model = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.dim_head = config.dim_head
-
- self.project_q = nn.Linear(self.dim_model, self.num_heads * self.dim_head, bias=False)
- self.project_k = nn.Linear(self.dim_model, self.num_heads * self.dim_head, bias=False)
- self.project_v = nn.Linear(self.dim_model, self.num_heads * self.dim_head, bias=False)
-
- self.attention_out = nn.Linear(self.num_heads * self.dim_head, self.dim_model, bias=False)
-
- self.softmax = torch.nn.Softmax(dim=-1)
-
- if config.dropout_p is not None:
- self.dropout = torch.nn.Dropout(p=config.dropout_p)
- else:
- self.dropout = None
-
- def forward(
- self,
- hidden_q: torch.Tensor,
- hidden_kv: torch.Tensor,
- attention_mask: torch.BoolTensor,
- position_bias: torch.Tensor,
- output_attentions: Optional[bool] = False,
- past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- use_cache: Optional[bool] = None,
- ):
- """
- Args:
- hidden_q (`torch.Tensor`):
- Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
- hidden_kv (`torch.Tensor` of shape `(batch, len_k, dim_model)`)):
- Tensor *key_value* and *query* of shape `(batch, len_k, dim_model)`
- attention_mask (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
- Avoid invalid areas to participate in the calculation of self-attention.
- position_bias (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
- Provide positional information to self-attention block.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers.
- past_key_values (`Tuple[torch.Tensor, torch.Tensor]`, *optional*):
- Cached past key and value projection states.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- """
- batch_size = hidden_q.size(0)
- len_q = hidden_q.size(1)
- len_k = hidden_kv.size(1)
-
- query = self.project_q(hidden_q)
- key = self.project_k(hidden_kv)
- value = self.project_v(hidden_kv)
-
- query = query.view(batch_size, len_q, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
- key = key.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
- value = value.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
-
- if past_key_values is not None:
- key = torch.cat([past_key_values[0], key], dim=-2)
- value = torch.cat([past_key_values[1], value], dim=-2)
- len_k = key.size(-2)
-
- # (batch_size, num_heads, len_q, dim_head) @ (batch_size, num_heads, dim_head, len_k) -> (batch_size, num_heads, len_q, len_k)
- score = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.dim_head)
- score = score + position_bias
-
- score = torch.masked_fill(
- score,
- attention_mask.view(batch_size, 1, len_q, len_k) == torch.tensor(False),
- torch.scalar_tensor(float("-inf"), device=score.device, dtype=score.dtype),
- )
- score = self.softmax(score)
-
- score = torch.masked_fill(
- score,
- attention_mask.view(batch_size, 1, len_q, len_k) == torch.tensor(False),
- torch.scalar_tensor(0, device=score.device, dtype=score.dtype),
- )
- if output_attentions:
- attn_weights = score
- else:
- attn_weights = None
-
- if self.dropout is not None:
- score = self.dropout(score)
-
- # (batch_size, num_heads, len_q, len_k) @ (batch_size, num_heads, len_k, dim_head) -> (batch_size, num_heads, len_q, dim_head)
- score = torch.matmul(score, value)
-
- score = score.view(batch_size, self.num_heads, len_q, self.dim_head).permute(0, 2, 1, 3)
- score = score.contiguous().view(batch_size, len_q, self.num_heads * self.dim_head)
-
- score = self.attention_out(score)
-
- past_key_values = None
- if use_cache:
- past_key_values = (key, value)
-
- return score, attn_weights, past_key_values
-
-
-class CpmAntSelfAttentionBlock(nn.Module):
- def __init__(self, config: CpmAntConfig):
- super().__init__()
- self.layernorm_before_attention = CpmAntLayerNorm(config)
- self.self_attention = CpmAntAttention(config)
- if config.dropout_p:
- self.dropout = torch.nn.Dropout(config.dropout_p)
- else:
- self.dropout = None
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: torch.Tensor,
- position_bias: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- use_cache: Optional[bool] = None,
- ):
- """
- Args:
- hidden_states (`torch.Tensor` of shape `(batch, len_seq, dim_model)`):
- Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
- attention_mask (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
- Avoid invalid areas to participate in the calculation of self-attention.
- position_bias (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
- Provide positional information to self-attention block.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers.
- past_key_values (`Tuple(torch.FloatTensor)`, *optional*):
- Cached past key and value projection states.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- """
- outputs = self.layernorm_before_attention(hidden_states)
- outputs = self.self_attention(
- outputs, outputs, attention_mask, position_bias, output_attentions, past_key_values, use_cache
- )
-
- outputs, attn_weights, current_key_value = outputs
-
- if self.dropout is not None:
- outputs = self.dropout(outputs)
- hidden_states = hidden_states + outputs
-
- return hidden_states, attn_weights, current_key_value
-
-
-class CpmAntDenseGatedACT(nn.Module):
- def __init__(self, config: CpmAntConfig):
- super().__init__()
- self.w_0 = nn.Linear(config.hidden_size, config.dim_ff, bias=False)
- self.w_1 = nn.Linear(config.hidden_size, config.dim_ff, bias=False)
- self.act = torch.nn.GELU()
-
- def forward(self, hidden_states: torch.Tensor):
- """Transform an input tensor from one feature space to another via a nonlinear operation
-
- Args:
- hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
- """
- gate_score = self.act(self.w_0(hidden_states))
- hidden_states = self.w_1(hidden_states)
-
- hidden_states = gate_score * hidden_states
- return hidden_states
-
-
-class CpmAntFeedForward(nn.Module):
- def __init__(self, config: CpmAntConfig):
- super().__init__()
- self.w_in = CpmAntDenseGatedACT(config)
- if config.dropout_p is not None:
- self.dropout = torch.nn.Dropout(config.dropout_p)
- else:
- self.dropout = None
-
- self.w_out = nn.Linear(config.dim_ff, config.hidden_size, bias=False)
-
- def forward(self, hidden_states: torch.Tensor):
- """
- Args:
- hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
- """
- hidden_states = self.w_in(hidden_states)
-
- if self.dropout is not None:
- hidden_states = self.dropout(hidden_states)
-
- hidden_states = self.w_out(hidden_states)
-
- return hidden_states
-
-
-class CpmAntFFNBlock(nn.Module):
- def __init__(self, config: CpmAntConfig):
- super().__init__()
- self.layernorm_before_ffn = CpmAntLayerNorm(config)
- self.ffn = CpmAntFeedForward(config)
- if config.dropout_p:
- self.dropout = torch.nn.Dropout(config.dropout_p)
- else:
- self.dropout = None
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- ):
- """
- Args:
- hidden_states (`torch.Tensor` of shape `(batch, len_seq, dim_model)`):
- Hidden states before feed forward layer.
- """
- ln_outputs = self.layernorm_before_ffn(hidden_states)
- outputs = self.ffn(ln_outputs)
- if self.dropout is not None:
- outputs = self.dropout(outputs)
- hidden_states = hidden_states + outputs
- return hidden_states
-
-
-class CpmAntTransformerBlock(nn.Module):
- def __init__(self, config: CpmAntConfig):
- super().__init__()
- self.self_att = CpmAntSelfAttentionBlock(config)
- self.ffn = CpmAntFFNBlock(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: torch.Tensor,
- position_bias: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- use_cache: Optional[bool] = None,
- ):
- """
- Args:
- hidden_states (`torch.Tensor`):
- Input to the layer of shape `(batch, seq_len, dim_model)`
- attention_mask (`torch.Tensor`):
- Avoid invalid areas to participate in the calculation of shape `(batch, seq_len, seq_len)`
- position_bias (`torch.Tensor`):
- Provides position information to attention mechanism of shape `(num_heads, seq_len, seq_len)`
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers.
- past_key_values (`Tuple[torch.Tensor, torch.Tensor])`, *optional*):
- Cached past key and value projection states
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- """
- hidden_states = self.self_att(
- hidden_states,
- attention_mask=attention_mask,
- position_bias=position_bias,
- output_attentions=output_attentions,
- past_key_values=past_key_values,
- use_cache=use_cache,
- )
-
- hidden_states, attn_weights, current_key_value = hidden_states
-
- hidden_states = self.ffn(hidden_states)
-
- return hidden_states, attn_weights, current_key_value
-
-
-class CpmAntEncoder(nn.Module):
- def __init__(self, config: CpmAntConfig):
- super().__init__()
- self.num_layers = config.num_hidden_layers
- self.layers = nn.ModuleList([CpmAntTransformerBlock(config) for ith in range(self.num_layers)])
-
- self.output_layernorm = CpmAntLayerNorm(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: torch.Tensor,
- position_bias: torch.Tensor,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- use_cache: Optional[bool] = None,
- ):
- """
- Args:
- hidden_states (`torch.Tensor`):
- Input to the layer of shape `(batch, seq_len, dim_model)`
- attention_mask (`torch.Tensor`):
- Avoid invalid areas to participate in the calculation of shape `(batch, seq_len, seq_len)`
- position_bias (`torch.Tensor`):
- Provides position information to attention mechanism of shape `(num_heads, seq_len, seq_len)`
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers.
- past_key_values (`Tuple[torch.Tensor, torch.Tensor])`, *optional*):
- Cached past key and value projection states
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- """
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- current_key_values = () if use_cache else None
-
- for i, layer in enumerate(self.layers):
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
- layer_outputs = layer(
- hidden_states,
- attention_mask,
- position_bias,
- output_attentions=output_attentions,
- past_key_values=past_key_values[i] if past_key_values else None,
- use_cache=use_cache,
- )
- hidden_states, attn_weights, current_key_value = layer_outputs
- if output_attentions:
- all_self_attns += (attn_weights,)
- if current_key_value is not None:
- current_key_values = current_key_values + (current_key_value,)
-
- hidden_states = self.output_layernorm(hidden_states)
-
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- return hidden_states, current_key_values, all_hidden_states, all_self_attns
-
-
-# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->CPMAnt
-class CpmAntIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-class CpmAntSegmentPositionEmbedding(nn.Module):
- def __init__(self, config: CpmAntConfig):
- super().__init__()
-
- self.num_heads = config.num_attention_heads
- self.num_buckets = config.position_bias_num_buckets
- self.max_distance = config.position_bias_max_distance
- self.num_segments = config.segment_types
-
- self.relative_attention_bias = nn.Parameter(
- torch.empty(
- config.segment_types * config.segment_types + config.position_bias_num_buckets,
- config.num_attention_heads,
- )
- )
-
- def forward(
- self,
- key_pos: torch.Tensor,
- query_pos: torch.Tensor,
- key_segment: torch.Tensor,
- query_segment: torch.Tensor,
- ):
- with torch.no_grad():
- batch = key_pos.size(0)
- keylen = key_pos.size(1)
- querylen = query_pos.size(1)
-
- if key_pos.size(0) != query_pos.size(0):
- raise AssertionError(
- f"key_pos.size(0) should be equal to query_pos.size(0), but got {key_pos.size(0)} and {query_pos.size(0)}!"
- )
- if keylen != key_segment.size(1) or querylen != query_segment.size(1):
- raise AssertionError(
- f"keylen should be equal to key_segment.size(1), but got {keylen} and {key_segment.size(1)}!"
- )
- if querylen != query_segment.size(1):
- raise AssertionError(
- f"querylen should be equal to query_segment.size(1), but got {querylen} and {query_segment.szie(1)}!"
- )
-
- key_pos = key_pos.view(batch, -1, keylen)
- query_pos = query_pos.view(batch, querylen, -1)
- key_segment = key_segment.view(batch, -1, keylen)
- query_segment = query_segment.view(batch, querylen, -1)
-
- relative_position_bucket = self._segment_relative_position_bucket(query_segment, key_segment)
- relative_position_bucket = relative_position_bucket + self.num_buckets
-
- # (batch, len_q, len_k)
- absolute_position_bucket = self._position_bucket(
- torch.arange(keylen, dtype=torch.int32, device=relative_position_bucket.device)[None, :]
- - torch.arange(querylen, dtype=torch.int32, device=relative_position_bucket.device)[:, None],
- num_buckets=self.num_buckets,
- max_distance=self.max_distance,
- )
- relative_position_bucket = torch.where(
- (key_segment == query_segment),
- absolute_position_bucket[None, :, :],
- relative_position_bucket,
- )
-
- # (batch, len_q, len_k, num_heads)
- embeds = F.embedding(relative_position_bucket, self.relative_attention_bias)
- # (batch, num_heads, len_q, len_k)
- embeds = embeds.permute(0, 3, 1, 2).contiguous()
- return embeds
-
- def _segment_relative_position_bucket(self, query_segment, key_segment):
- return query_segment * self.num_segments + key_segment
-
- def _position_bucket(self, relative_position, num_buckets=32, max_distance=128):
- relative_buckets = 0
- # always bidirectional in CPMAnt
- num_buckets //= 2
- relative_buckets = (relative_position > 0).to(torch.int32) * num_buckets
- relative_position = torch.abs(relative_position)
- max_exact = num_buckets // 2
- is_small = relative_position < max_exact
- relative_postion_if_large = max_exact + (
- torch.log(relative_position.float() / max_exact)
- / math.log(max_distance / max_exact)
- * (num_buckets - max_exact)
- ).to(torch.int32)
- relative_postion_if_large = torch.min(
- relative_postion_if_large,
- torch.full_like(relative_postion_if_large, num_buckets - 1),
- )
- relative_buckets += torch.where(is_small, relative_position.to(torch.int32), relative_postion_if_large)
- return relative_buckets
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->CPMAnt
-class CpmAntOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class CpmAntPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = CpmAntConfig
- base_model_prefix = "cpmant"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=self.config.init_std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.init_std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
- elif isinstance(module, CpmAntLayerNorm):
- module.weight.data.fill_(1.0)
- elif isinstance(module, CpmAntSegmentPositionEmbedding):
- module.relative_attention_bias.data.normal_(mean=0.0, std=self.config.init_std)
-
-
-CPMANT_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters
- config ([`~CpmAntConfig`]): Model configuration class with all the parameters of the
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CPMANT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`CPMAntTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
- blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare CPMAnt Model outputting raw hidden-states without any specific head on top.",
- CPMANT_START_DOCSTRING,
-)
-class CpmAntModel(CpmAntPreTrainedModel):
- def __init__(self, config: CpmAntConfig):
- super().__init__(config)
- self.encoder = CpmAntEncoder(config)
- self.segment_embedding = nn.Embedding(config.segment_types, config.hidden_size)
- self.input_embedding = nn.Embedding(
- config.vocab_size + config.prompt_types * config.prompt_length, config.hidden_size
- )
- self.position_bias = CpmAntSegmentPositionEmbedding(config)
- self.prompt_length = config.prompt_length
- self.vocab_size = config.vocab_size
-
- self.post_init()
-
- def get_input_embeddings(self):
- return self.input_embedding
-
- def set_input_embeddings(self, embeddings, **kwargs):
- self.input_embedding = embeddings
-
- def _prepare_attention_mask(self, input_ids, span, context, length):
- batch = input_ids.size(0)
- seqlen = input_ids.size(1)
- device = input_ids.device
- directional_mask_2d = torch.arange(seqlen, device=device) <= torch.arange(seqlen, device=device).view(-1, 1)
- attention_mask = context[:, None, :] | (
- context[:, :, None].logical_not() & directional_mask_2d.view(1, seqlen, seqlen)
- )
- attention_mask = attention_mask & (span[:, None, :] == span[:, :, None])
- # mask for left padding
- mask_1d = (
- torch.tensor(list(range(seqlen - self.prompt_length))[::-1], device=device)[None, :].repeat(batch, 1)
- < length[:, None]
- )
- mask_1d = torch.cat((torch.ones(batch, self.prompt_length, device=device).bool(), mask_1d), dim=1)
- attention_mask = mask_1d.view(batch, seqlen, 1) & mask_1d.view(batch, 1, seqlen) & attention_mask
- return attention_mask
-
- @add_start_docstrings_to_model_forward(CPMANT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- use_cache: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- **kwargs,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- use_cache = use_cache if use_cache is not None else self.config.use_cache
-
- # add prompts ahead
- if input_ids.dtype != torch.int32:
- input_ids = input_ids.to(torch.int32)
- dtype, device = input_ids.dtype, input_ids.device
- segment = torch.where(input_ids != 0, 2, 0).to(dtype=dtype, device=device)
- length = (segment != 0).sum(-1).to(dtype=dtype, device=device)
- input_ids = torch.cat(
- (
- torch.arange(
- self.prompt_length * 2 + self.vocab_size,
- self.prompt_length * 3 + self.vocab_size,
- dtype=dtype,
- device=device,
- ).repeat(input_ids.size(0), 1),
- input_ids,
- ),
- dim=1,
- )
- batch, seq_length = input_ids.size()
- segment = torch.cat((torch.zeros(batch, self.prompt_length, dtype=dtype, device=device), segment), dim=1)
- context = torch.full((batch, seq_length), 1, dtype=dtype, device=device)
- position = torch.arange(seq_length, dtype=dtype, device=device).repeat(batch, 1)
- span = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * self.encoder.num_layers)
- input_ids = input_ids.contiguous()
- hidden_states = self.input_embedding(input_ids)
- segment_states = self.segment_embedding(segment)
- hidden_states = hidden_states + segment_states
- else:
- past_length = past_key_values[0][0].size(-2)
- segment_states = self.segment_embedding(segment)
- hidden_states = self.input_embedding(input_ids) + segment_states[:, -1:, :]
-
- attention_mask = self._prepare_attention_mask(input_ids, span, context, length)
- position_bias = self.position_bias(position, position, segment, segment)
-
- attention_mask = attention_mask[:, past_length:, :]
- position_bias = position_bias[:, :, past_length:, :]
- hidden_states = hidden_states[:, past_length:, :]
-
- hidden_states, present_key_values, all_hidden_states, all_attentions = self.encoder(
- hidden_states,
- attention_mask,
- position_bias,
- output_attentions,
- output_hidden_states,
- past_key_values,
- use_cache,
- )
-
- if past_length == 0:
- hidden_states = hidden_states[:, self.prompt_length :, :]
- # drop the prompt
- if all_attentions is not None:
- new_attentions = ()
- for attention in all_attentions:
- new_attentions += (attention[:, :, self.prompt_length :, self.prompt_length :],)
- all_attentions = new_attentions
- if all_hidden_states is not None:
- new_hidden_states = ()
- for hidden_state in all_hidden_states:
- new_hidden_states += (hidden_state[:, self.prompt_length :, :],)
- all_hidden_states = new_hidden_states
-
- if not return_dict:
- return tuple(
- v for v in [hidden_states, present_key_values, all_hidden_states, all_attentions] if v is not None
- )
-
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=present_key_values,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The CPMAnt Model with a language modeling head on top (linear layer with weights tied to the input embeddings).
- """,
- CPMANT_START_DOCSTRING,
-)
-class CpmAntForCausalLM(CpmAntPreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config: CpmAntConfig):
- super().__init__(config)
- self.cpmant = CpmAntModel(config)
-
- # lm_head.weight is tied to cpmant.input_embedding.weight
- self.lm_head = nn.Linear(
- config.hidden_size, config.vocab_size + config.prompt_types * config.prompt_length, bias=False
- )
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CPMANT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- labels: Optional[torch.Tensor] = None,
- return_dict: Optional[bool] = None,
- attention_mask: Optional[torch.Tensor] = None, # dummy parameter for text-generation pipeline
- **kwargs,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- Args:
- input_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`CPMAntTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
- cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers.
- labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- CPMAnt will process attention mask automatically, this parameter is a dummy parameter for
- text-generation pipeline.
-
- Example:
-
- Text Generation with CpmAntForCausalLM.
- ```python
- >>> from transformers import CPMAntTokenizer, CpmAntForCausalLM
-
- >>> texts = "今天天气不错，"
- >>> model = CpmAntForCausalLM.from_pretrained("openbmb/cpm-ant-10b")
- >>> tokenizer = CPMAntTokenizer.from_pretrained("openbmb/cpm-ant-10b")
- >>> input_ids = tokenizer(texts, return_tensors="pt")
- >>> outputs = model.generate(**input_ids)
- >>> output_texts = tokenizer.batch_decode(outputs)
- >>> print(output_texts)
- ['今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的']
- ```
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- model_output = self.cpmant(
- input_ids, output_attentions, output_hidden_states, past_key_values, use_cache, return_dict
- )
- hidden_states = model_output.last_hidden_state if return_dict else model_output[0]
-
- logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- loss_func = CrossEntropyLoss()
- loss = loss_func(logits.view(-1, logits.size(-1)), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + model_output[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=model_output.past_key_values,
- hidden_states=model_output.hidden_states,
- attentions=model_output.attentions,
- )
-
- def get_input_embeddings(self):
- return self.cpmant.input_embedding
-
- def set_input_embeddings(self, embeddings):
- self.cpmant.input_embedding = embeddings
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, **kwargs):
- input_ids = input_ids.int()
- # save the memory usage of dummy attention mask
- if "attention_mask" in kwargs:
- kwargs["attention_mask"] = torch.zeros(1, 1)
-
- return {
- "input_ids": input_ids,
- "use_cache": kwargs["use_cache"],
- "past_key_values": kwargs.get("past_key_values", None),
- }
-
- def _reorder_cache(self, past_key_values, beam_idx):
- past_key_values = [list(each) if each is not None else each for each in past_key_values]
- for key_value_layer in past_key_values:
- key_value_layer[0] = key_value_layer[0][beam_idx]
- key_value_layer[1] = key_value_layer[1][beam_idx]
- return past_key_values
diff --git a/transformers/models/cpmant/tokenization_cpmant.py b/transformers/models/cpmant/tokenization_cpmant.py
deleted file mode 100644
index a5e66c7679c728320c00ddd9a60696477232f304..0000000000000000000000000000000000000000
--- a/transformers/models/cpmant/tokenization_cpmant.py
+++ /dev/null
@@ -1,266 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for CPMAnt."""
-import collections
-import os
-from typing import List, Optional, Tuple
-
-from transformers.utils import is_jieba_available, requires_backends
-
-
-if is_jieba_available():
- import jieba
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
-
-
-def load_vocab(vocab_file):
- """Loads a vocabulary file into a dictionary."""
- vocab = collections.OrderedDict()
- with open(vocab_file, "r", encoding="utf-8") as reader:
- tokens = reader.readlines()
- for index, token in enumerate(tokens):
- token = token.rstrip("\n")
- vocab[token] = index
- return vocab
-
-
-class WordpieceTokenizer(object):
- def __init__(self, vocab, unk_token="", max_input_chars_per_word=200):
- self.vocab = vocab
- self.unk_token = unk_token
- self.max_input_chars_per_word = max_input_chars_per_word
-
- def tokenize(self, token):
- chars = list(token)
- if len(chars) > self.max_input_chars_per_word:
- return [self.unk_token]
-
- start = 0
- sub_tokens = []
- while start < len(chars):
- end = len(chars)
- cur_substr = None
- while start < end:
- substr = "".join(chars[start:end])
- if substr in self.vocab:
- cur_substr = substr
- break
- end -= 1
- if cur_substr is None:
- sub_tokens.append(self.unk_token)
- start += 1
- else:
- sub_tokens.append(cur_substr)
- start = end
-
- return sub_tokens
-
-
-class CpmAntTokenizer(PreTrainedTokenizer):
- """
- Construct a CPMAnt tokenizer. Based on byte-level Byte-Pair-Encoding.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- bod_token (`str`, *optional*, defaults to `""`):
- The beginning of document token.
- eod_token (`str`, *optional*, defaults to `""`):
- The end of document token.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sequence token.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token.
- line_token (`str`, *optional*, defaults to `""`):
- The line token.
- space_token (`str`, *optional*, defaults to `""`):
- The space token.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
- add_prefix_space = False
-
- def __init__(
- self,
- vocab_file,
- bod_token="",
- eod_token="",
- bos_token="",
- eos_token="",
- pad_token="",
- unk_token="",
- line_token="",
- space_token="",
- padding_side="left",
- **kwargs,
- ):
- requires_backends(self, ["jieba"])
- self.bod_token = bod_token
- self.eod_token = eod_token
- self.encoder = load_vocab(vocab_file)
- self.encoder[" "] = self.encoder[space_token]
- self.encoder["\n"] = self.encoder[line_token]
-
- del self.encoder[space_token]
- del self.encoder[line_token]
-
- self.encoder = collections.OrderedDict(sorted(self.encoder.items(), key=lambda x: x[1]))
- self.decoder = {v: k for k, v in self.encoder.items()}
-
- self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.encoder, unk_token=unk_token)
-
- super().__init__(
- bod_token=bod_token,
- eod_token=eod_token,
- bos_token=bos_token,
- eos_token=eos_token,
- pad_token=pad_token,
- unk_token=unk_token,
- line_token=line_token,
- space_token=space_token,
- padding_side=padding_side,
- **kwargs,
- )
-
- @property
- def bod_token_id(self):
- return self.encoder[self.bod_token]
-
- @property
- def eod_token_id(self):
- return self.encoder[self.eod_token]
-
- @property
- def newline_id(self):
- return self.encoder["\n"]
-
- @property
- def vocab_size(self) -> int:
- return len(self.encoder)
-
- def get_vocab(self):
- return dict(self.encoder, **self.added_tokens_encoder)
-
- def _tokenize(self, text):
- """Tokenize a string."""
- output_tokens = []
- for x in jieba.cut(text, cut_all=False):
- output_tokens.extend(self.wordpiece_tokenizer.tokenize(x))
- return output_tokens
-
- def _decode(self, token_ids, **kwargs):
- """Decode ids into a string."""
- token_ids = [i for i in token_ids if i >= 0]
- token_ids = [
- x for x in token_ids if x != self.pad_token_id and x != self.eos_token_id and x != self.bos_token_id
- ]
- return super()._decode(token_ids, **kwargs)
-
- def check(self, token):
- return token in self.encoder
-
- def convert_tokens_to_string(self, tokens: List[str]) -> str:
- return "".join(tokens)
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index, self.unk_token)
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- else:
- vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
- index = 0
- if " " in self.encoder:
- self.encoder[""] = self.encoder[" "]
- del self.encoder[" "]
- if "\n" in self.encoder:
- self.encoder[""] = self.encoder["\n"]
- del self.encoder["\n"]
- self.encoder = collections.OrderedDict(sorted(self.encoder.items(), key=lambda x: x[1]))
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token, token_index in self.encoder.items():
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(token + "\n")
- index += 1
- return (vocab_file,)
-
- def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: List[int] = None) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A CPMAnt sequence has the following format:
-
- - single sequence: `[BOS] Sequence`.
-
- Args:
- token_ids_0 (`List[int]`): The first tokenized sequence that special tokens will be added.
- token_ids_1 (`List[int]`): The optional second tokenized sequence that special tokens will be added.
-
- Returns:
- `List[int]`: The model input with special tokens.
- """
- if token_ids_1 is None:
- return [self.bos_token_id] + token_ids_0
- return [self.bos_token_id] + token_ids_0 + [self.bos_token_id] + token_ids_1
-
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`): List of IDs.
- token_ids_1 (`List[int]`, *optional*): Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1))
- return [1] + ([0] * len(token_ids_0))
diff --git a/transformers/models/ctrl/__init__.py b/transformers/models/ctrl/__init__.py
deleted file mode 100644
index 7463117bfbc623a2c96019e9a7a3e864c11934db..0000000000000000000000000000000000000000
--- a/transformers/models/ctrl/__init__.py
+++ /dev/null
@@ -1,89 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available
-
-
-_import_structure = {
- "configuration_ctrl": ["CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP", "CTRLConfig"],
- "tokenization_ctrl": ["CTRLTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_ctrl"] = [
- "CTRL_PRETRAINED_MODEL_ARCHIVE_LIST",
- "CTRLForSequenceClassification",
- "CTRLLMHeadModel",
- "CTRLModel",
- "CTRLPreTrainedModel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_ctrl"] = [
- "TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFCTRLForSequenceClassification",
- "TFCTRLLMHeadModel",
- "TFCTRLModel",
- "TFCTRLPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_ctrl import CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP, CTRLConfig
- from .tokenization_ctrl import CTRLTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_ctrl import (
- CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
- CTRLForSequenceClassification,
- CTRLLMHeadModel,
- CTRLModel,
- CTRLPreTrainedModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_ctrl import (
- TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFCTRLForSequenceClassification,
- TFCTRLLMHeadModel,
- TFCTRLModel,
- TFCTRLPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/ctrl/__pycache__/__init__.cpython-310.pyc b/transformers/models/ctrl/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 8bfc639bc7b95f247efd43107748f5d9a7d85a2a..0000000000000000000000000000000000000000
Binary files a/transformers/models/ctrl/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ctrl/__pycache__/configuration_ctrl.cpython-310.pyc b/transformers/models/ctrl/__pycache__/configuration_ctrl.cpython-310.pyc
deleted file mode 100644
index 6a72bf0907800f405d23dad4dcf6eb73e6faf281..0000000000000000000000000000000000000000
Binary files a/transformers/models/ctrl/__pycache__/configuration_ctrl.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ctrl/__pycache__/modeling_ctrl.cpython-310.pyc b/transformers/models/ctrl/__pycache__/modeling_ctrl.cpython-310.pyc
deleted file mode 100644
index 7ddb7966403ce5f0af1e4147553618f26c4f00bd..0000000000000000000000000000000000000000
Binary files a/transformers/models/ctrl/__pycache__/modeling_ctrl.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ctrl/__pycache__/modeling_tf_ctrl.cpython-310.pyc b/transformers/models/ctrl/__pycache__/modeling_tf_ctrl.cpython-310.pyc
deleted file mode 100644
index ec533c2cd41070191580ce3d596895d6f17c8f1d..0000000000000000000000000000000000000000
Binary files a/transformers/models/ctrl/__pycache__/modeling_tf_ctrl.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ctrl/__pycache__/tokenization_ctrl.cpython-310.pyc b/transformers/models/ctrl/__pycache__/tokenization_ctrl.cpython-310.pyc
deleted file mode 100644
index e862f57a3d7a9095b9d65a13fd842d4ffb43025b..0000000000000000000000000000000000000000
Binary files a/transformers/models/ctrl/__pycache__/tokenization_ctrl.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ctrl/configuration_ctrl.py b/transformers/models/ctrl/configuration_ctrl.py
deleted file mode 100644
index 0c5a68bf6fcbdc6f3e282ccbba283edbba89b1df..0000000000000000000000000000000000000000
--- a/transformers/models/ctrl/configuration_ctrl.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Salesforce and HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Salesforce CTRL configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class CTRLConfig(PretrainedConfig):
- """
- This is the configuration class to store the configuration of a [`CTRLModel`] or a [`TFCTRLModel`]. It is used to
- instantiate a CTRL model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the
- [Salesforce/ctrl](https://huggingface.co/Salesforce/ctrl) architecture from SalesForce.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 246534):
- Vocabulary size of the CTRL model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`CTRLModel`] or [`TFCTRLModel`].
- n_positions (`int`, *optional*, defaults to 256):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- n_embd (`int`, *optional*, defaults to 1280):
- Dimensionality of the embeddings and hidden states.
- dff (`int`, *optional*, defaults to 8192):
- Dimensionality of the inner dimension of the feed forward networks (FFN).
- n_layer (`int`, *optional*, defaults to 48):
- Number of hidden layers in the Transformer encoder.
- n_head (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- resid_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- embd_pdrop (`int`, *optional*, defaults to 0.1):
- The dropout ratio for the embeddings.
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-06):
- The epsilon to use in the layer normalization layers
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
-
-
- Examples:
-
- ```python
- >>> from transformers import CTRLConfig, CTRLModel
-
- >>> # Initializing a CTRL configuration
- >>> configuration = CTRLConfig()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = CTRLModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "ctrl"
- keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {
- "max_position_embeddings": "n_positions",
- "hidden_size": "n_embd",
- "num_attention_heads": "n_head",
- "num_hidden_layers": "n_layer",
- }
-
- def __init__(
- self,
- vocab_size=246534,
- n_positions=256,
- n_embd=1280,
- dff=8192,
- n_layer=48,
- n_head=16,
- resid_pdrop=0.1,
- embd_pdrop=0.1,
- layer_norm_epsilon=1e-6,
- initializer_range=0.02,
- use_cache=True,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.n_positions = n_positions
- self.n_embd = n_embd
- self.n_layer = n_layer
- self.n_head = n_head
- self.dff = dff
- self.resid_pdrop = resid_pdrop
- self.embd_pdrop = embd_pdrop
- self.layer_norm_epsilon = layer_norm_epsilon
- self.initializer_range = initializer_range
-
- self.use_cache = use_cache
-
- super().__init__(**kwargs)
diff --git a/transformers/models/ctrl/modeling_ctrl.py b/transformers/models/ctrl/modeling_ctrl.py
deleted file mode 100644
index 7534a0e50c9a2335a543d37558e5495705b681ba..0000000000000000000000000000000000000000
--- a/transformers/models/ctrl/modeling_ctrl.py
+++ /dev/null
@@ -1,841 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Salesforce and HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch CTRL model."""
-
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import torch
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutput
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import Conv1D, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from .configuration_ctrl import CTRLConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "CTRLConfig"
-
-
-from ..deprecated._archive_maps import CTRL_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-def angle_defn(pos, i, d_model_size):
- angle_rates = 1 / torch.pow(10000, (2 * (i // 2)) / d_model_size)
- return pos * angle_rates
-
-
-def positional_encoding(position, d_model_size, dtype):
- # create the sinusoidal pattern for the positional encoding
- angle_rads = angle_defn(
- torch.arange(position, dtype=torch.int64).to(dtype).unsqueeze(1),
- torch.arange(d_model_size, dtype=torch.int64).to(dtype).unsqueeze(0),
- d_model_size,
- )
-
- sines = torch.sin(angle_rads[:, 0::2])
- cosines = torch.cos(angle_rads[:, 1::2])
-
- pos_encoding = torch.cat([sines, cosines], dim=-1)
- return pos_encoding
-
-
-def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=None):
- # calculate attention
- matmul_qk = torch.matmul(q, k.permute(0, 1, 3, 2))
-
- dk = k.shape[-1]
- scaled_attention_logits = matmul_qk / np.sqrt(dk)
-
- if mask is not None:
- nd, ns = scaled_attention_logits.size(-2), scaled_attention_logits.size(-1)
- scaled_attention_logits += mask[ns - nd : ns, :ns] * -1e4
-
- if attention_mask is not None:
- # Apply the attention mask
- scaled_attention_logits = scaled_attention_logits + attention_mask
-
- attention_weights = torch.softmax(scaled_attention_logits, dim=-1)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_weights = attention_weights * head_mask
-
- output = torch.matmul(attention_weights, v)
-
- return output, attention_weights
-
-
-class MultiHeadAttention(nn.Module):
- def __init__(self, d_model_size, num_heads):
- super().__init__()
- self.num_heads = num_heads
- self.d_model_size = d_model_size
-
- self.depth = int(d_model_size / self.num_heads)
-
- self.Wq = nn.Linear(d_model_size, d_model_size)
- self.Wk = nn.Linear(d_model_size, d_model_size)
- self.Wv = nn.Linear(d_model_size, d_model_size)
-
- self.dense = nn.Linear(d_model_size, d_model_size)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- attention_head_size = self.d_model_size // self.num_heads
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(heads, self.num_heads, attention_head_size, self.pruned_heads)
-
- # Prune linear layers
- self.Wq = prune_linear_layer(self.Wq, index)
- self.Wk = prune_linear_layer(self.Wk, index)
- self.Wv = prune_linear_layer(self.Wv, index)
- self.dense = prune_linear_layer(self.dense, index, dim=1)
-
- # Update hyper params
- self.num_heads = self.num_heads - len(heads)
- self.d_model_size = attention_head_size * self.num_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def split_into_heads(self, x, batch_size):
- x = x.reshape(batch_size, -1, self.num_heads, self.depth)
- return x.permute([0, 2, 1, 3])
-
- def forward(
- self,
- v,
- k,
- q,
- mask,
- layer_past=None,
- attention_mask=None,
- head_mask=None,
- use_cache=False,
- output_attentions=False,
- ):
- batch_size = q.shape[0]
-
- q = self.Wq(q)
- k = self.Wk(k)
- v = self.Wv(v)
-
- q = self.split_into_heads(q, batch_size)
- k = self.split_into_heads(k, batch_size)
- v = self.split_into_heads(v, batch_size)
- if layer_past is not None:
- past_key, past_value = layer_past[0], layer_past[1]
- k = torch.cat((past_key, k), dim=-2)
- v = torch.cat((past_value, v), dim=-2)
-
- if use_cache is True:
- present = torch.stack((k, v))
- else:
- present = (None,)
-
- output = scaled_dot_product_attention(q, k, v, mask, attention_mask, head_mask)
- scaled_attention = output[0].permute([0, 2, 1, 3])
- attn = output[1]
- original_size_attention = scaled_attention.reshape(batch_size, -1, self.d_model_size)
- output = self.dense(original_size_attention)
-
- outputs = (output, present)
- if output_attentions:
- outputs = outputs + (attn,)
- return outputs
-
-
-def point_wise_feed_forward_network(d_model_size, dff):
- return nn.Sequential(nn.Linear(d_model_size, dff), nn.ReLU(), nn.Linear(dff, d_model_size))
-
-
-class EncoderLayer(nn.Module):
- def __init__(self, d_model_size, num_heads, dff, rate=0.1):
- super().__init__()
-
- self.multi_head_attention = MultiHeadAttention(d_model_size, num_heads)
- self.ffn = point_wise_feed_forward_network(d_model_size, dff)
-
- self.layernorm1 = nn.LayerNorm(d_model_size, eps=1e-6)
- self.layernorm2 = nn.LayerNorm(d_model_size, eps=1e-6)
-
- self.dropout1 = nn.Dropout(rate)
- self.dropout2 = nn.Dropout(rate)
-
- def forward(
- self, x, mask, layer_past=None, attention_mask=None, head_mask=None, use_cache=False, output_attentions=False
- ):
- normed = self.layernorm1(x)
- attn_outputs = self.multi_head_attention(
- normed,
- normed,
- normed,
- mask,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0]
- attn_output = self.dropout1(attn_output)
- out1 = x + attn_output
-
- out2 = self.layernorm2(out1)
- ffn_output = self.ffn(out2)
- ffn_output = self.dropout2(ffn_output)
- out2 = out1 + ffn_output
-
- outputs = (out2,) + attn_outputs[1:]
- return outputs
-
-
-class CTRLPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = CTRLConfig
- base_model_prefix = "transformer"
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, (nn.Linear, Conv1D)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-CTRL_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`CTRLConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CTRL_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0].shape[-2]`
- (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
-
- If `past_key_values` is used, only input IDs that do not have their past calculated should be passed as
- `input_ids`.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`Tuple[Tuple[torch.FloatTensor]]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
- their past given to this model should not be passed as input ids as they have already been computed.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare CTRL Model transformer outputting raw hidden-states without any specific head on top.",
- CTRL_START_DOCSTRING,
-)
-class CTRLModel(CTRLPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.d_model_size = config.n_embd
- self.num_layers = config.n_layer
-
- self.pos_encoding = positional_encoding(config.n_positions, self.d_model_size, torch.float)
-
- self.w = nn.Embedding(config.vocab_size, config.n_embd)
-
- self.dropout = nn.Dropout(config.embd_pdrop)
- self.h = nn.ModuleList(
- [EncoderLayer(config.n_embd, config.n_head, config.dff, config.resid_pdrop) for _ in range(config.n_layer)]
- )
- self.layernorm = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.w
-
- def set_input_embeddings(self, new_embeddings):
- self.w = new_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
- """
- for layer, heads in heads_to_prune.items():
- self.h[layer].multi_head_attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BaseModelOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:
- r"""
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, CTRLModel
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")
- >>> model = CTRLModel.from_pretrained("Salesforce/ctrl")
-
- >>> # CTRL was trained with control codes as the first token
- >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")
- >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()
-
- >>> outputs = model(**inputs)
-
- >>> last_hidden_states = outputs.last_hidden_state
- >>> list(last_hidden_states.shape)
- [1, 5, 1280]
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- input_ids = input_ids.view(-1, input_shape[-1])
- batch_size = input_ids.shape[0]
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- batch_size = inputs_embeds.shape[0]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * len(self.h))
- else:
- past_length = past_key_values[0][0].size(-2)
- if position_ids is None:
- position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0)
-
- # Attention mask.
- if attention_mask is not None:
- if batch_size <= 0:
- raise ValueError("batch_size has to be defined and > 0")
- attention_mask = attention_mask.view(batch_size, -1)
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and the dtype's smallest value for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
- attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
- # Prepare head mask if needed
- head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-
- if token_type_ids is not None:
- token_type_ids = token_type_ids.view(-1, input_shape[-1])
- token_type_embeds = self.w(token_type_ids)
- token_type_embeds *= np.sqrt(self.d_model_size)
- else:
- token_type_embeds = 0
-
- if inputs_embeds is None:
- inputs_embeds = self.w(input_ids)
- # inputs_embeds = embedded.unsqueeze(0) if len(input_ids.shape)<2 else embedded
- seq_len = input_shape[-1]
- mask = torch.triu(torch.ones(seq_len + past_length, seq_len + past_length), 1).to(device)
-
- inputs_embeds *= np.sqrt(self.d_model_size)
-
- # `self.pos_encoding` won't be sent to the correct device along the model, so we do it manually.
- self.pos_encoding = self.pos_encoding.to(device)
- pos_embeds = self.pos_encoding[position_ids, :]
-
- hidden_states = inputs_embeds + pos_embeds + token_type_embeds
-
- hidden_states = self.dropout(hidden_states)
-
- presents = () if use_cache else None
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
- for i, (h, layer_past) in enumerate(zip(self.h, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
- outputs = h(
- hidden_states,
- mask,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask[i],
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- hidden_states, present = outputs[:2]
- if use_cache is True:
- presents = presents + (present,)
-
- if output_attentions:
- all_attentions += (outputs[2],)
-
- hidden_states = self.layernorm(hidden_states)
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)
-
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The CTRL Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- CTRL_START_DOCSTRING,
-)
-class CTRLLMHeadModel(CTRLPreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.transformer = CTRLModel(config)
- self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=True)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, use_cache=None, **kwargs):
- # only last tokens for inputs_ids if past is defined in kwargs
- if past_key_values is not None:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
-
- return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": use_cache}
-
- @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
-
- Returns:
-
- Example:
-
- ```python
- >>> import torch
- >>> from transformers import AutoTokenizer, CTRLLMHeadModel
-
- >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")
- >>> model = CTRLLMHeadModel.from_pretrained("Salesforce/ctrl")
-
- >>> # CTRL was trained with control codes as the first token
- >>> inputs = tokenizer("Wikipedia The llama is", return_tensors="pt")
- >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()
-
- >>> sequence_ids = model.generate(inputs["input_ids"])
- >>> sequences = tokenizer.batch_decode(sequence_ids)
- >>> sequences
- ['Wikipedia The llama is a member of the family Bovidae. It is native to the Andes of Peru,']
-
- >>> outputs = model(**inputs, labels=inputs["input_ids"])
- >>> round(outputs.loss.item(), 2)
- 9.21
-
- >>> list(outputs.logits.shape)
- [1, 5, 246534]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = transformer_outputs[0]
-
- lm_logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
- if not return_dict:
- output = (lm_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=lm_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- @staticmethod
- def _reorder_cache(
- past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
- ) -> Tuple[Tuple[torch.Tensor]]:
- """
- This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
- [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
- beam_idx at every generation step.
- """
- return tuple(
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
- for layer_past in past_key_values
- )
-
-
-@add_start_docstrings(
- """
- The CTRL Model transformer with a sequence classification head on top (linear layer).
- [`CTRLForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT-2) do. Since it does classification on the last token, it requires to know the position of the last
- token. If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in
- each row. If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot
- guess the padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last
- value in each row of the batch).
- """,
- CTRL_START_DOCSTRING,
-)
-class CTRLForSequenceClassification(CTRLPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = CTRLModel(config)
- self.classifier = nn.Linear(config.n_embd, self.num_labels, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=SequenceClassifierOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-
- Returns:
-
- Example of single-label classification:
-
- ```python
- >>> import torch
- >>> from transformers import AutoTokenizer, CTRLForSequenceClassification
-
- >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")
- >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl")
-
- >>> # CTRL was trained with control codes as the first token
- >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")
- >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()
-
- >>> with torch.no_grad():
- ... logits = model(**inputs).logits
-
- >>> predicted_class_id = logits.argmax().item()
- >>> model.config.id2label[predicted_class_id]
- 'LABEL_0'
- ```
-
- ```python
- >>> import torch
-
- >>> torch.manual_seed(42) # doctest: +IGNORE_RESULT
- >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
- >>> num_labels = len(model.config.id2label)
- >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl", num_labels=num_labels)
-
- >>> labels = torch.tensor(1)
- >>> loss = model(**inputs, labels=labels).loss
- >>> round(loss.item(), 2)
- 0.93
- ```
-
- Example of multi-label classification:
-
- ```python
- >>> import torch
- >>> from transformers import AutoTokenizer, CTRLForSequenceClassification
-
- >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")
- >>> model = CTRLForSequenceClassification.from_pretrained(
- ... "Salesforce/ctrl", problem_type="multi_label_classification"
- ... )
-
- >>> # CTRL was trained with control codes as the first token
- >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")
- >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()
-
- >>> with torch.no_grad():
- ... logits = model(**inputs).logits
-
- >>> predicted_class_id = logits.argmax().item()
- >>> model.config.id2label[predicted_class_id]
- 'LABEL_0'
- ```
-
- ```python
- >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
- >>> num_labels = len(model.config.id2label)
- >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl", num_labels=num_labels)
-
- >>> num_labels = len(model.config.id2label)
- >>> labels = torch.nn.functional.one_hot(torch.tensor([predicted_class_id]), num_classes=num_labels).to(
- ... torch.float
- ... )
- >>> loss = model(**inputs, labels=labels).loss
- >>> loss.backward() # doctest: +IGNORE_RESULT
- ```"""
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = transformer_outputs[0]
- logits = self.classifier(hidden_states)
-
- if input_ids is not None:
- batch_size, sequence_length = input_ids.shape[:2]
- else:
- batch_size, sequence_length = inputs_embeds.shape[:2]
-
- if self.config.pad_token_id is None and batch_size != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
-
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
-
- pooled_logits = logits[range(batch_size), sequence_lengths]
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=pooled_logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
diff --git a/transformers/models/ctrl/modeling_tf_ctrl.py b/transformers/models/ctrl/modeling_tf_ctrl.py
deleted file mode 100644
index 6569b9e7d7b788ce93d1f9186b10eaebcfbf3894..0000000000000000000000000000000000000000
--- a/transformers/models/ctrl/modeling_tf_ctrl.py
+++ /dev/null
@@ -1,931 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Salesforce and HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 CTRL model."""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...modeling_tf_outputs import TFBaseModelOutputWithPast, TFCausalLMOutputWithPast, TFSequenceClassifierOutput
-from ...modeling_tf_utils import (
- TFCausalLanguageModelingLoss,
- TFModelInputType,
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_ctrl import CTRLConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "Salesforce/ctrl"
-_CONFIG_FOR_DOC = "CTRLConfig"
-
-
-from ..deprecated._archive_maps import TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-def angle_defn(pos, i, d_model_size):
- angle_rates = 1 / np.power(10000, (2 * (i // 2)) / d_model_size)
- return pos * angle_rates
-
-
-def positional_encoding(position, d_model_size):
- # create the sinusoidal pattern for the positional encoding
- angle_rads = angle_defn(np.arange(position)[:, np.newaxis], np.arange(d_model_size)[np.newaxis, :], d_model_size)
-
- sines = np.sin(angle_rads[:, 0::2])
- cosines = np.cos(angle_rads[:, 1::2])
- pos_encoding = tf.convert_to_tensor(np.concatenate([sines, cosines], axis=-1))
-
- return pos_encoding
-
-
-def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=None):
- # calculate attention
- matmul_qk = tf.matmul(q, k, transpose_b=True)
-
- dk = tf.cast(shape_list(k)[-1], dtype=matmul_qk.dtype)
- scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
-
- if mask is not None:
- scaled_attention_logits += tf.cast(mask * -1e4, dtype=scaled_attention_logits.dtype)
-
- if attention_mask is not None:
- # Apply the attention mask
- attention_mask = tf.cast(attention_mask, dtype=scaled_attention_logits.dtype)
- scaled_attention_logits = scaled_attention_logits + attention_mask
-
- attention_weights = stable_softmax(scaled_attention_logits, axis=-1)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_weights = attention_weights * head_mask
-
- output = tf.matmul(attention_weights, v)
-
- return output, attention_weights
-
-
-class TFMultiHeadAttention(keras.layers.Layer):
- def __init__(self, d_model_size, num_heads, output_attentions=False, **kwargs):
- super().__init__(**kwargs)
- self.num_heads = num_heads
- self.d_model_size = d_model_size
- self.output_attentions = output_attentions
-
- self.depth = int(d_model_size / self.num_heads)
-
- self.Wq = keras.layers.Dense(d_model_size, name="Wq")
- self.Wk = keras.layers.Dense(d_model_size, name="Wk")
- self.Wv = keras.layers.Dense(d_model_size, name="Wv")
-
- self.dense = keras.layers.Dense(d_model_size, name="dense")
-
- def split_into_heads(self, x, batch_size):
- x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
- return tf.transpose(x, perm=[0, 2, 1, 3])
-
- def call(self, v, k, q, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions, training=False):
- batch_size = shape_list(q)[0]
-
- q = self.Wq(q)
- k = self.Wk(k)
- v = self.Wv(v)
-
- q = self.split_into_heads(q, batch_size)
- k = self.split_into_heads(k, batch_size)
- v = self.split_into_heads(v, batch_size)
-
- if layer_past is not None:
- past_key, past_value = tf.unstack(layer_past, axis=0)
- k = tf.concat((past_key, k), axis=-2)
- v = tf.concat((past_value, v), axis=-2)
-
- if use_cache:
- present = tf.stack((k, v), axis=0)
- else:
- present = (None,)
-
- output = scaled_dot_product_attention(q, k, v, mask, attention_mask, head_mask)
- scaled_attention = tf.transpose(output[0], perm=[0, 2, 1, 3])
- attn = output[1]
- original_size_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model_size))
- output = self.dense(original_size_attention)
- outputs = (output, present)
-
- if output_attentions:
- outputs = outputs + (attn,)
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "Wq", None) is not None:
- with tf.name_scope(self.Wq.name):
- self.Wq.build([None, None, self.d_model_size])
- if getattr(self, "Wk", None) is not None:
- with tf.name_scope(self.Wk.name):
- self.Wk.build([None, None, self.d_model_size])
- if getattr(self, "Wv", None) is not None:
- with tf.name_scope(self.Wv.name):
- self.Wv.build([None, None, self.d_model_size])
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.d_model_size])
-
-
-class TFPointWiseFeedForwardLayer(keras.layers.Layer):
- def __init__(self, d_model_size, dff, **kwargs):
- super().__init__(**kwargs)
-
- self.dense_0 = keras.layers.Dense(dff, activation="relu", name="0")
- self.dense_2 = keras.layers.Dense(d_model_size, name="2")
- self.d_model_size = d_model_size
- self.dff = dff
-
- def call(self, inputs, trainable=False):
- dense_0_output = self.dense_0(inputs)
- dense_2_output = self.dense_2(dense_0_output)
-
- return dense_2_output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense_0", None) is not None:
- with tf.name_scope(self.dense_0.name):
- self.dense_0.build([None, None, self.d_model_size])
- if getattr(self, "dense_2", None) is not None:
- with tf.name_scope(self.dense_2.name):
- self.dense_2.build([None, None, self.dff])
-
-
-class TFEncoderLayer(keras.layers.Layer):
- def __init__(
- self, d_model_size, num_heads, dff, rate=0.1, layer_norm_epsilon=1e-6, output_attentions=False, **kwargs
- ):
- super().__init__(**kwargs)
-
- self.output_attentions = output_attentions
-
- self.multi_head_attention = TFMultiHeadAttention(
- d_model_size, num_heads, output_attentions=self.output_attentions, name="multi_head_attention"
- )
- self.ffn = TFPointWiseFeedForwardLayer(d_model_size, dff, name="ffn")
-
- self.layernorm1 = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layernorm1")
- self.layernorm2 = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layernorm2")
-
- self.dropout1 = keras.layers.Dropout(rate)
- self.dropout2 = keras.layers.Dropout(rate)
- self.d_model_size = d_model_size
-
- def call(self, x, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions, training=False):
- normed = self.layernorm1(x)
- attn_outputs = self.multi_head_attention(
- normed,
- normed,
- normed,
- mask,
- layer_past,
- attention_mask,
- head_mask,
- use_cache,
- output_attentions,
- training=training,
- )
- attn_output = attn_outputs[0]
- attn_output = self.dropout1(attn_output, training=training)
- out1 = x + attn_output
-
- out2 = self.layernorm2(out1)
- ffn_output = self.ffn(out2)
- ffn_output = self.dropout2(ffn_output, training=training)
- out2 = out1 + ffn_output
-
- outputs = (out2,) + attn_outputs[1:]
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "multi_head_attention", None) is not None:
- with tf.name_scope(self.multi_head_attention.name):
- self.multi_head_attention.build(None)
- if getattr(self, "ffn", None) is not None:
- with tf.name_scope(self.ffn.name):
- self.ffn.build(None)
- if getattr(self, "layernorm1", None) is not None:
- with tf.name_scope(self.layernorm1.name):
- self.layernorm1.build([None, None, self.d_model_size])
- if getattr(self, "layernorm2", None) is not None:
- with tf.name_scope(self.layernorm2.name):
- self.layernorm2.build([None, None, self.d_model_size])
-
-
-@keras_serializable
-class TFCTRLMainLayer(keras.layers.Layer):
- config_class = CTRLConfig
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.output_hidden_states = config.output_hidden_states
- self.output_attentions = config.output_attentions
- self.use_cache = config.use_cache
- self.return_dict = config.use_return_dict
-
- self.d_model_size = config.n_embd
- self.num_layers = config.n_layer
-
- self.pos_encoding = positional_encoding(config.n_positions, self.d_model_size)
-
- self.w = keras.layers.Embedding(
- input_dim=config.vocab_size,
- output_dim=config.n_embd,
- embeddings_initializer=get_initializer(config.initializer_range),
- name="w",
- )
-
- self.dropout = keras.layers.Dropout(config.embd_pdrop)
- self.h = [
- TFEncoderLayer(
- config.n_embd,
- config.n_head,
- config.dff,
- config.resid_pdrop,
- config.layer_norm_epsilon,
- self.output_attentions,
- name=f"h_._{i}",
- )
- for i in range(config.n_layer)
- ]
- self.layernorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="layernorm")
-
- def get_input_embeddings(self):
- return self.w
-
- def set_input_embeddings(self, new_embeddings):
- self.w = new_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
- """
- raise NotImplementedError
-
- @unpack_inputs
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFBaseModelOutputWithPast]:
- # If using past key value states, only the last tokens
- # should be given as an input
- if past_key_values is not None:
- if input_ids is not None:
- input_ids = input_ids[:, -1:]
- if inputs_embeds is not None:
- inputs_embeds = inputs_embeds[:, -1:]
- if token_type_ids is not None:
- token_type_ids = token_type_ids[:, -1:]
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- input_ids = tf.reshape(input_ids, [-1, input_shape[-1]])
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if past_key_values is None:
- past_length = 0
- past_key_values = [None] * len(self.h)
- else:
- past_length = shape_list(past_key_values[0][0])[-2]
- if position_ids is None:
- position_ids = tf.expand_dims(tf.range(past_length, input_shape[-1] + past_length, dtype=tf.int32), axis=0)
- position_ids = tf.tile(position_ids, [input_shape[0], 1])
-
- # Attention mask.
- if attention_mask is not None:
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask = tf.reshape(attention_mask, (input_shape[0], 1, 1, input_shape[1] + past_length))
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and -10000.0 for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
-
- one_cst = tf.constant(1.0)
- ten_thousand_cst = tf.constant(-10000.0)
- attention_mask = tf.cast(attention_mask, dtype=one_cst.dtype)
- attention_mask = tf.multiply(tf.subtract(one_cst, attention_mask), ten_thousand_cst)
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # head_mask has shape n_layer x batch x n_heads x N x N
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.num_layers
-
- if token_type_ids is not None:
- token_type_ids = tf.reshape(token_type_ids, [-1, shape_list(token_type_ids)[-1]])
- token_type_embeds = self.w(token_type_ids)
- token_type_embeds *= tf.math.sqrt(tf.cast(self.d_model_size, dtype=token_type_embeds.dtype))
- else:
- token_type_embeds = tf.constant(0.0)
- position_ids = tf.reshape(position_ids, [-1, shape_list(position_ids)[-1]])
-
- if inputs_embeds is None:
- check_embeddings_within_bounds(input_ids, self.w.input_dim)
- inputs_embeds = self.w(input_ids)
- seq_len = input_shape[-1]
- mask = 1 - tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)
-
- inputs_embeds *= tf.math.sqrt(tf.cast(self.d_model_size, inputs_embeds.dtype))
-
- pos_embeds = tf.gather(self.pos_encoding, position_ids)
- pos_embeds = tf.cast(pos_embeds, dtype=token_type_embeds.dtype)
- hidden_states = inputs_embeds + pos_embeds + token_type_embeds
-
- hidden_states = self.dropout(hidden_states, training=training)
-
- output_shape = input_shape + [shape_list(hidden_states)[-1]]
- presents = () if use_cache else None
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
- for i, (h, layer_past) in enumerate(zip(self.h, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (tf.reshape(hidden_states, output_shape),)
- outputs = h(
- hidden_states,
- mask,
- layer_past,
- attention_mask,
- head_mask[i],
- use_cache,
- output_attentions,
- training=training,
- )
- hidden_states, present = outputs[:2]
-
- if use_cache:
- presents = presents + (present,)
-
- if output_attentions:
- all_attentions = all_attentions + (outputs[2],)
-
- hidden_states = self.layernorm(hidden_states)
- hidden_states = tf.reshape(hidden_states, output_shape)
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if output_attentions:
- # let the number of heads free (-1) so we can extract attention even after head pruning
- attention_output_shape = input_shape[:-1] + [-1] + shape_list(all_attentions[0])[-2:]
- all_attentions = tuple(tf.reshape(t, attention_output_shape) for t in all_attentions)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)
-
- return TFBaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "w", None) is not None:
- with tf.name_scope(self.w.name):
- self.w.build(None)
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, None, self.config.n_embd])
- if getattr(self, "h", None) is not None:
- for layer in self.h:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-class TFCTRLPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = CTRLConfig
- base_model_prefix = "transformer"
-
-
-CTRL_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`CTRLConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CTRL_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length` if `past` is `None` else `past[0].shape[-2]` (`sequence_length` of
- input past key value states).
-
- Indices of input sequence tokens in the vocabulary.
-
- If `past` is used, only input IDs that do not have their past calculated should be passed as `input_ids`.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past (`List[tf.Tensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past` output below). Can be used to speed up sequential decoding. The token ids which have their past
- given to this model should not be passed as input ids as they have already been computed.
- attention_mask (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- use_cache (`bool`, *optional*):
- If set to `True`, `past` key value states are returned and can be used to speed up decoding (see `past`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare CTRL Model transformer outputting raw hidden-states without any specific head on top.",
- CTRL_START_DOCSTRING,
-)
-class TFCTRLModel(TFCTRLPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFCTRLMainLayer(config, name="transformer")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFBaseModelOutputWithPast]:
- outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
-
-
-class TFCTRLBiasLayer(keras.layers.Layer):
- """
- Bias as a layer. It is used for serialization purposes: `keras.Model.save_weights` stores on a per-layer basis,
- so all weights have to be registered in a layer.
- """
-
- def __init__(self, shape, initializer, trainable, name, **kwargs):
- super().__init__(name=name, **kwargs)
- self.shape = shape
- self.initializer = initializer
- self.trainable = trainable
-
- def build(self, input_shape):
- self.bias = self.add_weight(
- name="bias", shape=self.shape, initializer=self.initializer, trainable=self.trainable
- )
- super().build(input_shape)
-
- def call(self, x):
- return x + self.bias
-
-
-@add_start_docstrings(
- """
- The CTRL Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- CTRL_START_DOCSTRING,
-)
-class TFCTRLLMHeadModel(TFCTRLPreTrainedModel, TFCausalLanguageModelingLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFCTRLMainLayer(config, name="transformer")
- self.bias_layer = TFCTRLBiasLayer(
- name="lm_head", shape=[1, config.vocab_size], initializer="zeros", trainable=True
- )
-
- def get_output_embeddings(self):
- return self.get_input_embeddings()
-
- def set_output_embeddings(self, value):
- self.set_input_embeddings(value)
-
- def get_bias(self):
- return {"lm_head.bias": self.bias_layer.bias}
-
- def set_bias(self, value):
- # Replaces the existing layers containing bias for correct (de)serialization.
- vocab_size = value["lm_head.bias"].shape[-1]
- self.bias_layer = TFCTRLBiasLayer(
- name="final_logits_bias", shape=[1, vocab_size], initializer="zeros", trainable=True
- )
- self.bias_layer.build(None)
- self.bias_layer.bias.assign(value["lm_head.bias"])
-
- # Copied from transformers.models.gpt2.modeling_tf_gpt2.TFGPT2LMHeadModel.prepare_inputs_for_generation
- def prepare_inputs_for_generation(self, inputs, past_key_values=None, use_cache=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # only last token for inputs_ids if past is defined in kwargs
- if past_key_values:
- inputs = tf.expand_dims(inputs[:, -1], -1)
- if token_type_ids is not None:
- token_type_ids = tf.expand_dims(token_type_ids[:, -1], -1)
-
- position_ids = kwargs.get("position_ids", None)
- attention_mask = kwargs.get("attention_mask", None)
-
- if attention_mask is not None and position_ids is None:
- position_ids = tf.math.cumsum(attention_mask, axis=-1, exclusive=True)
- if past_key_values:
- position_ids = tf.expand_dims(position_ids[:, -1], -1)
-
- return {
- "input_ids": inputs,
- "attention_mask": attention_mask,
- "position_ids": position_ids,
- "past_key_values": past_key_values,
- "use_cache": use_cache,
- "token_type_ids": token_type_ids,
- }
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFCausalLMOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFCausalLMOutputWithPast]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the cross entropy classification loss. Indices should be in `[0, ...,
- config.vocab_size - 1]`.
- """
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_states = transformer_outputs[0]
- logits = tf.matmul(hidden_states, self.transformer.w.weights, transpose_b=True)
- logits = self.bias_layer(logits)
-
- loss = None
- if labels is not None:
- # shift labels to the left and cut last logit token
- shifted_logits = logits[:, :-1]
- labels = labels[:, 1:]
- loss = self.hf_compute_loss(labels, shifted_logits)
-
- if not return_dict:
- output = (logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFCausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "bias_layer", None) is not None:
- with tf.name_scope(self.bias_layer.name):
- self.bias_layer.build(None)
-
-
-@add_start_docstrings(
- """
- The CTRL Model transformer with a sequence classification head on top (linear layer).
-
- [`TFCTRLForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT-1, GPT-2) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- CTRL_START_DOCSTRING,
-)
-class TFCTRLForSequenceClassification(TFCTRLPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
- self.classifier = keras.layers.Dense(
- config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- name="classifier",
- use_bias=False,
- )
- self.transformer = TFCTRLMainLayer(config, name="transformer")
- self.config = config
-
- def get_output_embeddings(self):
- # Remove after transformers v4.32. Fix this model's `test_model_common_attributes` test too.
- logger.warning(
- "Sequence classification models do not have output embeddings. `.get_output_embeddings` will be removed "
- "in transformers v4.32."
- )
- return self.transformer.w
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFSequenceClassifierOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the cross entropy classification loss. Indices should be in `[0, ...,
- config.vocab_size - 1]`.
- """
-
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- hidden_states = transformer_outputs[0]
- logits = self.classifier(hidden_states)
- in_logits = None
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- sequence_lengths = (
- tf.argmax(tf.cast(tf.math.equal(input_ids, self.config.pad_token_id), input_ids.dtype), axis=-1)
- - 1
- )
- sequence_lengths = tf.where(sequence_lengths >= 0, sequence_lengths, input_ids.shape[-1] - 1)
- in_logits = tf.gather(logits, sequence_lengths, batch_dims=1, axis=1)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
- loss = None
-
- if labels is not None:
- if input_ids is not None:
- batch_size, sequence_length = shape_list(input_ids)[:2]
- else:
- batch_size, sequence_length = shape_list(inputs_embeds)[:2]
- if self.config.pad_token_id is None and batch_size != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
-
- if not tf.is_tensor(sequence_lengths):
- in_logits = logits[0:batch_size, sequence_lengths]
-
- loss = self.hf_compute_loss(tf.reshape(labels, [-1, 1]), tf.reshape(in_logits, [-1, self.num_labels]))
-
- pooled_logits = in_logits if in_logits is not None else logits
-
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=pooled_logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.n_embd])
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
diff --git a/transformers/models/ctrl/tokenization_ctrl.py b/transformers/models/ctrl/tokenization_ctrl.py
deleted file mode 100644
index fdae22d2c3001963f7666bb1c4ff4e52d3e5db1a..0000000000000000000000000000000000000000
--- a/transformers/models/ctrl/tokenization_ctrl.py
+++ /dev/null
@@ -1,249 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Salesforce and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for Salesforce CTRL."""
-
-
-import json
-import os
-from typing import Optional, Tuple
-
-import regex as re
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {
- "vocab_file": "vocab.json",
- "merges_file": "merges.txt",
-}
-
-
-CONTROL_CODES = {
- "Pregnancy": 168629,
- "Christianity": 7675,
- "Explain": 106423,
- "Fitness": 63440,
- "Saving": 63163,
- "Ask": 27171,
- "Ass": 95985,
- "Joke": 163509,
- "Questions": 45622,
- "Thoughts": 49605,
- "Retail": 52342,
- "Feminism": 164338,
- "Writing": 11992,
- "Atheism": 192263,
- "Netflix": 48616,
- "Computing": 39639,
- "Opinion": 43213,
- "Alone": 44967,
- "Funny": 58917,
- "Gaming": 40358,
- "Human": 4088,
- "India": 1331,
- "Joker": 77138,
- "Diet": 36206,
- "Legal": 11859,
- "Norman": 4939,
- "Tip": 72689,
- "Weight": 52343,
- "Movies": 46273,
- "Running": 23425,
- "Science": 2090,
- "Horror": 37793,
- "Confession": 60572,
- "Finance": 12250,
- "Politics": 16360,
- "Scary": 191985,
- "Support": 12654,
- "Technologies": 32516,
- "Teenage": 66160,
- "Event": 32769,
- "Learned": 67460,
- "Notion": 182770,
- "Wikipedia": 37583,
- "Books": 6665,
- "Extract": 76050,
- "Confessions": 102701,
- "Conspiracy": 75932,
- "Links": 63674,
- "Narcissus": 150425,
- "Relationship": 54766,
- "Relationships": 134796,
- "Reviews": 41671,
- "News": 4256,
- "Translation": 26820,
- "multilingual": 128406,
-}
-
-
-def get_pairs(word):
- """
- Return set of symbol pairs in a word.
-
- Word is represented as tuple of symbols (symbols being variable-length strings).
- """
- pairs = set()
- prev_char = word[0]
- for char in word[1:]:
- pairs.add((prev_char, char))
- prev_char = char
-
- pairs = set(pairs)
- return pairs
-
-
-class CTRLTokenizer(PreTrainedTokenizer):
- """
- Construct a CTRL tokenizer. Based on Byte-Pair-Encoding.
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- merges_file (`str`):
- Path to the merges file.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- control_codes = CONTROL_CODES
-
- def __init__(self, vocab_file, merges_file, unk_token="", **kwargs):
- with open(vocab_file, encoding="utf-8") as vocab_handle:
- self.encoder = json.load(vocab_handle)
- self.decoder = {v: k for k, v in self.encoder.items()}
- with open(merges_file, encoding="utf-8") as merges_handle:
- merges = merges_handle.read().split("\n")[1:-1]
- merges = [tuple(merge.split()) for merge in merges]
- self.bpe_ranks = dict(zip(merges, range(len(merges))))
- self.cache = {}
- super().__init__(unk_token=unk_token, **kwargs)
-
- @property
- def vocab_size(self):
- return len(self.encoder)
-
- def get_vocab(self):
- return dict(self.encoder, **self.added_tokens_encoder)
-
- def bpe(self, token):
- if token in self.cache:
- return self.cache[token]
- word = tuple(token)
- word = tuple(list(word[:-1]) + [word[-1] + ""])
- pairs = get_pairs(word)
-
- if not pairs:
- return token
-
- while True:
- bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
- if bigram not in self.bpe_ranks:
- break
- first, second = bigram
- new_word = []
- i = 0
- while i < len(word):
- try:
- j = word.index(first, i)
- except ValueError:
- new_word.extend(word[i:])
- break
- else:
- new_word.extend(word[i:j])
- i = j
-
- if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
- new_word.append(first + second)
- i += 2
- else:
- new_word.append(word[i])
- i += 1
- new_word = tuple(new_word)
- word = new_word
- if len(word) == 1:
- break
- else:
- pairs = get_pairs(word)
- word = "@@ ".join(word)
- word = word[:-4]
- self.cache[token] = word
- return word
-
- def _tokenize(self, text):
- """Tokenize a string."""
- split_tokens = []
-
- words = re.findall(r"\S+\n?", text)
-
- for token in words:
- split_tokens.extend(list(self.bpe(token).split(" ")))
- return split_tokens
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index, self.unk_token)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- out_string = " ".join(tokens).replace("@@ ", "").strip()
- return out_string
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- merge_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
- )
-
- with open(vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
- index = 0
- with open(merge_file, "w", encoding="utf-8") as writer:
- writer.write("#version: 0.2\n")
- for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
- " Please check that the tokenizer is not corrupted!"
- )
- index = token_index
- writer.write(" ".join(bpe_tokens) + "\n")
- index += 1
-
- return vocab_file, merge_file
-
- # def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True):
- # filtered_tokens = ' '.join(self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens))
- # tokens_generated_so_far = re.sub('(@@ )', '', string=filtered_tokens)
- # tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
- # return ''.join(tokens_generated_so_far)
diff --git a/transformers/models/cvt/__init__.py b/transformers/models/cvt/__init__.py
deleted file mode 100644
index 5241bb5a5f3a7a5ace9c7786926e1ff212e751fe..0000000000000000000000000000000000000000
--- a/transformers/models/cvt/__init__.py
+++ /dev/null
@@ -1,81 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available
-
-
-_import_structure = {"configuration_cvt": ["CVT_PRETRAINED_CONFIG_ARCHIVE_MAP", "CvtConfig"]}
-
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_cvt"] = [
- "CVT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "CvtForImageClassification",
- "CvtModel",
- "CvtPreTrainedModel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_cvt"] = [
- "TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFCvtForImageClassification",
- "TFCvtModel",
- "TFCvtPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_cvt import CVT_PRETRAINED_CONFIG_ARCHIVE_MAP, CvtConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_cvt import (
- CVT_PRETRAINED_MODEL_ARCHIVE_LIST,
- CvtForImageClassification,
- CvtModel,
- CvtPreTrainedModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_cvt import (
- TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFCvtForImageClassification,
- TFCvtModel,
- TFCvtPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/cvt/__pycache__/__init__.cpython-310.pyc b/transformers/models/cvt/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 5d134bdf4dd7c0e001803b9ddf101d4f2f288288..0000000000000000000000000000000000000000
Binary files a/transformers/models/cvt/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cvt/__pycache__/configuration_cvt.cpython-310.pyc b/transformers/models/cvt/__pycache__/configuration_cvt.cpython-310.pyc
deleted file mode 100644
index 4a6f046d5d89e1ea5ea372f6a9c2a0cb3e0b6fe6..0000000000000000000000000000000000000000
Binary files a/transformers/models/cvt/__pycache__/configuration_cvt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cvt/__pycache__/convert_cvt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/cvt/__pycache__/convert_cvt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index d98fcf069bc27ec7e86fd51e4eb9becfc2acba84..0000000000000000000000000000000000000000
Binary files a/transformers/models/cvt/__pycache__/convert_cvt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cvt/__pycache__/modeling_cvt.cpython-310.pyc b/transformers/models/cvt/__pycache__/modeling_cvt.cpython-310.pyc
deleted file mode 100644
index 09663f9b99c03e1511e0ef4f9e940ab1efec937a..0000000000000000000000000000000000000000
Binary files a/transformers/models/cvt/__pycache__/modeling_cvt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cvt/__pycache__/modeling_tf_cvt.cpython-310.pyc b/transformers/models/cvt/__pycache__/modeling_tf_cvt.cpython-310.pyc
deleted file mode 100644
index db84db3e108accb39a413b4b565f868b750a2f7e..0000000000000000000000000000000000000000
Binary files a/transformers/models/cvt/__pycache__/modeling_tf_cvt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/cvt/configuration_cvt.py b/transformers/models/cvt/configuration_cvt.py
deleted file mode 100644
index 412387af5e8a7bf21975207b513eeff90ca01479..0000000000000000000000000000000000000000
--- a/transformers/models/cvt/configuration_cvt.py
+++ /dev/null
@@ -1,146 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" CvT model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import CVT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class CvtConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`CvtModel`]. It is used to instantiate a CvT model
- according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the CvT
- [microsoft/cvt-13](https://huggingface.co/microsoft/cvt-13) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- patch_sizes (`List[int]`, *optional*, defaults to `[7, 3, 3]`):
- The kernel size of each encoder's patch embedding.
- patch_stride (`List[int]`, *optional*, defaults to `[4, 2, 2]`):
- The stride size of each encoder's patch embedding.
- patch_padding (`List[int]`, *optional*, defaults to `[2, 1, 1]`):
- The padding size of each encoder's patch embedding.
- embed_dim (`List[int]`, *optional*, defaults to `[64, 192, 384]`):
- Dimension of each of the encoder blocks.
- num_heads (`List[int]`, *optional*, defaults to `[1, 3, 6]`):
- Number of attention heads for each attention layer in each block of the Transformer encoder.
- depth (`List[int]`, *optional*, defaults to `[1, 2, 10]`):
- The number of layers in each encoder block.
- mlp_ratios (`List[float]`, *optional*, defaults to `[4.0, 4.0, 4.0, 4.0]`):
- Ratio of the size of the hidden layer compared to the size of the input layer of the Mix FFNs in the
- encoder blocks.
- attention_drop_rate (`List[float]`, *optional*, defaults to `[0.0, 0.0, 0.0]`):
- The dropout ratio for the attention probabilities.
- drop_rate (`List[float]`, *optional*, defaults to `[0.0, 0.0, 0.0]`):
- The dropout ratio for the patch embeddings probabilities.
- drop_path_rate (`List[float]`, *optional*, defaults to `[0.0, 0.0, 0.1]`):
- The dropout probability for stochastic depth, used in the blocks of the Transformer encoder.
- qkv_bias (`List[bool]`, *optional*, defaults to `[True, True, True]`):
- The bias bool for query, key and value in attentions
- cls_token (`List[bool]`, *optional*, defaults to `[False, False, True]`):
- Whether or not to add a classification token to the output of each of the last 3 stages.
- qkv_projection_method (`List[string]`, *optional*, defaults to ["dw_bn", "dw_bn", "dw_bn"]`):
- The projection method for query, key and value Default is depth-wise convolutions with batch norm. For
- Linear projection use "avg".
- kernel_qkv (`List[int]`, *optional*, defaults to `[3, 3, 3]`):
- The kernel size for query, key and value in attention layer
- padding_kv (`List[int]`, *optional*, defaults to `[1, 1, 1]`):
- The padding size for key and value in attention layer
- stride_kv (`List[int]`, *optional*, defaults to `[2, 2, 2]`):
- The stride size for key and value in attention layer
- padding_q (`List[int]`, *optional*, defaults to `[1, 1, 1]`):
- The padding size for query in attention layer
- stride_q (`List[int]`, *optional*, defaults to `[1, 1, 1]`):
- The stride size for query in attention layer
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-6):
- The epsilon used by the layer normalization layers.
-
- Example:
-
- ```python
- >>> from transformers import CvtConfig, CvtModel
-
- >>> # Initializing a Cvt msft/cvt style configuration
- >>> configuration = CvtConfig()
-
- >>> # Initializing a model (with random weights) from the msft/cvt style configuration
- >>> model = CvtModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "cvt"
-
- def __init__(
- self,
- num_channels=3,
- patch_sizes=[7, 3, 3],
- patch_stride=[4, 2, 2],
- patch_padding=[2, 1, 1],
- embed_dim=[64, 192, 384],
- num_heads=[1, 3, 6],
- depth=[1, 2, 10],
- mlp_ratio=[4.0, 4.0, 4.0],
- attention_drop_rate=[0.0, 0.0, 0.0],
- drop_rate=[0.0, 0.0, 0.0],
- drop_path_rate=[0.0, 0.0, 0.1],
- qkv_bias=[True, True, True],
- cls_token=[False, False, True],
- qkv_projection_method=["dw_bn", "dw_bn", "dw_bn"],
- kernel_qkv=[3, 3, 3],
- padding_kv=[1, 1, 1],
- stride_kv=[2, 2, 2],
- padding_q=[1, 1, 1],
- stride_q=[1, 1, 1],
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.num_channels = num_channels
- self.patch_sizes = patch_sizes
- self.patch_stride = patch_stride
- self.patch_padding = patch_padding
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.depth = depth
- self.mlp_ratio = mlp_ratio
- self.attention_drop_rate = attention_drop_rate
- self.drop_rate = drop_rate
- self.drop_path_rate = drop_path_rate
- self.qkv_bias = qkv_bias
- self.cls_token = cls_token
- self.qkv_projection_method = qkv_projection_method
- self.kernel_qkv = kernel_qkv
- self.padding_kv = padding_kv
- self.stride_kv = stride_kv
- self.padding_q = padding_q
- self.stride_q = stride_q
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
diff --git a/transformers/models/cvt/convert_cvt_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/cvt/convert_cvt_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index ea4edac16cdbae353ea7b5f93f297164360b476f..0000000000000000000000000000000000000000
--- a/transformers/models/cvt/convert_cvt_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,362 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert CvT checkpoints from the original repository.
-
-URL: https://github.com/microsoft/CvT"""
-
-
-import argparse
-import json
-from collections import OrderedDict
-
-import torch
-from huggingface_hub import cached_download, hf_hub_url
-
-from transformers import AutoImageProcessor, CvtConfig, CvtForImageClassification
-
-
-def embeddings(idx):
- """
- The function helps in renaming embedding layer weights.
-
- Args:
- idx: stage number in original model
- """
- embed = []
- embed.append(
- (
- f"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.weight",
- f"stage{idx}.patch_embed.proj.weight",
- )
- )
- embed.append(
- (
- f"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.bias",
- f"stage{idx}.patch_embed.proj.bias",
- )
- )
- embed.append(
- (
- f"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.weight",
- f"stage{idx}.patch_embed.norm.weight",
- )
- )
- embed.append(
- (
- f"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.bias",
- f"stage{idx}.patch_embed.norm.bias",
- )
- )
- return embed
-
-
-def attention(idx, cnt):
- """
- The function helps in renaming attention block layers weights.
-
- Args:
- idx: stage number in original model
- cnt: count of blocks in each stage
- """
- attention_weights = []
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.convolution.weight",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.conv.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.weight",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.bias",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.bias",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_mean",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_mean",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_var",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_var",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.num_batches_tracked",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.num_batches_tracked",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.convolution.weight",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.conv.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.weight",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.bias",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.bias",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_mean",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_mean",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_var",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_var",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.num_batches_tracked",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.num_batches_tracked",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.convolution.weight",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.conv.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.weight",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.bias",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.bias",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_mean",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_mean",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_var",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_var",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.num_batches_tracked",
- f"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.num_batches_tracked",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.weight",
- f"stage{idx}.blocks.{cnt}.attn.proj_q.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.bias",
- f"stage{idx}.blocks.{cnt}.attn.proj_q.bias",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.weight",
- f"stage{idx}.blocks.{cnt}.attn.proj_k.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.bias",
- f"stage{idx}.blocks.{cnt}.attn.proj_k.bias",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.weight",
- f"stage{idx}.blocks.{cnt}.attn.proj_v.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.bias",
- f"stage{idx}.blocks.{cnt}.attn.proj_v.bias",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.weight",
- f"stage{idx}.blocks.{cnt}.attn.proj.weight",
- )
- )
- attention_weights.append(
- (
- f"cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.bias",
- f"stage{idx}.blocks.{cnt}.attn.proj.bias",
- )
- )
- attention_weights.append(
- (f"cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.weight", f"stage{idx}.blocks.{cnt}.mlp.fc1.weight")
- )
- attention_weights.append(
- (f"cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.bias", f"stage{idx}.blocks.{cnt}.mlp.fc1.bias")
- )
- attention_weights.append(
- (f"cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.weight", f"stage{idx}.blocks.{cnt}.mlp.fc2.weight")
- )
- attention_weights.append(
- (f"cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.bias", f"stage{idx}.blocks.{cnt}.mlp.fc2.bias")
- )
- attention_weights.append(
- (f"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.weight", f"stage{idx}.blocks.{cnt}.norm1.weight")
- )
- attention_weights.append(
- (f"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.bias", f"stage{idx}.blocks.{cnt}.norm1.bias")
- )
- attention_weights.append(
- (f"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.weight", f"stage{idx}.blocks.{cnt}.norm2.weight")
- )
- attention_weights.append(
- (f"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.bias", f"stage{idx}.blocks.{cnt}.norm2.bias")
- )
- return attention_weights
-
-
-def cls_token(idx):
- """
- Function helps in renaming cls_token weights
- """
- token = []
- token.append((f"cvt.encoder.stages.{idx}.cls_token", "stage2.cls_token"))
- return token
-
-
-def final():
- """
- Function helps in renaming final classification layer
- """
- head = []
- head.append(("layernorm.weight", "norm.weight"))
- head.append(("layernorm.bias", "norm.bias"))
- head.append(("classifier.weight", "head.weight"))
- head.append(("classifier.bias", "head.bias"))
- return head
-
-
-def convert_cvt_checkpoint(cvt_model, image_size, cvt_file_name, pytorch_dump_folder):
- """
- Fucntion to convert the microsoft cvt checkpoint to huggingface checkpoint
- """
- img_labels_file = "imagenet-1k-id2label.json"
- num_labels = 1000
-
- repo_id = "huggingface/label-files"
- num_labels = num_labels
- id2label = json.load(open(cached_download(hf_hub_url(repo_id, img_labels_file, repo_type="dataset")), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
-
- id2label = id2label
- label2id = {v: k for k, v in id2label.items()}
-
- config = config = CvtConfig(num_labels=num_labels, id2label=id2label, label2id=label2id)
-
- # For depth size 13 (13 = 1+2+10)
- if cvt_model.rsplit("/", 1)[-1][4:6] == "13":
- config.depth = [1, 2, 10]
-
- # For depth size 21 (21 = 1+4+16)
- elif cvt_model.rsplit("/", 1)[-1][4:6] == "21":
- config.depth = [1, 4, 16]
-
- # For wide cvt (similar to wide-resnet) depth size 24 (w24 = 2 + 2 20)
- else:
- config.depth = [2, 2, 20]
- config.num_heads = [3, 12, 16]
- config.embed_dim = [192, 768, 1024]
-
- model = CvtForImageClassification(config)
- image_processor = AutoImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k")
- image_processor.size["shortest_edge"] = image_size
- original_weights = torch.load(cvt_file_name, map_location=torch.device("cpu"))
-
- huggingface_weights = OrderedDict()
- list_of_state_dict = []
-
- for idx in range(len(config.depth)):
- if config.cls_token[idx]:
- list_of_state_dict = list_of_state_dict + cls_token(idx)
- list_of_state_dict = list_of_state_dict + embeddings(idx)
- for cnt in range(config.depth[idx]):
- list_of_state_dict = list_of_state_dict + attention(idx, cnt)
-
- list_of_state_dict = list_of_state_dict + final()
- for gg in list_of_state_dict:
- print(gg)
- for i in range(len(list_of_state_dict)):
- huggingface_weights[list_of_state_dict[i][0]] = original_weights[list_of_state_dict[i][1]]
-
- model.load_state_dict(huggingface_weights)
- model.save_pretrained(pytorch_dump_folder)
- image_processor.save_pretrained(pytorch_dump_folder)
-
-
-# Download the weights from zoo: https://1drv.ms/u/s!AhIXJn_J-blW9RzF3rMW7SsLHa8h?e=blQ0Al
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--cvt_model",
- default="cvt-w24",
- type=str,
- help="Name of the cvt model you'd like to convert.",
- )
- parser.add_argument(
- "--image_size",
- default=384,
- type=int,
- help="Input Image Size",
- )
- parser.add_argument(
- "--cvt_file_name",
- default=r"cvtmodels\CvT-w24-384x384-IN-22k.pth",
- type=str,
- help="Input Image Size",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
- )
-
- args = parser.parse_args()
- convert_cvt_checkpoint(args.cvt_model, args.image_size, args.cvt_file_name, args.pytorch_dump_folder_path)
diff --git a/transformers/models/cvt/modeling_cvt.py b/transformers/models/cvt/modeling_cvt.py
deleted file mode 100644
index 25cf3963cbe10c9f4e06d71154ae3a09a3e16d45..0000000000000000000000000000000000000000
--- a/transformers/models/cvt/modeling_cvt.py
+++ /dev/null
@@ -1,725 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Microsoft Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch CvT model."""
-
-
-import collections.abc
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_outputs import ImageClassifierOutputWithNoAttention, ModelOutput
-from ...modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import logging
-from .configuration_cvt import CvtConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "CvtConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "microsoft/cvt-13"
-_EXPECTED_OUTPUT_SHAPE = [1, 384, 14, 14]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "microsoft/cvt-13"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import CVT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class BaseModelOutputWithCLSToken(ModelOutput):
- """
- Base class for model's outputs, with potential hidden states and attentions.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- cls_token_value (`torch.FloatTensor` of shape `(batch_size, 1, hidden_size)`):
- Classification token at the output of the last layer of the model.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- """
-
- last_hidden_state: torch.FloatTensor = None
- cls_token_value: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitDropPath
-class CvtDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-class CvtEmbeddings(nn.Module):
- """
- Construct the CvT embeddings.
- """
-
- def __init__(self, patch_size, num_channels, embed_dim, stride, padding, dropout_rate):
- super().__init__()
- self.convolution_embeddings = CvtConvEmbeddings(
- patch_size=patch_size, num_channels=num_channels, embed_dim=embed_dim, stride=stride, padding=padding
- )
- self.dropout = nn.Dropout(dropout_rate)
-
- def forward(self, pixel_values):
- hidden_state = self.convolution_embeddings(pixel_values)
- hidden_state = self.dropout(hidden_state)
- return hidden_state
-
-
-class CvtConvEmbeddings(nn.Module):
- """
- Image to Conv Embedding.
- """
-
- def __init__(self, patch_size, num_channels, embed_dim, stride, padding):
- super().__init__()
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- self.patch_size = patch_size
- self.projection = nn.Conv2d(num_channels, embed_dim, kernel_size=patch_size, stride=stride, padding=padding)
- self.normalization = nn.LayerNorm(embed_dim)
-
- def forward(self, pixel_values):
- pixel_values = self.projection(pixel_values)
- batch_size, num_channels, height, width = pixel_values.shape
- hidden_size = height * width
- # rearrange "b c h w -> b (h w) c"
- pixel_values = pixel_values.view(batch_size, num_channels, hidden_size).permute(0, 2, 1)
- if self.normalization:
- pixel_values = self.normalization(pixel_values)
- # rearrange "b (h w) c" -> b c h w"
- pixel_values = pixel_values.permute(0, 2, 1).view(batch_size, num_channels, height, width)
- return pixel_values
-
-
-class CvtSelfAttentionConvProjection(nn.Module):
- def __init__(self, embed_dim, kernel_size, padding, stride):
- super().__init__()
- self.convolution = nn.Conv2d(
- embed_dim,
- embed_dim,
- kernel_size=kernel_size,
- padding=padding,
- stride=stride,
- bias=False,
- groups=embed_dim,
- )
- self.normalization = nn.BatchNorm2d(embed_dim)
-
- def forward(self, hidden_state):
- hidden_state = self.convolution(hidden_state)
- hidden_state = self.normalization(hidden_state)
- return hidden_state
-
-
-class CvtSelfAttentionLinearProjection(nn.Module):
- def forward(self, hidden_state):
- batch_size, num_channels, height, width = hidden_state.shape
- hidden_size = height * width
- # rearrange " b c h w -> b (h w) c"
- hidden_state = hidden_state.view(batch_size, num_channels, hidden_size).permute(0, 2, 1)
- return hidden_state
-
-
-class CvtSelfAttentionProjection(nn.Module):
- def __init__(self, embed_dim, kernel_size, padding, stride, projection_method="dw_bn"):
- super().__init__()
- if projection_method == "dw_bn":
- self.convolution_projection = CvtSelfAttentionConvProjection(embed_dim, kernel_size, padding, stride)
- self.linear_projection = CvtSelfAttentionLinearProjection()
-
- def forward(self, hidden_state):
- hidden_state = self.convolution_projection(hidden_state)
- hidden_state = self.linear_projection(hidden_state)
- return hidden_state
-
-
-class CvtSelfAttention(nn.Module):
- def __init__(
- self,
- num_heads,
- embed_dim,
- kernel_size,
- padding_q,
- padding_kv,
- stride_q,
- stride_kv,
- qkv_projection_method,
- qkv_bias,
- attention_drop_rate,
- with_cls_token=True,
- **kwargs,
- ):
- super().__init__()
- self.scale = embed_dim**-0.5
- self.with_cls_token = with_cls_token
- self.embed_dim = embed_dim
- self.num_heads = num_heads
-
- self.convolution_projection_query = CvtSelfAttentionProjection(
- embed_dim,
- kernel_size,
- padding_q,
- stride_q,
- projection_method="linear" if qkv_projection_method == "avg" else qkv_projection_method,
- )
- self.convolution_projection_key = CvtSelfAttentionProjection(
- embed_dim, kernel_size, padding_kv, stride_kv, projection_method=qkv_projection_method
- )
- self.convolution_projection_value = CvtSelfAttentionProjection(
- embed_dim, kernel_size, padding_kv, stride_kv, projection_method=qkv_projection_method
- )
-
- self.projection_query = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)
- self.projection_key = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)
- self.projection_value = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)
-
- self.dropout = nn.Dropout(attention_drop_rate)
-
- def rearrange_for_multi_head_attention(self, hidden_state):
- batch_size, hidden_size, _ = hidden_state.shape
- head_dim = self.embed_dim // self.num_heads
- # rearrange 'b t (h d) -> b h t d'
- return hidden_state.view(batch_size, hidden_size, self.num_heads, head_dim).permute(0, 2, 1, 3)
-
- def forward(self, hidden_state, height, width):
- if self.with_cls_token:
- cls_token, hidden_state = torch.split(hidden_state, [1, height * width], 1)
- batch_size, hidden_size, num_channels = hidden_state.shape
- # rearrange "b (h w) c -> b c h w"
- hidden_state = hidden_state.permute(0, 2, 1).view(batch_size, num_channels, height, width)
-
- key = self.convolution_projection_key(hidden_state)
- query = self.convolution_projection_query(hidden_state)
- value = self.convolution_projection_value(hidden_state)
-
- if self.with_cls_token:
- query = torch.cat((cls_token, query), dim=1)
- key = torch.cat((cls_token, key), dim=1)
- value = torch.cat((cls_token, value), dim=1)
-
- head_dim = self.embed_dim // self.num_heads
-
- query = self.rearrange_for_multi_head_attention(self.projection_query(query))
- key = self.rearrange_for_multi_head_attention(self.projection_key(key))
- value = self.rearrange_for_multi_head_attention(self.projection_value(value))
-
- attention_score = torch.einsum("bhlk,bhtk->bhlt", [query, key]) * self.scale
- attention_probs = torch.nn.functional.softmax(attention_score, dim=-1)
- attention_probs = self.dropout(attention_probs)
-
- context = torch.einsum("bhlt,bhtv->bhlv", [attention_probs, value])
- # rearrange"b h t d -> b t (h d)"
- _, _, hidden_size, _ = context.shape
- context = context.permute(0, 2, 1, 3).contiguous().view(batch_size, hidden_size, self.num_heads * head_dim)
- return context
-
-
-class CvtSelfOutput(nn.Module):
- """
- The residual connection is defined in CvtLayer instead of here (as is the case with other models), due to the
- layernorm applied before each block.
- """
-
- def __init__(self, embed_dim, drop_rate):
- super().__init__()
- self.dense = nn.Linear(embed_dim, embed_dim)
- self.dropout = nn.Dropout(drop_rate)
-
- def forward(self, hidden_state, input_tensor):
- hidden_state = self.dense(hidden_state)
- hidden_state = self.dropout(hidden_state)
- return hidden_state
-
-
-class CvtAttention(nn.Module):
- def __init__(
- self,
- num_heads,
- embed_dim,
- kernel_size,
- padding_q,
- padding_kv,
- stride_q,
- stride_kv,
- qkv_projection_method,
- qkv_bias,
- attention_drop_rate,
- drop_rate,
- with_cls_token=True,
- ):
- super().__init__()
- self.attention = CvtSelfAttention(
- num_heads,
- embed_dim,
- kernel_size,
- padding_q,
- padding_kv,
- stride_q,
- stride_kv,
- qkv_projection_method,
- qkv_bias,
- attention_drop_rate,
- with_cls_token,
- )
- self.output = CvtSelfOutput(embed_dim, drop_rate)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.attention.query = prune_linear_layer(self.attention.query, index)
- self.attention.key = prune_linear_layer(self.attention.key, index)
- self.attention.value = prune_linear_layer(self.attention.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
- self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(self, hidden_state, height, width):
- self_output = self.attention(hidden_state, height, width)
- attention_output = self.output(self_output, hidden_state)
- return attention_output
-
-
-class CvtIntermediate(nn.Module):
- def __init__(self, embed_dim, mlp_ratio):
- super().__init__()
- self.dense = nn.Linear(embed_dim, int(embed_dim * mlp_ratio))
- self.activation = nn.GELU()
-
- def forward(self, hidden_state):
- hidden_state = self.dense(hidden_state)
- hidden_state = self.activation(hidden_state)
- return hidden_state
-
-
-class CvtOutput(nn.Module):
- def __init__(self, embed_dim, mlp_ratio, drop_rate):
- super().__init__()
- self.dense = nn.Linear(int(embed_dim * mlp_ratio), embed_dim)
- self.dropout = nn.Dropout(drop_rate)
-
- def forward(self, hidden_state, input_tensor):
- hidden_state = self.dense(hidden_state)
- hidden_state = self.dropout(hidden_state)
- hidden_state = hidden_state + input_tensor
- return hidden_state
-
-
-class CvtLayer(nn.Module):
- """
- CvtLayer composed by attention layers, normalization and multi-layer perceptrons (mlps).
- """
-
- def __init__(
- self,
- num_heads,
- embed_dim,
- kernel_size,
- padding_q,
- padding_kv,
- stride_q,
- stride_kv,
- qkv_projection_method,
- qkv_bias,
- attention_drop_rate,
- drop_rate,
- mlp_ratio,
- drop_path_rate,
- with_cls_token=True,
- ):
- super().__init__()
- self.attention = CvtAttention(
- num_heads,
- embed_dim,
- kernel_size,
- padding_q,
- padding_kv,
- stride_q,
- stride_kv,
- qkv_projection_method,
- qkv_bias,
- attention_drop_rate,
- drop_rate,
- with_cls_token,
- )
-
- self.intermediate = CvtIntermediate(embed_dim, mlp_ratio)
- self.output = CvtOutput(embed_dim, mlp_ratio, drop_rate)
- self.drop_path = CvtDropPath(drop_prob=drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
- self.layernorm_before = nn.LayerNorm(embed_dim)
- self.layernorm_after = nn.LayerNorm(embed_dim)
-
- def forward(self, hidden_state, height, width):
- self_attention_output = self.attention(
- self.layernorm_before(hidden_state), # in Cvt, layernorm is applied before self-attention
- height,
- width,
- )
- attention_output = self_attention_output
- attention_output = self.drop_path(attention_output)
-
- # first residual connection
- hidden_state = attention_output + hidden_state
-
- # in Cvt, layernorm is also applied after self-attention
- layer_output = self.layernorm_after(hidden_state)
- layer_output = self.intermediate(layer_output)
-
- # second residual connection is done here
- layer_output = self.output(layer_output, hidden_state)
- layer_output = self.drop_path(layer_output)
- return layer_output
-
-
-class CvtStage(nn.Module):
- def __init__(self, config, stage):
- super().__init__()
- self.config = config
- self.stage = stage
- if self.config.cls_token[self.stage]:
- self.cls_token = nn.Parameter(torch.randn(1, 1, self.config.embed_dim[-1]))
-
- self.embedding = CvtEmbeddings(
- patch_size=config.patch_sizes[self.stage],
- stride=config.patch_stride[self.stage],
- num_channels=config.num_channels if self.stage == 0 else config.embed_dim[self.stage - 1],
- embed_dim=config.embed_dim[self.stage],
- padding=config.patch_padding[self.stage],
- dropout_rate=config.drop_rate[self.stage],
- )
-
- drop_path_rates = [x.item() for x in torch.linspace(0, config.drop_path_rate[self.stage], config.depth[stage])]
-
- self.layers = nn.Sequential(
- *[
- CvtLayer(
- num_heads=config.num_heads[self.stage],
- embed_dim=config.embed_dim[self.stage],
- kernel_size=config.kernel_qkv[self.stage],
- padding_q=config.padding_q[self.stage],
- padding_kv=config.padding_kv[self.stage],
- stride_kv=config.stride_kv[self.stage],
- stride_q=config.stride_q[self.stage],
- qkv_projection_method=config.qkv_projection_method[self.stage],
- qkv_bias=config.qkv_bias[self.stage],
- attention_drop_rate=config.attention_drop_rate[self.stage],
- drop_rate=config.drop_rate[self.stage],
- drop_path_rate=drop_path_rates[self.stage],
- mlp_ratio=config.mlp_ratio[self.stage],
- with_cls_token=config.cls_token[self.stage],
- )
- for _ in range(config.depth[self.stage])
- ]
- )
-
- def forward(self, hidden_state):
- cls_token = None
- hidden_state = self.embedding(hidden_state)
- batch_size, num_channels, height, width = hidden_state.shape
- # rearrange b c h w -> b (h w) c"
- hidden_state = hidden_state.view(batch_size, num_channels, height * width).permute(0, 2, 1)
- if self.config.cls_token[self.stage]:
- cls_token = self.cls_token.expand(batch_size, -1, -1)
- hidden_state = torch.cat((cls_token, hidden_state), dim=1)
-
- for layer in self.layers:
- layer_outputs = layer(hidden_state, height, width)
- hidden_state = layer_outputs
-
- if self.config.cls_token[self.stage]:
- cls_token, hidden_state = torch.split(hidden_state, [1, height * width], 1)
- hidden_state = hidden_state.permute(0, 2, 1).view(batch_size, num_channels, height, width)
- return hidden_state, cls_token
-
-
-class CvtEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.stages = nn.ModuleList([])
- for stage_idx in range(len(config.depth)):
- self.stages.append(CvtStage(config, stage_idx))
-
- def forward(self, pixel_values, output_hidden_states=False, return_dict=True):
- all_hidden_states = () if output_hidden_states else None
- hidden_state = pixel_values
-
- cls_token = None
- for _, (stage_module) in enumerate(self.stages):
- hidden_state, cls_token = stage_module(hidden_state)
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_state,)
-
- if not return_dict:
- return tuple(v for v in [hidden_state, cls_token, all_hidden_states] if v is not None)
-
- return BaseModelOutputWithCLSToken(
- last_hidden_state=hidden_state,
- cls_token_value=cls_token,
- hidden_states=all_hidden_states,
- )
-
-
-class CvtPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = CvtConfig
- base_model_prefix = "cvt"
- main_input_name = "pixel_values"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- module.weight.data = nn.init.trunc_normal_(module.weight.data, mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
- elif isinstance(module, CvtStage):
- if self.config.cls_token[module.stage]:
- module.cls_token.data = nn.init.trunc_normal_(
- torch.zeros(1, 1, self.config.embed_dim[-1]), mean=0.0, std=self.config.initializer_range
- )
-
-
-CVT_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`CvtConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-CVT_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`CvtImageProcessor.__call__`]
- for details.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Cvt Model transformer outputting raw hidden-states without any specific head on top.",
- CVT_START_DOCSTRING,
-)
-class CvtModel(CvtPreTrainedModel):
- def __init__(self, config, add_pooling_layer=True):
- super().__init__(config)
- self.config = config
- self.encoder = CvtEncoder(config)
- self.post_init()
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(CVT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithCLSToken,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithCLSToken]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- encoder_outputs = self.encoder(
- pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[1:]
-
- return BaseModelOutputWithCLSToken(
- last_hidden_state=sequence_output,
- cls_token_value=encoder_outputs.cls_token_value,
- hidden_states=encoder_outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- Cvt Model transformer with an image classification head on top (a linear layer on top of the final hidden state of
- the [CLS] token) e.g. for ImageNet.
- """,
- CVT_START_DOCSTRING,
-)
-class CvtForImageClassification(CvtPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.cvt = CvtModel(config, add_pooling_layer=False)
- self.layernorm = nn.LayerNorm(config.embed_dim[-1])
- # Classifier head
- self.classifier = (
- nn.Linear(config.embed_dim[-1], config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(CVT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=ImageClassifierOutputWithNoAttention,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- outputs = self.cvt(
- pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- cls_token = outputs[1]
- if self.config.cls_token[-1]:
- sequence_output = self.layernorm(cls_token)
- else:
- batch_size, num_channels, height, width = sequence_output.shape
- # rearrange "b c h w -> b (h w) c"
- sequence_output = sequence_output.view(batch_size, num_channels, height * width).permute(0, 2, 1)
- sequence_output = self.layernorm(sequence_output)
-
- sequence_output_mean = sequence_output.mean(dim=1)
- logits = self.classifier(sequence_output_mean)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.config.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.config.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
diff --git a/transformers/models/cvt/modeling_tf_cvt.py b/transformers/models/cvt/modeling_tf_cvt.py
deleted file mode 100644
index 5664412effb594852c86106afa705ac44fc8fb52..0000000000000000000000000000000000000000
--- a/transformers/models/cvt/modeling_tf_cvt.py
+++ /dev/null
@@ -1,1097 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Microsoft Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 Cvt model."""
-
-
-from __future__ import annotations
-
-import collections.abc
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import tensorflow as tf
-
-from ...modeling_tf_outputs import TFImageClassifierOutputWithNoAttention
-from ...modeling_tf_utils import (
- TFModelInputType,
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import shape_list, stable_softmax
-from ...utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_cvt import CvtConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "CvtConfig"
-
-
-from ..deprecated._archive_maps import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class TFBaseModelOutputWithCLSToken(ModelOutput):
- """
- Base class for model's outputs.
-
- Args:
- last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- cls_token_value (`tf.Tensor` of shape `(batch_size, 1, hidden_size)`):
- Classification token at the output of the last layer of the model.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus
- the initial embedding outputs.
- """
-
- last_hidden_state: tf.Tensor = None
- cls_token_value: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor, ...] | None = None
-
-
-class TFCvtDropPath(keras.layers.Layer):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
- References:
- (1) github.com:rwightman/pytorch-image-models
- """
-
- def __init__(self, drop_prob: float, **kwargs):
- super().__init__(**kwargs)
- self.drop_prob = drop_prob
-
- def call(self, x: tf.Tensor, training=None):
- if self.drop_prob == 0.0 or not training:
- return x
- keep_prob = 1 - self.drop_prob
- shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)
- random_tensor = keep_prob + tf.random.uniform(shape, 0, 1, dtype=self.compute_dtype)
- random_tensor = tf.floor(random_tensor)
- return (x / keep_prob) * random_tensor
-
-
-class TFCvtEmbeddings(keras.layers.Layer):
- """Construct the Convolutional Token Embeddings."""
-
- def __init__(
- self,
- config: CvtConfig,
- patch_size: int,
- num_channels: int,
- embed_dim: int,
- stride: int,
- padding: int,
- dropout_rate: float,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.convolution_embeddings = TFCvtConvEmbeddings(
- config,
- patch_size=patch_size,
- num_channels=num_channels,
- embed_dim=embed_dim,
- stride=stride,
- padding=padding,
- name="convolution_embeddings",
- )
- self.dropout = keras.layers.Dropout(dropout_rate)
-
- def call(self, pixel_values: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_state = self.convolution_embeddings(pixel_values)
- hidden_state = self.dropout(hidden_state, training=training)
- return hidden_state
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convolution_embeddings", None) is not None:
- with tf.name_scope(self.convolution_embeddings.name):
- self.convolution_embeddings.build(None)
-
-
-class TFCvtConvEmbeddings(keras.layers.Layer):
- """Image to Convolution Embeddings. This convolutional operation aims to model local spatial contexts."""
-
- def __init__(
- self,
- config: CvtConfig,
- patch_size: int,
- num_channels: int,
- embed_dim: int,
- stride: int,
- padding: int,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.padding = keras.layers.ZeroPadding2D(padding=padding)
- self.patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- self.projection = keras.layers.Conv2D(
- filters=embed_dim,
- kernel_size=patch_size,
- strides=stride,
- padding="valid",
- data_format="channels_last",
- kernel_initializer=get_initializer(config.initializer_range),
- name="projection",
- )
- # Using the same default epsilon as PyTorch
- self.normalization = keras.layers.LayerNormalization(epsilon=1e-5, name="normalization")
- self.num_channels = num_channels
- self.embed_dim = embed_dim
-
- def call(self, pixel_values: tf.Tensor) -> tf.Tensor:
- if isinstance(pixel_values, dict):
- pixel_values = pixel_values["pixel_values"]
-
- pixel_values = self.projection(self.padding(pixel_values))
-
- # "batch_size, height, width, num_channels -> batch_size, (height*width), num_channels"
- batch_size, height, width, num_channels = shape_list(pixel_values)
- hidden_size = height * width
- pixel_values = tf.reshape(pixel_values, shape=(batch_size, hidden_size, num_channels))
- pixel_values = self.normalization(pixel_values)
-
- # "batch_size, (height*width), num_channels -> batch_size, height, width, num_channels"
- pixel_values = tf.reshape(pixel_values, shape=(batch_size, height, width, num_channels))
- return pixel_values
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "projection", None) is not None:
- with tf.name_scope(self.projection.name):
- self.projection.build([None, None, None, self.num_channels])
- if getattr(self, "normalization", None) is not None:
- with tf.name_scope(self.normalization.name):
- self.normalization.build([None, None, self.embed_dim])
-
-
-class TFCvtSelfAttentionConvProjection(keras.layers.Layer):
- """Convolutional projection layer."""
-
- def __init__(self, config: CvtConfig, embed_dim: int, kernel_size: int, stride: int, padding: int, **kwargs):
- super().__init__(**kwargs)
- self.padding = keras.layers.ZeroPadding2D(padding=padding)
- self.convolution = keras.layers.Conv2D(
- filters=embed_dim,
- kernel_size=kernel_size,
- kernel_initializer=get_initializer(config.initializer_range),
- padding="valid",
- strides=stride,
- use_bias=False,
- name="convolution",
- groups=embed_dim,
- )
- # Using the same default epsilon as PyTorch, TF uses (1 - pytorch momentum)
- self.normalization = keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.9, name="normalization")
- self.embed_dim = embed_dim
-
- def call(self, hidden_state: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_state = self.convolution(self.padding(hidden_state))
- hidden_state = self.normalization(hidden_state, training=training)
- return hidden_state
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convolution", None) is not None:
- with tf.name_scope(self.convolution.name):
- self.convolution.build([None, None, None, self.embed_dim])
- if getattr(self, "normalization", None) is not None:
- with tf.name_scope(self.normalization.name):
- self.normalization.build([None, None, None, self.embed_dim])
-
-
-class TFCvtSelfAttentionLinearProjection(keras.layers.Layer):
- """Linear projection layer used to flatten tokens into 1D."""
-
- def call(self, hidden_state: tf.Tensor) -> tf.Tensor:
- # "batch_size, height, width, num_channels -> batch_size, (height*width), num_channels"
- batch_size, height, width, num_channels = shape_list(hidden_state)
- hidden_size = height * width
- hidden_state = tf.reshape(hidden_state, shape=(batch_size, hidden_size, num_channels))
- return hidden_state
-
-
-class TFCvtSelfAttentionProjection(keras.layers.Layer):
- """Convolutional Projection for Attention."""
-
- def __init__(
- self,
- config: CvtConfig,
- embed_dim: int,
- kernel_size: int,
- stride: int,
- padding: int,
- projection_method: str = "dw_bn",
- **kwargs,
- ):
- super().__init__(**kwargs)
- if projection_method == "dw_bn":
- self.convolution_projection = TFCvtSelfAttentionConvProjection(
- config, embed_dim, kernel_size, stride, padding, name="convolution_projection"
- )
- self.linear_projection = TFCvtSelfAttentionLinearProjection()
-
- def call(self, hidden_state: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_state = self.convolution_projection(hidden_state, training=training)
- hidden_state = self.linear_projection(hidden_state)
- return hidden_state
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convolution_projection", None) is not None:
- with tf.name_scope(self.convolution_projection.name):
- self.convolution_projection.build(None)
-
-
-class TFCvtSelfAttention(keras.layers.Layer):
- """
- Self-attention layer. A depth-wise separable convolution operation (Convolutional Projection), is applied for
- query, key, and value embeddings.
- """
-
- def __init__(
- self,
- config: CvtConfig,
- num_heads: int,
- embed_dim: int,
- kernel_size: int,
- stride_q: int,
- stride_kv: int,
- padding_q: int,
- padding_kv: int,
- qkv_projection_method: str,
- qkv_bias: bool,
- attention_drop_rate: float,
- with_cls_token: bool = True,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.scale = embed_dim**-0.5
- self.with_cls_token = with_cls_token
- self.embed_dim = embed_dim
- self.num_heads = num_heads
-
- self.convolution_projection_query = TFCvtSelfAttentionProjection(
- config,
- embed_dim,
- kernel_size,
- stride_q,
- padding_q,
- projection_method="linear" if qkv_projection_method == "avg" else qkv_projection_method,
- name="convolution_projection_query",
- )
- self.convolution_projection_key = TFCvtSelfAttentionProjection(
- config,
- embed_dim,
- kernel_size,
- stride_kv,
- padding_kv,
- projection_method=qkv_projection_method,
- name="convolution_projection_key",
- )
- self.convolution_projection_value = TFCvtSelfAttentionProjection(
- config,
- embed_dim,
- kernel_size,
- stride_kv,
- padding_kv,
- projection_method=qkv_projection_method,
- name="convolution_projection_value",
- )
-
- self.projection_query = keras.layers.Dense(
- units=embed_dim,
- kernel_initializer=get_initializer(config.initializer_range),
- use_bias=qkv_bias,
- bias_initializer="zeros",
- name="projection_query",
- )
- self.projection_key = keras.layers.Dense(
- units=embed_dim,
- kernel_initializer=get_initializer(config.initializer_range),
- use_bias=qkv_bias,
- bias_initializer="zeros",
- name="projection_key",
- )
- self.projection_value = keras.layers.Dense(
- units=embed_dim,
- kernel_initializer=get_initializer(config.initializer_range),
- use_bias=qkv_bias,
- bias_initializer="zeros",
- name="projection_value",
- )
- self.dropout = keras.layers.Dropout(attention_drop_rate)
-
- def rearrange_for_multi_head_attention(self, hidden_state: tf.Tensor) -> tf.Tensor:
- batch_size, hidden_size, _ = shape_list(hidden_state)
- head_dim = self.embed_dim // self.num_heads
- hidden_state = tf.reshape(hidden_state, shape=(batch_size, hidden_size, self.num_heads, head_dim))
- hidden_state = tf.transpose(hidden_state, perm=(0, 2, 1, 3))
- return hidden_state
-
- def call(self, hidden_state: tf.Tensor, height: int, width: int, training: bool = False) -> tf.Tensor:
- if self.with_cls_token:
- cls_token, hidden_state = tf.split(hidden_state, [1, height * width], 1)
-
- # "batch_size, (height*width), num_channels -> batch_size, height, width, num_channels"
- batch_size, hidden_size, num_channels = shape_list(hidden_state)
- hidden_state = tf.reshape(hidden_state, shape=(batch_size, height, width, num_channels))
-
- key = self.convolution_projection_key(hidden_state, training=training)
- query = self.convolution_projection_query(hidden_state, training=training)
- value = self.convolution_projection_value(hidden_state, training=training)
-
- if self.with_cls_token:
- query = tf.concat((cls_token, query), axis=1)
- key = tf.concat((cls_token, key), axis=1)
- value = tf.concat((cls_token, value), axis=1)
-
- head_dim = self.embed_dim // self.num_heads
-
- query = self.rearrange_for_multi_head_attention(self.projection_query(query))
- key = self.rearrange_for_multi_head_attention(self.projection_key(key))
- value = self.rearrange_for_multi_head_attention(self.projection_value(value))
-
- attention_score = tf.matmul(query, key, transpose_b=True) * self.scale
- attention_probs = stable_softmax(logits=attention_score, axis=-1)
- attention_probs = self.dropout(attention_probs, training=training)
-
- context = tf.matmul(attention_probs, value)
- # "batch_size, num_heads, hidden_size, head_dim -> batch_size, hidden_size, (num_heads*head_dim)"
- _, _, hidden_size, _ = shape_list(context)
- context = tf.transpose(context, perm=(0, 2, 1, 3))
- context = tf.reshape(context, (batch_size, hidden_size, self.num_heads * head_dim))
- return context
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convolution_projection_query", None) is not None:
- with tf.name_scope(self.convolution_projection_query.name):
- self.convolution_projection_query.build(None)
- if getattr(self, "convolution_projection_key", None) is not None:
- with tf.name_scope(self.convolution_projection_key.name):
- self.convolution_projection_key.build(None)
- if getattr(self, "convolution_projection_value", None) is not None:
- with tf.name_scope(self.convolution_projection_value.name):
- self.convolution_projection_value.build(None)
- if getattr(self, "projection_query", None) is not None:
- with tf.name_scope(self.projection_query.name):
- self.projection_query.build([None, None, self.embed_dim])
- if getattr(self, "projection_key", None) is not None:
- with tf.name_scope(self.projection_key.name):
- self.projection_key.build([None, None, self.embed_dim])
- if getattr(self, "projection_value", None) is not None:
- with tf.name_scope(self.projection_value.name):
- self.projection_value.build([None, None, self.embed_dim])
-
-
-class TFCvtSelfOutput(keras.layers.Layer):
- """Output of the Attention layer ."""
-
- def __init__(self, config: CvtConfig, embed_dim: int, drop_rate: float, **kwargs):
- super().__init__(**kwargs)
- self.dense = keras.layers.Dense(
- units=embed_dim, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.dropout = keras.layers.Dropout(drop_rate)
- self.embed_dim = embed_dim
-
- def call(self, hidden_state: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_state = self.dense(inputs=hidden_state)
- hidden_state = self.dropout(inputs=hidden_state, training=training)
- return hidden_state
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.embed_dim])
-
-
-class TFCvtAttention(keras.layers.Layer):
- """Attention layer. First chunk of the convolutional transformer block."""
-
- def __init__(
- self,
- config: CvtConfig,
- num_heads: int,
- embed_dim: int,
- kernel_size: int,
- stride_q: int,
- stride_kv: int,
- padding_q: int,
- padding_kv: int,
- qkv_projection_method: str,
- qkv_bias: bool,
- attention_drop_rate: float,
- drop_rate: float,
- with_cls_token: bool = True,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.attention = TFCvtSelfAttention(
- config,
- num_heads,
- embed_dim,
- kernel_size,
- stride_q,
- stride_kv,
- padding_q,
- padding_kv,
- qkv_projection_method,
- qkv_bias,
- attention_drop_rate,
- with_cls_token,
- name="attention",
- )
- self.dense_output = TFCvtSelfOutput(config, embed_dim, drop_rate, name="output")
-
- def prune_heads(self, heads):
- raise NotImplementedError
-
- def call(self, hidden_state: tf.Tensor, height: int, width: int, training: bool = False):
- self_output = self.attention(hidden_state, height, width, training=training)
- attention_output = self.dense_output(self_output, training=training)
- return attention_output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "dense_output", None) is not None:
- with tf.name_scope(self.dense_output.name):
- self.dense_output.build(None)
-
-
-class TFCvtIntermediate(keras.layers.Layer):
- """Intermediate dense layer. Second chunk of the convolutional transformer block."""
-
- def __init__(self, config: CvtConfig, embed_dim: int, mlp_ratio: int, **kwargs):
- super().__init__(**kwargs)
- self.dense = keras.layers.Dense(
- units=int(embed_dim * mlp_ratio),
- kernel_initializer=get_initializer(config.initializer_range),
- activation="gelu",
- name="dense",
- )
- self.embed_dim = embed_dim
-
- def call(self, hidden_state: tf.Tensor) -> tf.Tensor:
- hidden_state = self.dense(hidden_state)
- return hidden_state
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.embed_dim])
-
-
-class TFCvtOutput(keras.layers.Layer):
- """
- Output of the Convolutional Transformer Block (last chunk). It consists of a MLP and a residual connection.
- """
-
- def __init__(self, config: CvtConfig, embed_dim: int, mlp_ratio: int, drop_rate: int, **kwargs):
- super().__init__(**kwargs)
- self.dense = keras.layers.Dense(
- units=embed_dim, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.dropout = keras.layers.Dropout(drop_rate)
- self.embed_dim = embed_dim
- self.mlp_ratio = mlp_ratio
-
- def call(self, hidden_state: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_state = self.dense(inputs=hidden_state)
- hidden_state = self.dropout(inputs=hidden_state, training=training)
- hidden_state = hidden_state + input_tensor
- return hidden_state
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, int(self.embed_dim * self.mlp_ratio)])
-
-
-class TFCvtLayer(keras.layers.Layer):
- """
- Convolutional Transformer Block composed by attention layers, normalization and multi-layer perceptrons (mlps). It
- consists of 3 chunks : an attention layer, an intermediate dense layer and an output layer. This corresponds to the
- `Block` class in the original implementation.
- """
-
- def __init__(
- self,
- config: CvtConfig,
- num_heads: int,
- embed_dim: int,
- kernel_size: int,
- stride_q: int,
- stride_kv: int,
- padding_q: int,
- padding_kv: int,
- qkv_projection_method: str,
- qkv_bias: bool,
- attention_drop_rate: float,
- drop_rate: float,
- mlp_ratio: float,
- drop_path_rate: float,
- with_cls_token: bool = True,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.attention = TFCvtAttention(
- config,
- num_heads,
- embed_dim,
- kernel_size,
- stride_q,
- stride_kv,
- padding_q,
- padding_kv,
- qkv_projection_method,
- qkv_bias,
- attention_drop_rate,
- drop_rate,
- with_cls_token,
- name="attention",
- )
- self.intermediate = TFCvtIntermediate(config, embed_dim, mlp_ratio, name="intermediate")
- self.dense_output = TFCvtOutput(config, embed_dim, mlp_ratio, drop_rate, name="output")
- # Using `layers.Activation` instead of `tf.identity` to better control `training` behaviour.
- self.drop_path = (
- TFCvtDropPath(drop_path_rate, name="drop_path")
- if drop_path_rate > 0.0
- else keras.layers.Activation("linear", name="drop_path")
- )
- # Using the same default epsilon as PyTorch
- self.layernorm_before = keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_before")
- self.layernorm_after = keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_after")
- self.embed_dim = embed_dim
-
- def call(self, hidden_state: tf.Tensor, height: int, width: int, training: bool = False) -> tf.Tensor:
- # in Cvt, layernorm is applied before self-attention
- attention_output = self.attention(self.layernorm_before(hidden_state), height, width, training=training)
- attention_output = self.drop_path(attention_output, training=training)
-
- # first residual connection
- hidden_state = attention_output + hidden_state
-
- # in Cvt, layernorm is also applied after self-attention
- layer_output = self.layernorm_after(hidden_state)
- layer_output = self.intermediate(layer_output)
-
- # second residual connection is done here
- layer_output = self.dense_output(layer_output, hidden_state)
- layer_output = self.drop_path(layer_output, training=training)
- return layer_output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "intermediate", None) is not None:
- with tf.name_scope(self.intermediate.name):
- self.intermediate.build(None)
- if getattr(self, "dense_output", None) is not None:
- with tf.name_scope(self.dense_output.name):
- self.dense_output.build(None)
- if getattr(self, "drop_path", None) is not None:
- with tf.name_scope(self.drop_path.name):
- self.drop_path.build(None)
- if getattr(self, "layernorm_before", None) is not None:
- with tf.name_scope(self.layernorm_before.name):
- self.layernorm_before.build([None, None, self.embed_dim])
- if getattr(self, "layernorm_after", None) is not None:
- with tf.name_scope(self.layernorm_after.name):
- self.layernorm_after.build([None, None, self.embed_dim])
-
-
-class TFCvtStage(keras.layers.Layer):
- """
- Cvt stage (encoder block). Each stage has 2 parts :
- - (1) A Convolutional Token Embedding layer
- - (2) A Convolutional Transformer Block (layer).
- The classification token is added only in the last stage.
-
- Args:
- config ([`CvtConfig`]): Model configuration class.
- stage (`int`): Stage number.
- """
-
- def __init__(self, config: CvtConfig, stage: int, **kwargs):
- super().__init__(**kwargs)
- self.config = config
- self.stage = stage
- if self.config.cls_token[self.stage]:
- self.cls_token = self.add_weight(
- shape=(1, 1, self.config.embed_dim[-1]),
- initializer=get_initializer(self.config.initializer_range),
- trainable=True,
- name="cvt.encoder.stages.2.cls_token",
- )
-
- self.embedding = TFCvtEmbeddings(
- self.config,
- patch_size=config.patch_sizes[self.stage],
- num_channels=config.num_channels if self.stage == 0 else config.embed_dim[self.stage - 1],
- stride=config.patch_stride[self.stage],
- embed_dim=config.embed_dim[self.stage],
- padding=config.patch_padding[self.stage],
- dropout_rate=config.drop_rate[self.stage],
- name="embedding",
- )
-
- drop_path_rates = tf.linspace(0.0, config.drop_path_rate[self.stage], config.depth[stage])
- drop_path_rates = [x.numpy().item() for x in drop_path_rates]
- self.layers = [
- TFCvtLayer(
- config,
- num_heads=config.num_heads[self.stage],
- embed_dim=config.embed_dim[self.stage],
- kernel_size=config.kernel_qkv[self.stage],
- stride_q=config.stride_q[self.stage],
- stride_kv=config.stride_kv[self.stage],
- padding_q=config.padding_q[self.stage],
- padding_kv=config.padding_kv[self.stage],
- qkv_projection_method=config.qkv_projection_method[self.stage],
- qkv_bias=config.qkv_bias[self.stage],
- attention_drop_rate=config.attention_drop_rate[self.stage],
- drop_rate=config.drop_rate[self.stage],
- mlp_ratio=config.mlp_ratio[self.stage],
- drop_path_rate=drop_path_rates[self.stage],
- with_cls_token=config.cls_token[self.stage],
- name=f"layers.{j}",
- )
- for j in range(config.depth[self.stage])
- ]
-
- def call(self, hidden_state: tf.Tensor, training: bool = False):
- cls_token = None
- hidden_state = self.embedding(hidden_state, training)
-
- # "batch_size, height, width, num_channels -> batch_size, (height*width), num_channels"
- batch_size, height, width, num_channels = shape_list(hidden_state)
- hidden_size = height * width
- hidden_state = tf.reshape(hidden_state, shape=(batch_size, hidden_size, num_channels))
-
- if self.config.cls_token[self.stage]:
- cls_token = tf.repeat(self.cls_token, repeats=batch_size, axis=0)
- hidden_state = tf.concat((cls_token, hidden_state), axis=1)
-
- for layer in self.layers:
- layer_outputs = layer(hidden_state, height, width, training=training)
- hidden_state = layer_outputs
-
- if self.config.cls_token[self.stage]:
- cls_token, hidden_state = tf.split(hidden_state, [1, height * width], 1)
-
- # "batch_size, (height*width), num_channels -> batch_size, height, width, num_channels"
- hidden_state = tf.reshape(hidden_state, shape=(batch_size, height, width, num_channels))
- return hidden_state, cls_token
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embedding", None) is not None:
- with tf.name_scope(self.embedding.name):
- self.embedding.build(None)
- if getattr(self, "layers", None) is not None:
- for layer in self.layers:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-class TFCvtEncoder(keras.layers.Layer):
- """
- Convolutional Vision Transformer encoder. CVT has 3 stages of encoder blocks with their respective number of layers
- (depth) being 1, 2 and 10.
-
- Args:
- config ([`CvtConfig`]): Model configuration class.
- """
-
- config_class = CvtConfig
-
- def __init__(self, config: CvtConfig, **kwargs):
- super().__init__(**kwargs)
- self.config = config
- self.stages = [
- TFCvtStage(config, stage_idx, name=f"stages.{stage_idx}") for stage_idx in range(len(config.depth))
- ]
-
- def call(
- self,
- pixel_values: TFModelInputType,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithCLSToken, Tuple[tf.Tensor]]:
- all_hidden_states = () if output_hidden_states else None
- hidden_state = pixel_values
- # When running on CPU, `keras.layers.Conv2D` doesn't support (batch_size, num_channels, height, width)
- # as input format. So change the input format to (batch_size, height, width, num_channels).
- hidden_state = tf.transpose(hidden_state, perm=(0, 2, 3, 1))
-
- cls_token = None
- for _, (stage_module) in enumerate(self.stages):
- hidden_state, cls_token = stage_module(hidden_state, training=training)
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_state,)
-
- # Change back to (batch_size, num_channels, height, width) format to have uniformity in the modules
- hidden_state = tf.transpose(hidden_state, perm=(0, 3, 1, 2))
- if output_hidden_states:
- all_hidden_states = tuple([tf.transpose(hs, perm=(0, 3, 1, 2)) for hs in all_hidden_states])
-
- if not return_dict:
- return tuple(v for v in [hidden_state, cls_token, all_hidden_states] if v is not None)
-
- return TFBaseModelOutputWithCLSToken(
- last_hidden_state=hidden_state,
- cls_token_value=cls_token,
- hidden_states=all_hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "stages", None) is not None:
- for layer in self.stages:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-@keras_serializable
-class TFCvtMainLayer(keras.layers.Layer):
- """Construct the Cvt model."""
-
- config_class = CvtConfig
-
- def __init__(self, config: CvtConfig, **kwargs):
- super().__init__(**kwargs)
- self.config = config
- self.encoder = TFCvtEncoder(config, name="encoder")
-
- @unpack_inputs
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithCLSToken, Tuple[tf.Tensor]]:
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- encoder_outputs = self.encoder(
- pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = encoder_outputs[0]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[1:]
-
- return TFBaseModelOutputWithCLSToken(
- last_hidden_state=sequence_output,
- cls_token_value=encoder_outputs.cls_token_value,
- hidden_states=encoder_outputs.hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
-
-
-class TFCvtPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = CvtConfig
- base_model_prefix = "cvt"
- main_input_name = "pixel_values"
-
-
-TFCVT_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TF 2.0 models accepts two formats as inputs:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional arguments.
-
- This second option is useful when using [`keras.Model.fit`] method which currently requires having all the
- tensors in the first argument of the model call function: `model(inputs)`.
-
-
-
- Args:
- config ([`CvtConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-TFCVT_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`np.ndarray`, `tf.Tensor`, `List[tf.Tensor]` ``Dict[str, tf.Tensor]` or `Dict[str, np.ndarray]` and each example must have the shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`CvtImageProcessor.__call__`]
- for details.
-
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False``):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare Cvt Model transformer outputting raw hidden-states without any specific head on top.",
- TFCVT_START_DOCSTRING,
-)
-class TFCvtModel(TFCvtPreTrainedModel):
- def __init__(self, config: CvtConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.cvt = TFCvtMainLayer(config, name="cvt")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(TFCVT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFBaseModelOutputWithCLSToken, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithCLSToken, Tuple[tf.Tensor]]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, TFCvtModel
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/cvt-13")
- >>> model = TFCvtModel.from_pretrained("microsoft/cvt-13")
-
- >>> inputs = image_processor(images=image, return_tensors="tf")
- >>> outputs = model(**inputs)
- >>> last_hidden_states = outputs.last_hidden_state
- ```"""
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- outputs = self.cvt(
- pixel_values=pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- if not return_dict:
- return (outputs[0],) + outputs[1:]
-
- return TFBaseModelOutputWithCLSToken(
- last_hidden_state=outputs.last_hidden_state,
- cls_token_value=outputs.cls_token_value,
- hidden_states=outputs.hidden_states,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "cvt", None) is not None:
- with tf.name_scope(self.cvt.name):
- self.cvt.build(None)
-
-
-@add_start_docstrings(
- """
- Cvt Model transformer with an image classification head on top (a linear layer on top of the final hidden state of
- the [CLS] token) e.g. for ImageNet.
- """,
- TFCVT_START_DOCSTRING,
-)
-class TFCvtForImageClassification(TFCvtPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: CvtConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
- self.cvt = TFCvtMainLayer(config, name="cvt")
- # Using same default epsilon as in the original implementation.
- self.layernorm = keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm")
-
- # Classifier head
- self.classifier = keras.layers.Dense(
- units=config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- use_bias=True,
- bias_initializer="zeros",
- name="classifier",
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(TFCVT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFImageClassifierOutputWithNoAttention, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- labels: tf.Tensor | None = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFImageClassifierOutputWithNoAttention, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, TFCvtForImageClassification
- >>> import tensorflow as tf
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/cvt-13")
- >>> model = TFCvtForImageClassification.from_pretrained("microsoft/cvt-13")
-
- >>> inputs = image_processor(images=image, return_tensors="tf")
- >>> outputs = model(**inputs)
- >>> logits = outputs.logits
- >>> # model predicts one of the 1000 ImageNet classes
- >>> predicted_class_idx = tf.math.argmax(logits, axis=-1)[0]
- >>> print("Predicted class:", model.config.id2label[int(predicted_class_idx)])
- ```"""
-
- outputs = self.cvt(
- pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = outputs[0]
- cls_token = outputs[1]
- if self.config.cls_token[-1]:
- sequence_output = self.layernorm(cls_token)
- else:
- # rearrange "batch_size, num_channels, height, width -> batch_size, (height*width), num_channels"
- batch_size, num_channels, height, width = shape_list(sequence_output)
- sequence_output = tf.reshape(sequence_output, shape=(batch_size, num_channels, height * width))
- sequence_output = tf.transpose(sequence_output, perm=(0, 2, 1))
- sequence_output = self.layernorm(sequence_output)
-
- sequence_output_mean = tf.reduce_mean(sequence_output, axis=1)
- logits = self.classifier(sequence_output_mean)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "cvt", None) is not None:
- with tf.name_scope(self.cvt.name):
- self.cvt.build(None)
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, None, self.config.embed_dim[-1]])
- if getattr(self, "classifier", None) is not None:
- if hasattr(self.classifier, "name"):
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.embed_dim[-1]])
diff --git a/transformers/models/data2vec/__init__.py b/transformers/models/data2vec/__init__.py
deleted file mode 100644
index 45522f4ba893a154b3400b76b4bb280fd00b692a..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/__init__.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available
-
-
-_import_structure = {
- "configuration_data2vec_audio": ["DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP", "Data2VecAudioConfig"],
- "configuration_data2vec_text": [
- "DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "Data2VecTextConfig",
- "Data2VecTextOnnxConfig",
- ],
- "configuration_data2vec_vision": [
- "DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "Data2VecVisionConfig",
- "Data2VecVisionOnnxConfig",
- ],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_data2vec_audio"] = [
- "DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST",
- "Data2VecAudioForAudioFrameClassification",
- "Data2VecAudioForCTC",
- "Data2VecAudioForSequenceClassification",
- "Data2VecAudioForXVector",
- "Data2VecAudioModel",
- "Data2VecAudioPreTrainedModel",
- ]
- _import_structure["modeling_data2vec_text"] = [
- "DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "Data2VecTextForCausalLM",
- "Data2VecTextForMaskedLM",
- "Data2VecTextForMultipleChoice",
- "Data2VecTextForQuestionAnswering",
- "Data2VecTextForSequenceClassification",
- "Data2VecTextForTokenClassification",
- "Data2VecTextModel",
- "Data2VecTextPreTrainedModel",
- ]
- _import_structure["modeling_data2vec_vision"] = [
- "DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST",
- "Data2VecVisionForImageClassification",
- "Data2VecVisionForMaskedImageModeling",
- "Data2VecVisionForSemanticSegmentation",
- "Data2VecVisionModel",
- "Data2VecVisionPreTrainedModel",
- ]
-
-if is_tf_available():
- _import_structure["modeling_tf_data2vec_vision"] = [
- "TFData2VecVisionForImageClassification",
- "TFData2VecVisionForSemanticSegmentation",
- "TFData2VecVisionModel",
- "TFData2VecVisionPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_data2vec_audio import DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP, Data2VecAudioConfig
- from .configuration_data2vec_text import (
- DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP,
- Data2VecTextConfig,
- Data2VecTextOnnxConfig,
- )
- from .configuration_data2vec_vision import (
- DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP,
- Data2VecVisionConfig,
- Data2VecVisionOnnxConfig,
- )
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_data2vec_audio import (
- DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST,
- Data2VecAudioForAudioFrameClassification,
- Data2VecAudioForCTC,
- Data2VecAudioForSequenceClassification,
- Data2VecAudioForXVector,
- Data2VecAudioModel,
- Data2VecAudioPreTrainedModel,
- )
- from .modeling_data2vec_text import (
- DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
- Data2VecTextForCausalLM,
- Data2VecTextForMaskedLM,
- Data2VecTextForMultipleChoice,
- Data2VecTextForQuestionAnswering,
- Data2VecTextForSequenceClassification,
- Data2VecTextForTokenClassification,
- Data2VecTextModel,
- Data2VecTextPreTrainedModel,
- )
- from .modeling_data2vec_vision import (
- DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST,
- Data2VecVisionForImageClassification,
- Data2VecVisionForMaskedImageModeling,
- Data2VecVisionForSemanticSegmentation,
- Data2VecVisionModel,
- Data2VecVisionPreTrainedModel,
- )
- if is_tf_available():
- from .modeling_tf_data2vec_vision import (
- TFData2VecVisionForImageClassification,
- TFData2VecVisionForSemanticSegmentation,
- TFData2VecVisionModel,
- TFData2VecVisionPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/data2vec/__pycache__/__init__.cpython-310.pyc b/transformers/models/data2vec/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 0d8ef2713bc4f06c362f8011f6b9f4121ad80d46..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/configuration_data2vec_audio.cpython-310.pyc b/transformers/models/data2vec/__pycache__/configuration_data2vec_audio.cpython-310.pyc
deleted file mode 100644
index bc6a64cc4d89915e1c3584a9ae72f96cd5f9e02d..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/configuration_data2vec_audio.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/configuration_data2vec_text.cpython-310.pyc b/transformers/models/data2vec/__pycache__/configuration_data2vec_text.cpython-310.pyc
deleted file mode 100644
index a5ff3f20a7823e6390c5120a5fedb3e3eb867aa9..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/configuration_data2vec_text.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/configuration_data2vec_vision.cpython-310.pyc b/transformers/models/data2vec/__pycache__/configuration_data2vec_vision.cpython-310.pyc
deleted file mode 100644
index 03c042519829f4e08984ba35e43c74bdd5bd64bb..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/configuration_data2vec_vision.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/convert_data2vec_audio_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/data2vec/__pycache__/convert_data2vec_audio_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 936e52a28b6bed500ebba174b512b11aff863c6c..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/convert_data2vec_audio_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/convert_data2vec_text_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/data2vec/__pycache__/convert_data2vec_text_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 0c3d174510fe6bfbba68aa285c13aebb5ef4af83..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/convert_data2vec_text_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/data2vec/__pycache__/convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index ae5a40fd82f149b9a96f40ad48397d88eb39ae93..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/modeling_data2vec_audio.cpython-310.pyc b/transformers/models/data2vec/__pycache__/modeling_data2vec_audio.cpython-310.pyc
deleted file mode 100644
index b3b0ac240a556a5d93cd416b05c62cae64b8468c..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/modeling_data2vec_audio.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/modeling_data2vec_text.cpython-310.pyc b/transformers/models/data2vec/__pycache__/modeling_data2vec_text.cpython-310.pyc
deleted file mode 100644
index e3face458e2d77851f91f7cec29cb5bec0de5c6f..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/modeling_data2vec_text.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/modeling_data2vec_vision.cpython-310.pyc b/transformers/models/data2vec/__pycache__/modeling_data2vec_vision.cpython-310.pyc
deleted file mode 100644
index 4b719af75beea53004d28e89c923470c3a0d9865..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/modeling_data2vec_vision.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/__pycache__/modeling_tf_data2vec_vision.cpython-310.pyc b/transformers/models/data2vec/__pycache__/modeling_tf_data2vec_vision.cpython-310.pyc
deleted file mode 100644
index 4317c35a02d68fbc7a3afabc6fc539b698bc8823..0000000000000000000000000000000000000000
Binary files a/transformers/models/data2vec/__pycache__/modeling_tf_data2vec_vision.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/data2vec/configuration_data2vec_audio.py b/transformers/models/data2vec/configuration_data2vec_audio.py
deleted file mode 100644
index 32d505f157d63f628fc10c5226b0c823e843fbb8..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/configuration_data2vec_audio.py
+++ /dev/null
@@ -1,285 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Data2VecText configuration"""
-
-import math
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class Data2VecAudioConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`Data2VecAudioModel`]. It is used to instantiate
- an Data2VecAudio model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the Data2VecAudio
- [facebook/data2vec-audio-base-960h](https://huggingface.co/facebook/data2vec-audio-base-960h) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 32):
- Vocabulary size of the Data2VecAudio model. Defines the number of different tokens that can be represented
- by the `inputs_ids` passed when calling [`Data2VecAudioModel`] or [`TFData2VecAudioModel`]. Vocabulary size
- of the model. Defines the different tokens that can be represented by the *inputs_ids* passed to the
- forward method of [`Data2VecAudioModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- activation_dropout (`float`, *optional*, defaults to 0.1):
- The dropout ratio for activations inside the fully connected layer.
- attention_dropout (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- final_dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for the final projection layer of [`Data2VecAudioForCTC`].
- layerdrop (`float`, *optional*, defaults to 0.1):
- The LayerDrop probability. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556) for more
- details.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- feat_proj_dropout (`float`, *optional*, defaults to 0.0):
- The dropout probability for output of the feature encoder.
- feat_extract_activation (`str, `optional`, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the 1D convolutional layers of the feature
- extractor. If string, `"gelu"`, `"relu"`, `"selu"` and `"gelu_new"` are supported.
- conv_dim (`Tuple[int]` or `List[int]`, *optional*, defaults to `(512, 512, 512, 512, 512, 512, 512)`):
- A tuple of integers defining the number of input and output channels of each 1D convolutional layer in the
- feature encoder. The length of *conv_dim* defines the number of 1D convolutional layers.
- conv_stride (`Tuple[int]` or `List[int]`, *optional*, defaults to `(5, 2, 2, 2, 2, 2, 2)`):
- A tuple of integers defining the stride of each 1D convolutional layer in the feature encoder. The length
- of *conv_stride* defines the number of convolutional layers and has to match the length of *conv_dim*.
- conv_kernel (`Tuple[int]` or `List[int]`, *optional*, defaults to `(10, 3, 3, 3, 3, 3, 3)`):
- A tuple of integers defining the kernel size of each 1D convolutional layer in the feature encoder. The
- length of *conv_kernel* defines the number of convolutional layers and has to match the length of
- *conv_dim*.
- conv_bias (`bool`, *optional*, defaults to `False`):
- Whether the 1D convolutional layers have a bias.
- num_conv_pos_embeddings (`int`, *optional*, defaults to 128):
- Number of convolutional positional embeddings. Defines the kernel size of 1D convolutional positional
- embeddings layer.
- num_conv_pos_embedding_groups (`int`, *optional*, defaults to 16):
- Number of groups of 1D convolutional positional embeddings layer.
- mask_time_prob (`float`, *optional*, defaults to 0.05):
- Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked. The masking
- procecure generates ''mask_time_prob*len(time_axis)/mask_time_length'' independent masks over the axis. If
- reasoning from the propability of each feature vector to be chosen as the start of the vector span to be
- masked, *mask_time_prob* should be `prob_vector_start*mask_time_length`. Note that overlap may decrease the
- mask_time_length (`int`, *optional*, defaults to 10):
- Length of vector span along the time axis.
- mask_time_min_masks (`int`, *optional*, defaults to 2),:
- The minimum number of masks of length `mask_feature_length` generated along the time axis, each time step,
- irrespectively of `mask_feature_prob`. Only relevant if ''mask_time_prob*len(time_axis)/mask_time_length <
- mask_time_min_masks''
- mask_feature_prob (`float`, *optional*, defaults to 0.0):
- Percentage (between 0 and 1) of all feature vectors along the feature axis which will be masked. The
- masking procecure generates ''mask_feature_prob*len(feature_axis)/mask_time_length'' independent masks over
- the axis. If reasoning from the propability of each feature vector to be chosen as the start of the vector
- span to be masked, *mask_feature_prob* should be `prob_vector_start*mask_feature_length`. Note that overlap
- may decrease the actual percentage of masked vectors. This is only relevant if `apply_spec_augment is
- True`.
- mask_feature_length (`int`, *optional*, defaults to 10):
- Length of vector span along the feature axis.
- mask_feature_min_masks (`int`, *optional*, defaults to 0),:
- The minimum number of masks of length `mask_feature_length` generated along the feature axis, each time
- step, irrespectively of `mask_feature_prob`. Only relevant if
- ''mask_feature_prob*len(feature_axis)/mask_feature_length < mask_feature_min_masks''
- ctc_loss_reduction (`str`, *optional*, defaults to `"sum"`):
- Specifies the reduction to apply to the output of `torch.nn.CTCLoss`. Only relevant when training an
- instance of [`Data2VecAudioForCTC`].
- ctc_zero_infinity (`bool`, *optional*, defaults to `False`):
- Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly
- occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance
- of [`Data2VecAudioForCTC`].
- use_weighted_layer_sum (`bool`, *optional*, defaults to `False`):
- Whether to use a weighted average of layer outputs with learned weights. Only relevant when using an
- instance of [`Data2VecAudioForSequenceClassification`].
- classifier_proj_size (`int`, *optional*, defaults to 256):
- Dimensionality of the projection before token mean-pooling for classification.
- tdnn_dim (`Tuple[int]` or `List[int]`, *optional*, defaults to `(512, 512, 512, 512, 1500)`):
- A tuple of integers defining the number of output channels of each 1D convolutional layer in the *TDNN*
- module of the *XVector* model. The length of *tdnn_dim* defines the number of *TDNN* layers.
- tdnn_kernel (`Tuple[int]` or `List[int]`, *optional*, defaults to `(5, 3, 3, 1, 1)`):
- A tuple of integers defining the kernel size of each 1D convolutional layer in the *TDNN* module of the
- *XVector* model. The length of *tdnn_kernel* has to match the length of *tdnn_dim*.
- tdnn_dilation (`Tuple[int]` or `List[int]`, *optional*, defaults to `(1, 2, 3, 1, 1)`):
- A tuple of integers defining the dilation factor of each 1D convolutional layer in *TDNN* module of the
- *XVector* model. The length of *tdnn_dilation* has to match the length of *tdnn_dim*.
- xvector_output_dim (`int`, *optional*, defaults to 512):
- Dimensionality of the *XVector* embedding vectors.
- add_adapter (`bool`, *optional*, defaults to `False`):
- Whether a convolutional network should be stacked on top of the Data2VecAudio Encoder. Can be very useful
- for warm-starting Data2VecAudio for SpeechEncoderDecoder models.
- adapter_kernel_size (`int`, *optional*, defaults to 3):
- Kernel size of the convolutional layers in the adapter network. Only relevant if `add_adapter is True`.
- adapter_stride (`int`, *optional*, defaults to 2):
- Stride of the convolutional layers in the adapter network. Only relevant if `add_adapter is True`.
- num_adapter_layers (`int`, *optional*, defaults to 3):
- Number of convolutional layers that should be used in the adapter network. Only relevant if `add_adapter is
- True`.
- output_hidden_size (`int`, *optional*):
- Dimensionality of the encoder output layer. If not defined, this defaults to *hidden-size*. Only relevant
- if `add_adapter is True`.
-
- Example:
-
- ```python
- >>> from transformers import Data2VecAudioConfig, Data2VecAudioModel
-
- >>> # Initializing a Data2VecAudio facebook/data2vec-audio-base-960h style configuration
- >>> configuration = Data2VecAudioConfig()
-
- >>> # Initializing a model (with random weights) from the facebook/data2vec-audio-base-960h style configuration
- >>> model = Data2VecAudioModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "data2vec-audio"
-
- def __init__(
- self,
- vocab_size=32,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout=0.1,
- activation_dropout=0.1,
- attention_dropout=0.1,
- feat_proj_dropout=0.0,
- final_dropout=0.1,
- layerdrop=0.1,
- initializer_range=0.02,
- layer_norm_eps=1e-5,
- feat_extract_activation="gelu",
- conv_dim=(512, 512, 512, 512, 512, 512, 512),
- conv_stride=(5, 2, 2, 2, 2, 2, 2),
- conv_kernel=(10, 3, 3, 3, 3, 2, 2),
- conv_bias=False,
- num_conv_pos_embedding_groups=16,
- conv_pos_kernel_size=19,
- num_conv_pos_embeddings=5,
- mask_time_prob=0.05,
- mask_time_length=10,
- mask_time_min_masks=2,
- mask_feature_prob=0.0,
- mask_feature_length=10,
- mask_feature_min_masks=0,
- ctc_loss_reduction="sum",
- ctc_zero_infinity=False,
- use_weighted_layer_sum=False,
- classifier_proj_size=256,
- tdnn_dim=(512, 512, 512, 512, 1500),
- tdnn_kernel=(5, 3, 3, 1, 1),
- tdnn_dilation=(1, 2, 3, 1, 1),
- xvector_output_dim=512,
- pad_token_id=0,
- bos_token_id=1,
- eos_token_id=2,
- add_adapter=False,
- adapter_kernel_size=3,
- adapter_stride=2,
- num_adapter_layers=3,
- output_hidden_size=None,
- **kwargs,
- ):
- super().__init__(**kwargs, pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id)
- self.hidden_size = hidden_size
- self.feat_extract_activation = feat_extract_activation
- self.conv_dim = list(conv_dim)
- self.conv_stride = list(conv_stride)
- self.conv_kernel = list(conv_kernel)
- self.conv_bias = conv_bias
- self.num_conv_pos_embeddings = num_conv_pos_embeddings
- self.num_conv_pos_embedding_groups = num_conv_pos_embedding_groups
- self.conv_pos_kernel_size = conv_pos_kernel_size
- self.num_feat_extract_layers = len(self.conv_dim)
- self.num_hidden_layers = num_hidden_layers
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.num_attention_heads = num_attention_heads
- self.hidden_dropout = hidden_dropout
- self.attention_dropout = attention_dropout
- self.activation_dropout = activation_dropout
- self.feat_proj_dropout = feat_proj_dropout
- self.final_dropout = final_dropout
- self.layerdrop = layerdrop
- self.layer_norm_eps = layer_norm_eps
- self.initializer_range = initializer_range
- self.vocab_size = vocab_size
- self.use_weighted_layer_sum = use_weighted_layer_sum
-
- if (
- (len(self.conv_stride) != self.num_feat_extract_layers)
- or (len(self.conv_kernel) != self.num_feat_extract_layers)
- or (len(self.conv_dim) != self.num_feat_extract_layers)
- ):
- raise ValueError(
- "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="
- " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="
- f" {len(self.conv_dim)}`, `len(config.conv_stride) = {len(self.conv_stride)}`,"
- f" `len(config.conv_kernel) = {len(self.conv_kernel)}`."
- )
-
- # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
- self.mask_time_prob = mask_time_prob
- self.mask_time_length = mask_time_length
- self.mask_time_min_masks = mask_time_min_masks
- self.mask_feature_prob = mask_feature_prob
- self.mask_feature_length = mask_feature_length
- self.mask_feature_min_masks = mask_feature_min_masks
-
- # ctc loss
- self.ctc_loss_reduction = ctc_loss_reduction
- self.ctc_zero_infinity = ctc_zero_infinity
-
- # adapter
- self.add_adapter = add_adapter
- self.adapter_kernel_size = adapter_kernel_size
- self.adapter_stride = adapter_stride
- self.num_adapter_layers = num_adapter_layers
- self.output_hidden_size = output_hidden_size or hidden_size
-
- # SequenceClassification-specific parameter. Feel free to ignore for other classes.
- self.classifier_proj_size = classifier_proj_size
-
- # XVector-specific parameters. Feel free to ignore for other classes.
- self.tdnn_dim = list(tdnn_dim)
- self.tdnn_kernel = list(tdnn_kernel)
- self.tdnn_dilation = list(tdnn_dilation)
- self.xvector_output_dim = xvector_output_dim
-
- @property
- def inputs_to_logits_ratio(self):
- return math.prod(self.conv_stride)
diff --git a/transformers/models/data2vec/configuration_data2vec_text.py b/transformers/models/data2vec/configuration_data2vec_text.py
deleted file mode 100644
index cd52db2d326e9f5a9f7e6392815e5f63185352af..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/configuration_data2vec_text.py
+++ /dev/null
@@ -1,153 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Data2VecText configuration"""
-from collections import OrderedDict
-from typing import Mapping
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class Data2VecTextConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`Data2VecTextModel`] and [`Data2VecTextModel`]. It
- is used to instantiate a Data2VecText model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the Data2VecText
- [facebook/data2vec-text-base](https://huggingface.co/facebook/data2vec-text-base) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the DATA2VEC model. Defines the number of different tokens that can be represented by
- the `inputs_ids` passed when calling [`Data2VecModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 2):
- The vocabulary size of the `token_type_ids` passed when calling [`Data2VecModel`].
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
- Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
- positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
- [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
- For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
- with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
- is_decoder (`bool`, *optional*, defaults to `False`):
- Whether the model is used as a decoder or not. If `False`, the model is used as an encoder.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- classifier_dropout (`float`, *optional*):
- The dropout ratio for the classification head.
-
- Examples:
-
- ```python
- >>> from transformers import Data2VecTextConfig, Data2VecTextModel
-
- >>> # Initializing a Data2VecText facebook/data2vec-text-base style configuration
- >>> configuration = Data2VecTextConfig()
-
- >>> # Initializing a model (with random weights) from the facebook/data2vec-text-base style configuration
- >>> model = Data2VecTextModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "data2vec-text"
-
- def __init__(
- self,
- vocab_size=30522,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=2,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- pad_token_id=1,
- bos_token_id=0,
- eos_token_id=2,
- position_embedding_type="absolute",
- use_cache=True,
- classifier_dropout=None,
- **kwargs,
- ):
- super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
-
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.hidden_act = hidden_act
- self.intermediate_size = intermediate_size
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.position_embedding_type = position_embedding_type
- self.use_cache = use_cache
- self.classifier_dropout = classifier_dropout
-
-
-class Data2VecTextOnnxConfig(OnnxConfig):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- if self.task == "multiple-choice":
- dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
- else:
- dynamic_axis = {0: "batch", 1: "sequence"}
- return OrderedDict(
- [
- ("input_ids", dynamic_axis),
- ("attention_mask", dynamic_axis),
- ]
- )
diff --git a/transformers/models/data2vec/configuration_data2vec_vision.py b/transformers/models/data2vec/configuration_data2vec_vision.py
deleted file mode 100644
index 9a9de9c4be5a0dc10dc35598544f3baed29cda62..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/configuration_data2vec_vision.py
+++ /dev/null
@@ -1,193 +0,0 @@
-# coding=utf-8
-# Copyright Meta Platforms and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Data2VecVision model configuration"""
-from collections import OrderedDict
-from typing import Mapping
-
-from packaging import version
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class Data2VecVisionConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`Data2VecVisionModel`]. It is used to instantiate
- an Data2VecVision model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the Data2VecVision
- [facebook/data2vec-vision-base](https://huggingface.co/facebook/data2vec-vision-base) architecture.
-
- Args:
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 16):
- The size (resolution) of each patch.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- use_mask_token (`bool`, *optional*, defaults to `False`):
- Whether to use a mask token for masked image modeling.
- use_absolute_position_embeddings (`bool`, *optional*, defaults to `False`):
- Whether to use BERT-style absolute position embeddings.
- use_relative_position_bias (`bool`, *optional*, defaults to `False`):
- Whether to use T5-style relative position embeddings in the self-attention layers.
- use_shared_relative_position_bias (`bool`, *optional*, defaults to `False`):
- Whether to use the same relative position embeddings across all self-attention layers of the Transformer.
- layer_scale_init_value (`float`, *optional*, defaults to 0.1):
- Scale to use in the self-attention layers. 0.1 for base, 1e-5 for large. Set 0 to disable layer scale.
- drop_path_rate (`float`, *optional*, defaults to 0.1):
- Stochastic depth rate per sample (when applied in the main path of residual layers).
- use_mean_pooling (`bool`, *optional*, defaults to `True`):
- Whether to mean pool the final hidden states of the patches instead of using the final hidden state of the
- CLS token, before applying the classification head.
- out_indices (`List[int]`, *optional*, defaults to `[3, 5, 7, 11]`):
- Indices of the feature maps to use for semantic segmentation.
- pool_scales (`Tuple[int]`, *optional*, defaults to `[1, 2, 3, 6]`):
- Pooling scales used in Pooling Pyramid Module applied on the last feature map.
- use_auxiliary_head (`bool`, *optional*, defaults to `True`):
- Whether to use an auxiliary head during training.
- auxiliary_loss_weight (`float`, *optional*, defaults to 0.4):
- Weight of the cross-entropy loss of the auxiliary head.
- auxiliary_channels (`int`, *optional*, defaults to 256):
- Number of channels to use in the auxiliary head.
- auxiliary_num_convs (`int`, *optional*, defaults to 1):
- Number of convolutional layers to use in the auxiliary head.
- auxiliary_concat_input (`bool`, *optional*, defaults to `False`):
- Whether to concatenate the output of the auxiliary head with the input before the classification layer.
- semantic_loss_ignore_index (`int`, *optional*, defaults to 255):
- The index that is ignored by the loss function of the semantic segmentation model.
-
- Example:
-
- ```python
- >>> from transformers import Data2VecVisionConfig, Data2VecVisionModel
-
- >>> # Initializing a Data2VecVision data2vec_vision-base-patch16-224-in22k style configuration
- >>> configuration = Data2VecVisionConfig()
-
- >>> # Initializing a model (with random weights) from the data2vec_vision-base-patch16-224-in22k style configuration
- >>> model = Data2VecVisionModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "data2vec-vision"
-
- def __init__(
- self,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.0,
- attention_probs_dropout_prob=0.0,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- image_size=224,
- patch_size=16,
- num_channels=3,
- use_mask_token=False,
- use_absolute_position_embeddings=False,
- use_relative_position_bias=False,
- use_shared_relative_position_bias=False,
- layer_scale_init_value=0.1,
- drop_path_rate=0.1,
- use_mean_pooling=True,
- out_indices=[3, 5, 7, 11],
- pool_scales=[1, 2, 3, 6],
- use_auxiliary_head=True,
- auxiliary_loss_weight=0.4,
- auxiliary_channels=256,
- auxiliary_num_convs=1,
- auxiliary_concat_input=False,
- semantic_loss_ignore_index=255,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
-
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.use_mask_token = use_mask_token
- self.use_absolute_position_embeddings = use_absolute_position_embeddings
- self.use_relative_position_bias = use_relative_position_bias
- self.use_shared_relative_position_bias = use_shared_relative_position_bias
- self.layer_scale_init_value = layer_scale_init_value
- self.drop_path_rate = drop_path_rate
- self.use_mean_pooling = use_mean_pooling
- # decode head attributes (semantic segmentation)
- self.out_indices = out_indices
- self.pool_scales = pool_scales
- # auxiliary head attributes (semantic segmentation)
- self.use_auxiliary_head = use_auxiliary_head
- self.auxiliary_loss_weight = auxiliary_loss_weight
- self.auxiliary_channels = auxiliary_channels
- self.auxiliary_num_convs = auxiliary_num_convs
- self.auxiliary_concat_input = auxiliary_concat_input
- self.semantic_loss_ignore_index = semantic_loss_ignore_index
-
-
-# Copied from transformers.models.vit.configuration_vit.ViTOnnxConfig
-class Data2VecVisionOnnxConfig(OnnxConfig):
- torch_onnx_minimum_version = version.parse("1.11")
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- return OrderedDict(
- [
- ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
- ]
- )
-
- @property
- def atol_for_validation(self) -> float:
- return 1e-4
diff --git a/transformers/models/data2vec/convert_data2vec_audio_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/data2vec/convert_data2vec_audio_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index 01c2d8cab27894b8f6cc91572d3c9fdd55aafcab..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/convert_data2vec_audio_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,286 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Wav2Vec2 checkpoint."""
-
-
-import argparse
-import os
-from functools import reduce
-
-import fairseq
-import torch
-from datasets import load_dataset
-
-from transformers import Wav2Vec2Processor, logging
-from transformers.models.data2vec.configuration_data2vec_audio import Data2VecAudioConfig
-
-# Copied from https://github.com/pytorch/fairseq/blob/main/examples/data2vec/models/data2vec_audio.py
-from transformers.models.data2vec.data2vec_audio import Data2VecAudioModel as Dummy # noqa: F401
-from transformers.models.data2vec.modeling_data2vec_audio import Data2VecAudioForCTC, Data2VecAudioModel
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-MAPPING = {
- "post_extract_proj": "feature_projection.projection",
- "models.0.layer_norm": "feature_projection.layer_norm",
- "self_attn.k_proj": "encoder.layers.*.attention.k_proj",
- "self_attn.v_proj": "encoder.layers.*.attention.v_proj",
- "self_attn.q_proj": "encoder.layers.*.attention.q_proj",
- "self_attn.out_proj": "encoder.layers.*.attention.out_proj",
- "self_attn_layer_norm": "encoder.layers.*.layer_norm",
- "fc1": "encoder.layers.*.feed_forward.intermediate_dense",
- "fc2": "encoder.layers.*.feed_forward.output_dense",
- "final_layer_norm": "encoder.layers.*.final_layer_norm",
- "encoder.layer_norm": "encoder.layer_norm",
- "w2v_model.layer_norm": "feature_projection.layer_norm",
- "w2v_encoder.proj": "lm_head",
- "mask_emb": "masked_spec_embed",
-}
-TOP_LEVEL_KEYS = [
- "lm_head",
-]
-
-
-def set_recursively(hf_pointer, key, value, full_name, weight_type):
- for attribute in key.split("."):
- hf_pointer = getattr(hf_pointer, attribute)
-
- if weight_type is not None:
- hf_shape = getattr(hf_pointer, weight_type).shape
- else:
- hf_shape = hf_pointer.shape
-
- if hf_shape != value.shape:
- raise ValueError(
- f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
- f" {value.shape} for {full_name}"
- )
-
- if weight_type == "weight":
- hf_pointer.weight.data = value
- elif weight_type == "weight_g":
- hf_pointer.weight_g.data = value
- elif weight_type == "weight_v":
- hf_pointer.weight_v.data = value
- elif weight_type == "bias":
- hf_pointer.bias.data = value
- else:
- hf_pointer.data = value
-
- logger.info(f"{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.")
-
-
-def recursively_load_weights(fairseq_model, hf_model, is_headless):
- unused_weights = []
- fairseq_dict = fairseq_model.state_dict()
-
- if not is_headless:
- feature_extractor = hf_model.data2vec_audio.feature_extractor
- pos_conv_embedding = hf_model.data2vec_audio.encoder.pos_conv_embed
-
- else:
- feature_extractor = hf_model.feature_extractor
- pos_conv_embedding = hf_model.encoder.pos_conv_embed
-
- for name, value in fairseq_dict.items():
- is_used = False
- if "conv_layers" in name:
- load_conv_layer(
- name,
- value,
- feature_extractor,
- unused_weights,
- )
- is_used = True
- elif "pos_conv" in name:
- load_pos_conv_layer(
- name,
- value,
- pos_conv_embedding,
- unused_weights,
- )
- is_used = True
- else:
- for key, mapped_key in MAPPING.items():
- if not is_headless:
- mapped_key = "data2vec_audio." + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
- if key in name or key.split("w2v_model.")[-1] == name.split(".")[0]:
- is_used = True
- if "*" in mapped_key:
- layer_index = name.split(key)[0].split(".")[-2]
- mapped_key = mapped_key.replace("*", layer_index)
- if "weight_g" in name:
- weight_type = "weight_g"
- elif "weight_v" in name:
- weight_type = "weight_v"
- elif "bias" in name:
- weight_type = "bias"
- elif "weight" in name:
- # TODO: don't match quantizer.weight_proj
- weight_type = "weight"
- else:
- weight_type = None
- set_recursively(hf_model, mapped_key, value, name, weight_type)
- continue
- if not is_used:
- unused_weights.append(name)
-
- logger.warning(f"Unused weights: {unused_weights}")
-
-
-def access_by_string(module, path):
- names = path.split(".")
- return reduce(getattr, names, module)
-
-
-def set_weights(full_name, module, fsq_value, hf_weight_path):
- hf_weight = access_by_string(module, hf_weight_path)
- hf_value = hf_weight.data
-
- if fsq_value.shape != hf_value.shape:
- raise ValueError(f"{full_name} has size {fsq_value.shape}, but {hf_value.shape} was found.")
- hf_weight.data = fsq_value
- logger.info(f"{full_name} was correctly initialized from {hf_weight_path}.")
-
-
-def load_conv_layer(full_name, value, feature_extractor, unused_weights):
- name = full_name.split("conv_layers.")[-1]
- items = name.split(".")
- layer_id = int(items[0])
- type_id = int(items[1])
-
- weight_type = name.split(".")[-1]
- if type_id == 0:
- layer_type = "conv"
- elif type_id == 2:
- layer_type = "layer_norm"
- else:
- unused_weights.append(full_name)
- return
-
- set_weights(full_name, feature_extractor, value, f"conv_layers.{layer_id}.{layer_type}.{weight_type}")
-
-
-def load_pos_conv_layer(full_name, value, pos_conv_embeddings, unused_weights):
- name = full_name.split("pos_conv.")[-1]
- items = name.split(".")
- layer_id = int(items[0])
- type_id = int(items[1])
-
- weight_type = name.split(".")[-1]
- if type_id != 0:
- unused_weights.append(full_name)
- return
- else:
- layer_type = "conv"
-
- set_weights(full_name, pos_conv_embeddings, value, f"layers.{layer_id}.{layer_type}.{weight_type}")
-
-
-@torch.no_grad()
-def convert_wav2vec2_checkpoint(
- checkpoint_path, pytorch_dump_folder_path, config_path=None, dict_path=None, is_finetuned=True
-):
- """
- Copy/paste/tweak model's weights to transformers design.
- """
- if config_path is not None:
- config = Data2VecAudioConfig.from_pretrained(config_path)
- else:
- config = Data2VecAudioConfig()
-
- if not is_finetuned:
- # Modify final_proj layer name
- hf_wav2vec = Data2VecAudioModel(config)
- data2vec_checkpoint_dir = os.path.dirname(checkpoint_path)
-
- state_dict = torch.load(checkpoint_path)
- state_dict["model"]["final_proj.weight"] = state_dict["model"].pop("final_proj.0.weight")
- state_dict["model"]["final_proj.bias"] = state_dict["model"].pop("final_proj.0.bias")
- converted_ckpt = os.path.join(data2vec_checkpoint_dir, "converted.pt")
- torch.save(state_dict, converted_ckpt)
- else:
- hf_wav2vec = Data2VecAudioForCTC(config)
- converted_ckpt = checkpoint_path
-
- def load_data2vec(path):
- model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([path])
- return model[0].eval()
-
- model = load_data2vec(converted_ckpt)
-
- recursively_load_weights(model, hf_wav2vec, not is_finetuned)
-
- processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-lv60")
-
- ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
- input_audio = [x["array"] for x in ds[:4]["audio"]]
-
- inputs = processor(input_audio, return_tensors="pt", padding=True)
-
- input_values = inputs.input_values
- attention_mask = inputs.attention_mask
- # input_values = inputs.input_values[:, :-1]
- # attention_mask = inputs.attention_mask[:, :-1]
-
- hf_wav2vec.eval()
- model.eval()
- if is_finetuned:
- their_output = model(source=input_values, padding_mask=(1 - attention_mask), mask=False, features_only=True)[
- "encoder_out"
- ].transpose(0, 1)
- our_output = hf_wav2vec(input_values, attention_mask=attention_mask)["logits"]
-
- pred_ids = torch.argmax(our_output, dim=-1)
- output_string = processor.batch_decode(pred_ids)
-
- print(f"Expected Output: {ds[:4]['text']}, Pred: {output_string}")
- else:
- their_output = model(source=input_values, padding_mask=(1 - attention_mask), mask=False, features_only=True)[
- "layer_results"
- ][-1][0].transpose(0, 1)
- our_output = hf_wav2vec(input_values, attention_mask=attention_mask)["last_hidden_state"]
-
- print(our_output.shape, their_output.shape)
- max_absolute_diff = torch.max(torch.abs(our_output - their_output)).item()
- print(f"max_absolute_diff = {max_absolute_diff}") # ~ 1e-7
- success = torch.allclose(our_output, their_output, atol=1e-3)
- print("Do both models output the same tensors?", "🔥" if success else "💩")
- if not success:
- raise Exception("Something went wRoNg")
-
- hf_wav2vec.save_pretrained(pytorch_dump_folder_path)
-
- if is_finetuned:
- processor.save_pretrained(pytorch_dump_folder_path)
- else:
- processor.feature_extractor.save_pretrained(pytorch_dump_folder_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
- parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
- parser.add_argument("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model")
- parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
- parser.add_argument(
- "--not_finetuned", action="store_true", help="Whether the model to convert is a fine-tuned model or not"
- )
- args = parser.parse_args()
- convert_wav2vec2_checkpoint(
- args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
- )
diff --git a/transformers/models/data2vec/convert_data2vec_text_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/data2vec/convert_data2vec_text_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index 81f5cd23fb9ef8ba045c1b363bfba3acbcffd876..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/convert_data2vec_text_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,208 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert data2vec checkpoint."""
-
-
-import argparse
-import os
-import pathlib
-
-import fairseq
-import torch
-from fairseq.modules import TransformerSentenceEncoderLayer
-from packaging import version
-
-from transformers import (
- Data2VecTextConfig,
- Data2VecTextForMaskedLM,
- Data2VecTextForSequenceClassification,
- Data2VecTextModel,
-)
-from transformers.models.bert.modeling_bert import (
- BertIntermediate,
- BertLayer,
- BertOutput,
- BertSelfAttention,
- BertSelfOutput,
-)
-
-# IMPORTANT: In order for this script to run, please make sure to download the dictionary: `dict.txt` from wget https://dl.fbaipublicfiles.com/fairseq/models/roberta.large.tar.gz
-# File copied from https://github.com/pytorch/fairseq/blob/main/examples/data2vec/models/data2vec_text.py
-from transformers.utils import logging
-
-
-if version.parse(fairseq.__version__) < version.parse("0.9.0"):
- raise Exception("requires fairseq >= 0.9.0")
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-SAMPLE_TEXT = "Hello world! cécé herlolip"
-
-
-def convert_data2vec_checkpoint_to_pytorch(
- data2vec_checkpoint_path: str, pytorch_dump_folder_path: str, classification_head: bool
-):
- """
- Copy/paste/tweak data2vec's weights to our BERT structure.
- """
- data2vec_checkpoint_dir, data2vec_checkpoint_file_name = os.path.split(data2vec_checkpoint_path)
- data2vec = Data2VecTextModel.from_pretrained(
- data2vec_checkpoint_dir, checkpoint_file=data2vec_checkpoint_file_name
- )
- data2vec.eval() # disable dropout
- data2vec_model = data2vec.models[0]
- data2vec_sent_encoder = data2vec_model.encoder.sentence_encoder
- config = Data2VecTextConfig(
- vocab_size=data2vec_sent_encoder.embed_tokens.num_embeddings,
- hidden_size=data2vec_model.args.encoder_embed_dim,
- num_hidden_layers=data2vec_model.args.encoder_layers,
- num_attention_heads=data2vec_model.args.encoder_attention_heads,
- intermediate_size=data2vec_model.args.encoder_ffn_embed_dim,
- max_position_embeddings=514,
- type_vocab_size=1,
- layer_norm_eps=1e-5, # PyTorch default used in fairseq
- )
- if classification_head:
- config.num_labels = data2vec.model.classification_heads["mnli"].out_proj.weight.shape[0]
- print("Our BERT config:", config)
-
- model = Data2VecTextForSequenceClassification(config) if classification_head else Data2VecTextForMaskedLM(config)
- model.eval()
-
- # Now let's copy all the weights.
- # Embeddings
- model.data2vec_text.embeddings.word_embeddings.weight = data2vec_sent_encoder.embed_tokens.weight
- model.data2vec_text.embeddings.position_embeddings.weight = data2vec_sent_encoder.embed_positions.weight
- model.data2vec_text.embeddings.token_type_embeddings.weight.data = torch.zeros_like(
- model.data2vec_text.embeddings.token_type_embeddings.weight
- ) # just zero them out b/c data2vec doesn't use them.
- model.data2vec_text.embeddings.LayerNorm.weight = data2vec_sent_encoder.layernorm_embedding.weight
- model.data2vec_text.embeddings.LayerNorm.bias = data2vec_sent_encoder.layernorm_embedding.bias
-
- for i in range(config.num_hidden_layers):
- # Encoder: start of layer
- layer: BertLayer = model.data2vec_text.encoder.layer[i]
- data2vec_layer: TransformerSentenceEncoderLayer = data2vec_sent_encoder.layers[i]
-
- # self attention
- self_attn: BertSelfAttention = layer.attention.self
- assert data2vec_layer.self_attn.k_proj.weight.data.shape == torch.Size(
- (config.hidden_size, config.hidden_size)
- ), (
- "Shape for data2vec_layer.self_attn.k_proj.weight.data should be"
- f" {torch.Size((config.hidden_size, config.hidden_size))}"
- )
- assert data2vec_layer.self_attn.q_proj.weight.data.shape == torch.Size(
- (config.hidden_size, config.hidden_size)
- ), (
- "Shape for data2vec_layer.self_attn.q_proj.weight.data should be"
- f" {torch.Size((config.hidden_size, config.hidden_size))}"
- )
- assert data2vec_layer.self_attn.v_proj.weight.data.shape == torch.Size(
- (config.hidden_size, config.hidden_size)
- ), (
- "Shape for data2vec_layer.self_attn.v_proj.weight.data should be"
- f" {torch.Size((config.hidden_size, config.hidden_size))}"
- )
-
- self_attn.query.weight.data = data2vec_layer.self_attn.q_proj.weight
- self_attn.query.bias.data = data2vec_layer.self_attn.q_proj.bias
- self_attn.key.weight.data = data2vec_layer.self_attn.k_proj.weight
- self_attn.key.bias.data = data2vec_layer.self_attn.k_proj.bias
- self_attn.value.weight.data = data2vec_layer.self_attn.v_proj.weight
- self_attn.value.bias.data = data2vec_layer.self_attn.v_proj.bias
-
- # self-attention output
- self_output: BertSelfOutput = layer.attention.output
- assert (
- self_output.dense.weight.shape == data2vec_layer.self_attn.out_proj.weight.shape
- ), f"Shape for self_output.dense.weight should be {data2vec_layer.self_attn.out_proj.weight.shape}"
- self_output.dense.weight = data2vec_layer.self_attn.out_proj.weight
- self_output.dense.bias = data2vec_layer.self_attn.out_proj.bias
- self_output.LayerNorm.weight = data2vec_layer.self_attn_layer_norm.weight
- self_output.LayerNorm.bias = data2vec_layer.self_attn_layer_norm.bias
-
- # intermediate
- intermediate: BertIntermediate = layer.intermediate
- assert (
- intermediate.dense.weight.shape == data2vec_layer.fc1.weight.shape
- ), f"Shape for intermediate.dense.weight should be {data2vec_layer.fc1.weight.shape}"
- intermediate.dense.weight = data2vec_layer.fc1.weight
- intermediate.dense.bias = data2vec_layer.fc1.bias
-
- # output
- bert_output: BertOutput = layer.output
- assert (
- bert_output.dense.weight.shape == data2vec_layer.fc2.weight.shape
- ), f"Shape for bert_output.dense.weight should be {data2vec_layer.fc2.weight.shape}"
- bert_output.dense.weight = data2vec_layer.fc2.weight
- bert_output.dense.bias = data2vec_layer.fc2.bias
- bert_output.LayerNorm.weight = data2vec_layer.final_layer_norm.weight
- bert_output.LayerNorm.bias = data2vec_layer.final_layer_norm.bias
- # end of layer
-
- if classification_head:
- model.classifier.dense.weight = data2vec.model.classification_heads["mnli"].dense.weight
- model.classifier.dense.bias = data2vec.model.classification_heads["mnli"].dense.bias
- model.classifier.out_proj.weight = data2vec.model.classification_heads["mnli"].out_proj.weight
- model.classifier.out_proj.bias = data2vec.model.classification_heads["mnli"].out_proj.bias
- else:
- # LM Head
- model.lm_head.dense.weight = data2vec_model.encoder.lm_head.dense.weight
- model.lm_head.dense.bias = data2vec_model.encoder.lm_head.dense.bias
- model.lm_head.layer_norm.weight = data2vec_model.encoder.lm_head.layer_norm.weight
- model.lm_head.layer_norm.bias = data2vec_model.encoder.lm_head.layer_norm.bias
- model.lm_head.decoder.weight = data2vec_model.encoder.lm_head.weight
- model.lm_head.decoder.bias = data2vec_model.encoder.lm_head.bias
-
- # Let's check that we get the same results.
- input_ids: torch.Tensor = data2vec.encode(SAMPLE_TEXT).unsqueeze(0) # batch of size 1
-
- our_output = model(input_ids)[0]
- if classification_head:
- their_output = data2vec.model.classification_heads["mnli"](data2vec.extract_features(input_ids))
- else:
- their_output = data2vec_model(input_ids)[0]
- print(our_output.shape, their_output.shape)
- max_absolute_diff = torch.max(torch.abs(our_output - their_output)).item()
- print(f"max_absolute_diff = {max_absolute_diff}") # ~ 1e-7
- success = torch.allclose(our_output, their_output, atol=1e-3)
- print("Do both models output the same tensors?", "🔥" if success else "💩")
- if not success:
- raise Exception("Something went wRoNg")
-
- pathlib.Path(pytorch_dump_folder_path).mkdir(parents=True, exist_ok=True)
- print(f"Saving model to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--checkpoint_path", default=None, type=str, required=True, help="Path the official PyTorch dump."
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
- parser.add_argument(
- "--classification_head", action="store_true", help="Whether to convert a final classification head."
- )
- args = parser.parse_args()
- convert_data2vec_checkpoint_to_pytorch(
- args.checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
- )
diff --git a/transformers/models/data2vec/convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/data2vec/convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index 0c6f42f4ba7f1b6a2afea7a9d03b9b89c1a21f25..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,374 +0,0 @@
-#!/usr/bin/env python3
-import argparse
-import json
-
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-from timm.models import create_model
-
-from transformers import (
- BeitImageProcessor,
- Data2VecVisionConfig,
- Data2VecVisionForImageClassification,
- Data2VecVisionModel,
-)
-
-
-def create_rename_keys(config, has_lm_head=False, is_semantic=False, hf_prefix="data2vec."):
- prefix = "backbone." if is_semantic else ""
-
- rename_keys = []
- for i in range(config.num_hidden_layers):
- # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
- rename_keys.append(
- (f"{prefix}blocks.{i}.norm1.weight", f"{hf_prefix}encoder.layer.{i}.layernorm_before.weight")
- )
- rename_keys.append((f"{prefix}blocks.{i}.norm1.bias", f"{hf_prefix}encoder.layer.{i}.layernorm_before.bias"))
- rename_keys.append(
- (f"{prefix}blocks.{i}.attn.proj.weight", f"{hf_prefix}encoder.layer.{i}.attention.output.dense.weight")
- )
- rename_keys.append(
- (f"{prefix}blocks.{i}.attn.proj.bias", f"{hf_prefix}encoder.layer.{i}.attention.output.dense.bias")
- )
- rename_keys.append(
- (f"{prefix}blocks.{i}.norm2.weight", f"{hf_prefix}encoder.layer.{i}.layernorm_after.weight")
- )
- rename_keys.append((f"{prefix}blocks.{i}.norm2.bias", f"{hf_prefix}encoder.layer.{i}.layernorm_after.bias"))
- rename_keys.append(
- (f"{prefix}blocks.{i}.mlp.fc1.weight", f"{hf_prefix}encoder.layer.{i}.intermediate.dense.weight")
- )
- rename_keys.append(
- (f"{prefix}blocks.{i}.mlp.fc1.bias", f"{hf_prefix}encoder.layer.{i}.intermediate.dense.bias")
- )
- rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.weight", f"{hf_prefix}encoder.layer.{i}.output.dense.weight"))
- rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.bias", f"{hf_prefix}encoder.layer.{i}.output.dense.bias"))
-
- # projection layer + position embeddings
- rename_keys.extend(
- [
- (f"{prefix}cls_token", f"{hf_prefix}embeddings.cls_token"),
- (f"{prefix}patch_embed.proj.weight", f"{hf_prefix}embeddings.patch_embeddings.projection.weight"),
- (f"{prefix}patch_embed.proj.bias", f"{hf_prefix}embeddings.patch_embeddings.projection.bias"),
- ]
- )
-
- if has_lm_head:
- # mask token + shared relative position bias + layernorm
- rename_keys.extend(
- [
- ("mask_token", f"{hf_prefix}embeddings.mask_token"),
- (
- "rel_pos_bias.relative_position_bias_table",
- f"{hf_prefix}encoder.relative_position_bias.relative_position_bias_table",
- ),
- (
- "rel_pos_bias.relative_position_index",
- f"{hf_prefix}encoder.relative_position_bias.relative_position_index",
- ),
- ("norm.weight", "layernorm.weight"),
- ("norm.bias", "layernorm.bias"),
- ]
- )
- elif is_semantic:
- # semantic segmentation classification heads
- rename_keys.extend(
- [
- ("decode_head.conv_seg.weight", "decode_head.classifier.weight"),
- ("decode_head.conv_seg.bias", "decode_head.classifier.bias"),
- ("auxiliary_head.conv_seg.weight", "auxiliary_head.classifier.weight"),
- ("auxiliary_head.conv_seg.bias", "auxiliary_head.classifier.bias"),
- ]
- )
- else:
- # layernorm + classification head
- rename_keys.extend(
- [
- ("fc_norm.weight", f"{hf_prefix}pooler.layernorm.weight"),
- ("fc_norm.bias", f"{hf_prefix}pooler.layernorm.bias"),
- ("head.weight", "classifier.weight"),
- ("head.bias", "classifier.bias"),
- ]
- )
-
- return rename_keys
-
-
-def read_in_q_k_v(state_dict, config, has_lm_head=False, is_semantic=False, hf_prefix="data2vec_vision."):
- for i in range(config.num_hidden_layers):
- prefix = "backbone." if is_semantic else ""
- # queries, keys and values
- in_proj_weight = state_dict.pop(f"{prefix}blocks.{i}.attn.qkv.weight")
- q_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.q_bias")
- v_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.v_bias")
-
- state_dict[f"{hf_prefix}encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[
- : config.hidden_size, :
- ]
- state_dict[f"{hf_prefix}encoder.layer.{i}.attention.attention.query.bias"] = q_bias
- state_dict[f"{hf_prefix}encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- config.hidden_size : config.hidden_size * 2, :
- ]
- state_dict[f"{hf_prefix}encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
- -config.hidden_size :, :
- ]
- state_dict[f"{hf_prefix}encoder.layer.{i}.attention.attention.value.bias"] = v_bias
-
- # gamma_1 and gamma_2
- # we call them lambda because otherwise they are renamed when using .from_pretrained
- gamma_1 = state_dict.pop(f"{prefix}blocks.{i}.gamma_1")
- gamma_2 = state_dict.pop(f"{prefix}blocks.{i}.gamma_2")
-
- state_dict[f"{hf_prefix}encoder.layer.{i}.lambda_1"] = gamma_1
- state_dict[f"{hf_prefix}encoder.layer.{i}.lambda_2"] = gamma_2
-
- # relative_position bias table + index
- if not has_lm_head:
- # each layer has its own relative position bias
- table = state_dict.pop(f"{prefix}blocks.{i}.attn.relative_position_bias_table")
- index = state_dict.pop(f"{prefix}blocks.{i}.attn.relative_position_index")
-
- state_dict[
- f"{hf_prefix}encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_bias_table"
- ] = table
- state_dict[
- f"{hf_prefix}encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_index"
- ] = index
-
-
-def get_args():
- parser = argparse.ArgumentParser(
- "Convert Data2VecVision to HF for image classification and pretraining", add_help=False
- )
- parser.add_argument("--hf_checkpoint_name", type=str)
- parser.add_argument("--input_size", default=224, type=int, help="images input size")
- parser.add_argument("--beit_checkpoint", default="", help="beit checkpoint")
-
- return parser.parse_args()
-
-
-def load_beit_model(args, is_finetuned, is_large):
- def load_state_dict(model, state_dict, prefix="", ignore_missing="relative_position_index"):
- missing_keys = []
- unexpected_keys = []
- error_msgs = []
- # copy state_dict so _load_from_state_dict can modify it
- metadata = getattr(state_dict, "_metadata", None)
- state_dict = state_dict.copy()
- if metadata is not None:
- state_dict._metadata = metadata
-
- def load(module, prefix=""):
- local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
- module._load_from_state_dict(
- state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs
- )
- for name, child in module._modules.items():
- if child is not None:
- load(child, prefix + name + ".")
-
- load(model, prefix=prefix)
-
- warn_missing_keys = []
- ignore_missing_keys = []
- for key in missing_keys:
- keep_flag = True
- for ignore_key in ignore_missing.split("|"):
- if ignore_key in key:
- keep_flag = False
- break
- if keep_flag:
- warn_missing_keys.append(key)
- else:
- ignore_missing_keys.append(key)
-
- missing_keys = warn_missing_keys
-
- if len(missing_keys) > 0:
- print(
- "Weights of {} not initialized from pretrained model: {}".format(
- model.__class__.__name__, missing_keys
- )
- )
- if len(unexpected_keys) > 0:
- print("Weights from pretrained model not used in {}: {}".format(model.__class__.__name__, unexpected_keys))
- if len(ignore_missing_keys) > 0:
- print(
- "Ignored weights of {} not initialized from pretrained model: {}".format(
- model.__class__.__name__, ignore_missing_keys
- )
- )
- if len(error_msgs) > 0:
- print("\n".join(error_msgs))
-
- model_kwargs = {
- "pretrained": False,
- "use_shared_rel_pos_bias": True,
- "use_abs_pos_emb": False,
- "init_values": 0.1,
- }
-
- if is_finetuned:
- model_kwargs.update(
- {
- "num_classes": 1000,
- "use_mean_pooling": True,
- "init_scale": 0.001,
- "use_rel_pos_bias": True,
- }
- )
-
- model = create_model(
- "beit_large_patch16_224" if is_large else "beit_base_patch16_224",
- **model_kwargs,
- )
- patch_size = model.patch_embed.patch_size
- args.window_size = (args.input_size // patch_size[0], args.input_size // patch_size[1])
- checkpoint = torch.load(args.beit_checkpoint, map_location="cpu")
-
- print(f"Load ckpt from {args.beit_checkpoint}")
- checkpoint_model = None
- for model_key in ("model", "module"):
- if model_key in checkpoint:
- checkpoint_model = checkpoint[model_key]
- print(f"Load state_dict by model_key = {model_key}")
- break
-
- all_keys = list(checkpoint_model.keys())
- for key in all_keys:
- if "relative_position_index" in key:
- checkpoint_model.pop(key)
-
- if "relative_position_bias_table" in key:
- rel_pos_bias = checkpoint_model[key]
- src_num_pos, num_attn_heads = rel_pos_bias.size()
- dst_num_pos, _ = model.state_dict()[key].size()
- dst_patch_shape = model.patch_embed.patch_shape
- if dst_patch_shape[0] != dst_patch_shape[1]:
- raise NotImplementedError()
-
- load_state_dict(model, checkpoint_model, prefix="")
-
- return model
-
-
-def main():
- args = get_args()
-
- is_finetuned = "ft1k" in args.hf_checkpoint_name
- is_large = "large" in args.hf_checkpoint_name
-
- if is_finetuned:
- # To convert Beit's data2vec_vision to HF you need to copy
- # https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_finetune.py
- # into this folder.
- import modeling_finetune # noqa: F401
- else:
- # To convert Beit's data2vec_vision to HF you need to copy
- # https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_cyclical.py
- # into this folder
- # IMPORTANT: Note that for now we've only converted the down-stream
- # model and not the full pretrained model. This means for the integration
- # test you need to add a `return x` after the following line:
- # https://github.com/facebookresearch/data2vec_vision/blob/af9a36349aaed59ae66e69b5dabeef2d62fdc5da/beit/modeling_cyclical.py#L197
- # to make the integration test pass.
- import modeling_cyclical # noqa: F401
-
- # 1. Create model config
- config = Data2VecVisionConfig()
- if is_finetuned:
- config.use_relative_position_bias = True
- config.use_shared_relative_position_bias = False
- config.use_mean_pooling = True
- config.num_labels = 1000
-
- repo_id = "huggingface/label-files"
- filename = "imagenet-1k-id2label.json"
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
- else:
- config.use_relative_position_bias = False
- config.use_shared_relative_position_bias = True
- config.use_mean_pooling = False
-
- if is_large:
- config.hidden_size = 1024
- config.intermediate_size = 4096
- config.num_hidden_layers = 24
- config.num_attention_heads = 16
-
- # 2. Load Beit model
- orig_model = load_beit_model(args, is_finetuned, is_large)
- orig_model.eval()
-
- # 3. Forward Beit model
- image_processor = BeitImageProcessor(size=config.image_size, do_center_crop=False)
- image = Image.open("../../../../tests/fixtures/tests_samples/COCO/000000039769.png")
- encoding = image_processor(images=image, return_tensors="pt")
- pixel_values = encoding["pixel_values"]
-
- orig_args = (pixel_values,) if is_finetuned else (pixel_values, None)
- with torch.no_grad():
- orig_model_output = orig_model(*orig_args)
-
- # 4. Load HF Data2VecVision model
- if is_finetuned:
- hf_model = Data2VecVisionForImageClassification(config)
- hf_model.eval()
- has_lm_head = False
- hf_prefix = "data2vec_vision."
- else:
- hf_model = Data2VecVisionModel(config)
- hf_model.eval()
- has_lm_head = True
- hf_prefix = ""
-
- rename_keys = create_rename_keys(config, hf_prefix=hf_prefix, has_lm_head=has_lm_head)
- state_dict = orig_model.state_dict()
- for src, dest in rename_keys:
- val = state_dict.pop(src)
- state_dict[dest] = val
-
- read_in_q_k_v(state_dict, config, hf_prefix=hf_prefix, has_lm_head=has_lm_head)
- missing_keys, unexpected_keys = hf_model.load_state_dict(state_dict, strict=False)
- print("HF missing", missing_keys)
- print("HF unexpected_keys", unexpected_keys)
-
- # 5. Forward HF Data2VecVision model
- with torch.no_grad():
- hf_model_output = hf_model(pixel_values)
-
- hf_output = hf_model_output.logits if is_finetuned else hf_model_output.last_hidden_state
-
- # 6. Compare
- max_absolute_diff = torch.max(torch.abs(hf_output - orig_model_output)).item()
-
- print(f"max_absolute_diff = {max_absolute_diff}")
- success = torch.allclose(hf_output, orig_model_output, atol=1e-3)
- print("Do both models output the same tensors?", "🔥" if success else "💩")
- if not success:
- raise Exception("Something went wRoNg")
-
- # 7. Save
- print(f"Saving to {args.hf_checkpoint_name}")
- hf_model.save_pretrained(args.hf_checkpoint_name)
- image_processor.save_pretrained(args.hf_checkpoint_name)
-
-
-if __name__ == "__main__":
- main()
- # Run the following to convert checkpoints
- # python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
- # --beit_checkpoint ./pretrained_base.pt \
- # --hf_checkpoint_name "./data2vec-vision-base"
- # python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
- # --beit_checkpoint ./finetuned_base.pt \
- # --hf_checkpoint_name "./data2vec-vision-base-ft1k"
- # python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
- # --beit_checkpoint ./pretrained_large.pt \
- # --hf_checkpoint_name "./data2vec-vision-large"
- # python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
- # --beit_checkpoint ./finetuned_large.pt \
- # --hf_checkpoint_name "./data2vec-vision-large-ft1k"
diff --git a/transformers/models/data2vec/modeling_data2vec_audio.py b/transformers/models/data2vec/modeling_data2vec_audio.py
deleted file mode 100644
index b5300cca084fa6d3c4f86b9154962c38464c1331..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/modeling_data2vec_audio.py
+++ /dev/null
@@ -1,1514 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The Fairseq Authors and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Data2VecAudio model."""
-
-import math
-import warnings
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from ...activations import ACT2FN
-from ...integrations.deepspeed import is_deepspeed_zero3_enabled
-from ...modeling_outputs import (
- BaseModelOutput,
- CausalLMOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
- Wav2Vec2BaseModelOutput,
- XVectorOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_peft_available,
- logging,
-)
-from .configuration_data2vec_audio import Data2VecAudioConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-_HIDDEN_STATES_START_POSITION = 2
-
-# General docstring
-_CONFIG_FOR_DOC = "Data2VecAudioConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/data2vec-audio-base-960h"
-_EXPECTED_OUTPUT_SHAPE = [1, 292, 768]
-
-# CTC docstring
-_CTC_EXPECTED_OUTPUT = "'MISTER QUILTER IS THE APOSTLE OF THE MIDDLE CLASSES AND WE ARE GLAD TO WELCOME HIS GOSPEL'"
-_CTC_EXPECTED_LOSS = 66.95
-
-
-from ..deprecated._archive_maps import DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2._compute_mask_indices
-def _compute_mask_indices(
- shape: Tuple[int, int],
- mask_prob: float,
- mask_length: int,
- attention_mask: Optional[torch.LongTensor] = None,
- min_masks: int = 0,
-) -> np.ndarray:
- """
- Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for
- ASR](https://arxiv.org/abs/1904.08779). Note that this method is not optimized to run on TPU and should be run on
- CPU as part of the preprocessing during training.
-
- Args:
- shape: The shape for which to compute masks. This should be of a tuple of size 2 where
- the first element is the batch size and the second element is the length of the axis to span.
- mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of
- independently generated mask spans of length `mask_length` is computed by
- `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the
- actual percentage will be smaller.
- mask_length: size of the mask
- min_masks: minimum number of masked spans
- attention_mask: A (right-padded) attention mask which independently shortens the feature axis of
- each batch dimension.
- """
- batch_size, sequence_length = shape
-
- if mask_length < 1:
- raise ValueError("`mask_length` has to be bigger than 0.")
-
- if mask_length > sequence_length:
- raise ValueError(
- f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}"
- f" and `sequence_length`: {sequence_length}`"
- )
-
- # epsilon is used for probabilistic rounding
- epsilon = np.random.rand(1).item()
-
- def compute_num_masked_span(input_length):
- """Given input length, compute how many spans should be masked"""
- num_masked_span = int(mask_prob * input_length / mask_length + epsilon)
- num_masked_span = max(num_masked_span, min_masks)
-
- # make sure num masked span <= sequence_length
- if num_masked_span * mask_length > sequence_length:
- num_masked_span = sequence_length // mask_length
-
- # make sure num_masked span is also <= input_length - (mask_length - 1)
- if input_length - (mask_length - 1) < num_masked_span:
- num_masked_span = max(input_length - (mask_length - 1), 0)
-
- return num_masked_span
-
- # compute number of masked spans in batch
- input_lengths = (
- attention_mask.sum(-1).detach().tolist()
- if attention_mask is not None
- else [sequence_length for _ in range(batch_size)]
- )
-
- # SpecAugment mask to fill
- spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)
- spec_aug_mask_idxs = []
-
- max_num_masked_span = compute_num_masked_span(sequence_length)
-
- if max_num_masked_span == 0:
- return spec_aug_mask
-
- for input_length in input_lengths:
- # compute num of masked spans for this input
- num_masked_span = compute_num_masked_span(input_length)
-
- # get random indices to mask
- spec_aug_mask_idx = np.random.choice(
- np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False
- )
-
- # pick first sampled index that will serve as a dummy index to pad vector
- # to ensure same dimension for all batches due to probabilistic rounding
- # Picking first sample just pads those vectors twice.
- if len(spec_aug_mask_idx) == 0:
- # this case can only happen if `input_length` is strictly smaller then
- # `sequence_length` in which case the last token has to be a padding
- # token which we can use as a dummy mask id
- dummy_mask_idx = sequence_length - 1
- else:
- dummy_mask_idx = spec_aug_mask_idx[0]
-
- spec_aug_mask_idx = np.concatenate(
- [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx]
- )
- spec_aug_mask_idxs.append(spec_aug_mask_idx)
-
- spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)
-
- # expand masked indices to masked spans
- spec_aug_mask_idxs = np.broadcast_to(
- spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length)
- )
- spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)
-
- # add offset to the starting indexes so that indexes now create a span
- offsets = np.arange(mask_length)[None, None, :]
- offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape(
- batch_size, max_num_masked_span * mask_length
- )
- spec_aug_mask_idxs = spec_aug_mask_idxs + offsets
-
- # ensure that we cannot have indices larger than sequence_length
- if spec_aug_mask_idxs.max() > sequence_length - 1:
- spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1
-
- # scatter indices to mask
- np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)
-
- return spec_aug_mask
-
-
-class Data2VecAudioConvLayer(nn.Module):
- def __init__(self, config, layer_id=0):
- super().__init__()
- self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1
- self.out_conv_dim = config.conv_dim[layer_id]
-
- self.conv = nn.Conv1d(
- self.in_conv_dim,
- self.out_conv_dim,
- kernel_size=config.conv_kernel[layer_id],
- stride=config.conv_stride[layer_id],
- bias=config.conv_bias,
- )
- self.layer_norm = nn.LayerNorm(self.out_conv_dim, elementwise_affine=True)
- self.activation = ACT2FN[config.feat_extract_activation]
-
- def forward(self, hidden_states):
- hidden_states = self.conv(hidden_states)
-
- hidden_states = hidden_states.transpose(-2, -1)
- hidden_states = self.layer_norm(hidden_states)
- hidden_states = hidden_states.transpose(-2, -1)
-
- hidden_states = self.activation(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2SamePadLayer with Wav2Vec2->Data2VecAudio
-class Data2VecAudioPadLayer(nn.Module):
- def __init__(self, num_conv_pos_embeddings):
- super().__init__()
- self.num_pad_remove = 1 if num_conv_pos_embeddings % 2 == 0 else 0
-
- def forward(self, hidden_states):
- if self.num_pad_remove > 0:
- hidden_states = hidden_states[:, :, : -self.num_pad_remove]
- return hidden_states
-
-
-class Data2VecAudioPositionalConvLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.conv = nn.Conv1d(
- config.hidden_size,
- config.hidden_size,
- kernel_size=config.conv_pos_kernel_size,
- padding=config.conv_pos_kernel_size // 2,
- groups=config.num_conv_pos_embedding_groups,
- )
-
- self.padding = Data2VecAudioPadLayer(config.conv_pos_kernel_size)
- self.activation = ACT2FN[config.feat_extract_activation]
- # no learnable parameters
- self.layer_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False)
-
- def forward(self, hidden_states):
- hidden_states = self.conv(hidden_states)
- hidden_states = self.padding(hidden_states)
-
- hidden_states = hidden_states.transpose(1, 2)
- hidden_states = self.layer_norm(hidden_states)
- hidden_states = hidden_states.transpose(1, 2)
- hidden_states = self.activation(hidden_states)
- return hidden_states
-
-
-class Data2VecAudioPositionalConvEmbedding(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.layers = nn.ModuleList(
- [Data2VecAudioPositionalConvLayer(config) for _ in range(config.num_conv_pos_embeddings)]
- )
-
- def forward(self, hidden_states):
- hidden_states = hidden_states.transpose(1, 2)
- for layer in self.layers:
- hidden_states = layer(hidden_states)
- hidden_states = hidden_states.transpose(1, 2)
- return hidden_states
-
-
-class Data2VecAudioFeatureEncoder(nn.Module):
- """Construct the features from raw audio waveform"""
-
- def __init__(self, config):
- super().__init__()
- self.conv_layers = nn.ModuleList(
- [Data2VecAudioConvLayer(config, layer_id=i) for i in range(config.num_feat_extract_layers)]
- )
- self.gradient_checkpointing = False
- self._requires_grad = True
-
- # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2FeatureEncoder._freeze_parameters
- def _freeze_parameters(self):
- for param in self.parameters():
- param.requires_grad = False
- self._requires_grad = False
-
- # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2FeatureEncoder.forward
- def forward(self, input_values):
- hidden_states = input_values[:, None]
-
- # make sure hidden_states require grad for gradient_checkpointing
- if self._requires_grad and self.training:
- hidden_states.requires_grad = True
-
- for conv_layer in self.conv_layers:
- if self._requires_grad and self.gradient_checkpointing and self.training:
- hidden_states = self._gradient_checkpointing_func(
- conv_layer.__call__,
- hidden_states,
- )
- else:
- hidden_states = conv_layer(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2FeatureProjection with Wav2Vec2->Data2VecAudio
-class Data2VecAudioFeatureProjection(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.layer_norm = nn.LayerNorm(config.conv_dim[-1], eps=config.layer_norm_eps)
- self.projection = nn.Linear(config.conv_dim[-1], config.hidden_size)
- self.dropout = nn.Dropout(config.feat_proj_dropout)
-
- def forward(self, hidden_states):
- # non-projected hidden states are needed for quantization
- norm_hidden_states = self.layer_norm(hidden_states)
- hidden_states = self.projection(norm_hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states, norm_hidden_states
-
-
-# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->Data2VecAudio
-class Data2VecAudioAttention(nn.Module):
- """Multi-headed attention from 'Attention Is All You Need' paper"""
-
- def __init__(
- self,
- embed_dim: int,
- num_heads: int,
- dropout: float = 0.0,
- is_decoder: bool = False,
- bias: bool = True,
- is_causal: bool = False,
- config: Optional[Data2VecAudioConfig] = None,
- ):
- super().__init__()
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.dropout = dropout
- self.head_dim = embed_dim // num_heads
- self.config = config
-
- if (self.head_dim * num_heads) != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
- f" and `num_heads`: {num_heads})."
- )
- self.scaling = self.head_dim**-0.5
- self.is_decoder = is_decoder
- self.is_causal = is_causal
-
- self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
-
- def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
- return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- key_value_states: Optional[torch.Tensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- layer_head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel"""
-
- # if key_value_states are provided this layer is used as a cross-attention layer
- # for the decoder
- is_cross_attention = key_value_states is not None
-
- bsz, tgt_len, _ = hidden_states.size()
-
- # get query proj
- query_states = self.q_proj(hidden_states) * self.scaling
- # get key, value proj
- # `past_key_value[0].shape[2] == key_value_states.shape[1]`
- # is checking that the `sequence_length` of the `past_key_value` is the same as
- # the provided `key_value_states` to support prefix tuning
- if (
- is_cross_attention
- and past_key_value is not None
- and past_key_value[0].shape[2] == key_value_states.shape[1]
- ):
- # reuse k,v, cross_attentions
- key_states = past_key_value[0]
- value_states = past_key_value[1]
- elif is_cross_attention:
- # cross_attentions
- key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
- value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
- elif past_key_value is not None:
- # reuse k, v, self_attention
- key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
- value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
- else:
- # self_attention
- key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
- value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-
- if self.is_decoder:
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_states, value_states)
-
- proj_shape = (bsz * self.num_heads, -1, self.head_dim)
- query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
- key_states = key_states.reshape(*proj_shape)
- value_states = value_states.reshape(*proj_shape)
-
- src_len = key_states.size(1)
- attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-
- if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
- raise ValueError(
- f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- if attention_mask is not None:
- if attention_mask.size() != (bsz, 1, tgt_len, src_len):
- raise ValueError(
- f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
- )
- attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if layer_head_mask is not None:
- if layer_head_mask.size() != (self.num_heads,):
- raise ValueError(
- f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
- f" {layer_head_mask.size()}"
- )
- attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- if output_attentions:
- # this operation is a bit awkward, but it's required to
- # make sure that attn_weights keeps its gradient.
- # In order to do so, attn_weights have to be reshaped
- # twice and have to be reused in the following
- attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
- attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-
- attn_output = torch.bmm(attn_probs, value_states)
-
- if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
- attn_output = attn_output.transpose(1, 2)
-
- # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
- # partitioned across GPUs when using tensor-parallelism.
- attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped, past_key_value
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2FeedForward with Wav2Vec2->Data2VecAudio
-class Data2VecAudioFeedForward(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.intermediate_dropout = nn.Dropout(config.activation_dropout)
-
- self.intermediate_dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- self.output_dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.output_dropout = nn.Dropout(config.hidden_dropout)
-
- def forward(self, hidden_states):
- hidden_states = self.intermediate_dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- hidden_states = self.intermediate_dropout(hidden_states)
-
- hidden_states = self.output_dense(hidden_states)
- hidden_states = self.output_dropout(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2EncoderLayer with Wav2Vec2->Data2VecAudio
-class Data2VecAudioEncoderLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.attention = Data2VecAudioAttention(
- embed_dim=config.hidden_size,
- num_heads=config.num_attention_heads,
- dropout=config.attention_dropout,
- is_decoder=False,
- )
- self.dropout = nn.Dropout(config.hidden_dropout)
- self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.feed_forward = Data2VecAudioFeedForward(config)
- self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_states, attention_mask=None, output_attentions=False):
- attn_residual = hidden_states
- hidden_states, attn_weights, _ = self.attention(
- hidden_states, attention_mask=attention_mask, output_attentions=output_attentions
- )
- hidden_states = self.dropout(hidden_states)
- hidden_states = attn_residual + hidden_states
-
- hidden_states = self.layer_norm(hidden_states)
- hidden_states = hidden_states + self.feed_forward(hidden_states)
- hidden_states = self.final_layer_norm(hidden_states)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Encoder with Wav2Vec2->Data2VecAudio
-class Data2VecAudioEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.pos_conv_embed = Data2VecAudioPositionalConvEmbedding(config)
- self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout)
- self.layers = nn.ModuleList([Data2VecAudioEncoderLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.tensor,
- attention_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- if attention_mask is not None:
- # make sure padded tokens output 0
- expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2])
- hidden_states[~expand_attention_mask] = 0
-
- # extend attention_mask
- attention_mask = 1.0 - attention_mask[:, None, None, :].to(dtype=hidden_states.dtype)
- attention_mask = attention_mask * torch.finfo(hidden_states.dtype).min
- attention_mask = attention_mask.expand(
- attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1]
- )
-
- position_embeddings = self.pos_conv_embed(hidden_states)
- hidden_states = hidden_states + position_embeddings
- hidden_states = self.layer_norm(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled()
-
- for layer in self.layers:
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
- dropout_probability = torch.rand([])
-
- skip_the_layer = True if self.training and (dropout_probability < self.config.layerdrop) else False
- if not skip_the_layer or deepspeed_zero3_is_enabled:
- # under deepspeed zero3 all gpus must run in sync
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer.__call__,
- hidden_states,
- attention_mask,
- output_attentions,
- )
- else:
- layer_outputs = layer(
- hidden_states, attention_mask=attention_mask, output_attentions=output_attentions
- )
- hidden_states = layer_outputs[0]
-
- if skip_the_layer:
- layer_outputs = (None, None)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Adapter with Wav2Vec2->Data2VecAudio
-class Data2VecAudioAdapter(nn.Module):
- def __init__(self, config):
- super().__init__()
-
- # feature dim might need to be down-projected
- if config.output_hidden_size != config.hidden_size:
- self.proj = nn.Linear(config.hidden_size, config.output_hidden_size)
- self.proj_layer_norm = nn.LayerNorm(config.output_hidden_size)
- else:
- self.proj = self.proj_layer_norm = None
-
- self.layers = nn.ModuleList(Data2VecAudioAdapterLayer(config) for _ in range(config.num_adapter_layers))
- self.layerdrop = config.layerdrop
-
- def forward(self, hidden_states):
- # down project hidden_states if necessary
- if self.proj is not None and self.proj_layer_norm is not None:
- hidden_states = self.proj(hidden_states)
- hidden_states = self.proj_layer_norm(hidden_states)
-
- hidden_states = hidden_states.transpose(1, 2)
-
- for layer in self.layers:
- layerdrop_prob = np.random.random()
- if not self.training or (layerdrop_prob > self.layerdrop):
- hidden_states = layer(hidden_states)
-
- hidden_states = hidden_states.transpose(1, 2)
- return hidden_states
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2AdapterLayer with Wav2Vec2->Data2VecAudio
-class Data2VecAudioAdapterLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.conv = nn.Conv1d(
- config.output_hidden_size,
- 2 * config.output_hidden_size,
- config.adapter_kernel_size,
- stride=config.adapter_stride,
- padding=1,
- )
-
- def forward(self, hidden_states):
- hidden_states = self.conv(hidden_states)
- hidden_states = nn.functional.glu(hidden_states, dim=1)
-
- return hidden_states
-
-
-class Data2VecAudioPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = Data2VecAudioConfig
- base_model_prefix = "data2vec_audio"
- main_input_name = "input_values"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, Data2VecAudioFeatureProjection):
- k = math.sqrt(1 / module.projection.in_features)
- nn.init.uniform_(module.projection.weight, a=-k, b=k)
- nn.init.uniform_(module.projection.bias, a=-k, b=k)
- elif isinstance(module, Data2VecAudioPositionalConvLayer):
- nn.init.constant_(module.conv.bias, 0)
- elif isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
- if module.bias is not None:
- module.bias.data.zero_()
- if module.weight is not None:
- module.weight.data.fill_(1.0)
- elif isinstance(module, nn.Conv1d):
- nn.init.kaiming_normal_(module.weight)
-
- if module.bias is not None:
- k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))
- nn.init.uniform_(module.bias, a=-k, b=k)
-
- # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2PreTrainedModel._get_feat_extract_output_lengths with
- def _get_feat_extract_output_lengths(
- self, input_lengths: Union[torch.LongTensor, int], add_adapter: Optional[bool] = None
- ):
- """
- Computes the output length of the convolutional layers
- """
-
- add_adapter = self.config.add_adapter if add_adapter is None else add_adapter
-
- def _conv_out_length(input_length, kernel_size, stride):
- # 1D convolutional layer output length formula taken
- # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
- return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 1
-
- for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride):
- input_lengths = _conv_out_length(input_lengths, kernel_size, stride)
-
- if add_adapter:
- for _ in range(self.config.num_adapter_layers):
- input_lengths = _conv_out_length(input_lengths, 1, self.config.adapter_stride)
-
- return input_lengths
-
- # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2PreTrainedModel._get_feature_vector_attention_mask
- def _get_feature_vector_attention_mask(
- self, feature_vector_length: int, attention_mask: torch.LongTensor, add_adapter=None
- ):
- # Effectively attention_mask.sum(-1), but not inplace to be able to run
- # on inference mode.
- non_padded_lengths = attention_mask.cumsum(dim=-1)[:, -1]
-
- output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths, add_adapter=add_adapter)
- output_lengths = output_lengths.to(torch.long)
-
- batch_size = attention_mask.shape[0]
-
- attention_mask = torch.zeros(
- (batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device
- )
- # these two operations makes sure that all values before the output lengths idxs are attended to
- attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 1
- attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()
- return attention_mask
-
-
-DATA2VEC_AUDIO_START_DOCSTRING = r"""
- Data2VecAudio was proposed in [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and
- Language](https://arxiv.org/pdf/2202.03555) by Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu and
- Michael Auli.
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving etc.).
-
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`Data2VecAudioConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-DATA2VEC_AUDIO_INPUTS_DOCSTRING = r"""
- Args:
- input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
- Float values of input raw speech waveform. Values can be obtained by loading a *.flac* or *.wav* audio file
- into an array of type *List[float]* or a *numpy.ndarray*, *e.g.* via the soundfile library (*pip install
- soundfile*). To prepare the array into *input_values*, the [`AutoProcessor`] should be used for padding and
- conversion into a tensor of type *torch.FloatTensor*. See [`Wav2Vec2Processor.__call__`] for details.
- attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing convolution and attention on padding token indices. Mask values selected in `[0,
- 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
-
-
-
- `attention_mask` should be passed if the corresponding processor has `config.return_attention_mask ==
- True`, which is the case for all pre-trained Data2Vec Audio models. Be aware that that even with
- `attention_mask`, zero-padded inputs will have slightly different outputs compared to non-padded inputs
- because there are more than one convolutional layer in the positional encodings. For a more detailed
- explanation, see [here](https://github.com/huggingface/transformers/issues/25621#issuecomment-1713759349).
-
-
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Data2VecAudio Model transformer outputting raw hidden-states without any specific head on top.",
- DATA2VEC_AUDIO_START_DOCSTRING,
-)
-class Data2VecAudioModel(Data2VecAudioPreTrainedModel):
- def __init__(self, config: Data2VecAudioConfig):
- super().__init__(config)
- self.config = config
- self.feature_extractor = Data2VecAudioFeatureEncoder(config)
- self.feature_projection = Data2VecAudioFeatureProjection(config)
-
- # model only needs masking vector if mask prob is > 0.0
- if config.mask_time_prob > 0.0 or config.mask_feature_prob > 0.0:
- self.masked_spec_embed = nn.Parameter(torch.FloatTensor(config.hidden_size).uniform_())
-
- self.encoder = Data2VecAudioEncoder(config)
-
- self.adapter = Data2VecAudioAdapter(config) if config.add_adapter else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def freeze_feature_encoder(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameter will
- not be updated during training.
- """
- self.feature_extractor._freeze_parameters()
-
- def _mask_hidden_states(
- self,
- hidden_states: torch.FloatTensor,
- mask_time_indices: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.LongTensor] = None,
- ):
- """
- Masks extracted features along time axis and/or along feature axis according to
- [SpecAugment](https://arxiv.org/abs/1904.08779).
- """
-
- # `config.apply_spec_augment` can set masking to False
- if not getattr(self.config, "apply_spec_augment", True):
- return hidden_states
-
- # generate indices & apply SpecAugment along time axis
- batch_size, sequence_length, hidden_size = hidden_states.size()
-
- if mask_time_indices is not None:
- # apply SpecAugment along time axis with given mask_time_indices
- hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype)
- elif self.config.mask_time_prob > 0 and self.training:
- mask_time_indices = _compute_mask_indices(
- (batch_size, sequence_length),
- mask_prob=self.config.mask_time_prob,
- mask_length=self.config.mask_time_length,
- attention_mask=attention_mask,
- min_masks=self.config.mask_time_min_masks,
- )
- mask_time_indices = torch.tensor(mask_time_indices, device=hidden_states.device, dtype=torch.bool)
- hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype)
-
- if self.config.mask_feature_prob > 0 and self.training:
- # generate indices & apply SpecAugment along feature axis
- mask_feature_indices = _compute_mask_indices(
- (batch_size, hidden_size),
- mask_prob=self.config.mask_feature_prob,
- mask_length=self.config.mask_feature_length,
- min_masks=self.config.mask_feature_min_masks,
- )
- mask_feature_indices = torch.tensor(mask_feature_indices, device=hidden_states.device, dtype=torch.bool)
- mask_feature_indices = mask_feature_indices[:, None].expand(-1, sequence_length, -1)
- hidden_states[mask_feature_indices] = 0
-
- return hidden_states
-
- @add_start_docstrings_to_model_forward(DATA2VEC_AUDIO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=Wav2Vec2BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="audio",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- input_values: Optional[torch.Tensor],
- attention_mask: Optional[torch.Tensor] = None,
- mask_time_indices: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, Wav2Vec2BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- extract_features = self.feature_extractor(input_values)
- extract_features = extract_features.transpose(1, 2)
-
- if attention_mask is not None:
- # compute reduced attention_mask corresponding to feature vectors
- attention_mask = self._get_feature_vector_attention_mask(
- extract_features.shape[1], attention_mask, add_adapter=False
- )
-
- hidden_states, extract_features = self.feature_projection(extract_features)
- hidden_states = self._mask_hidden_states(
- hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
- )
-
- encoder_outputs = self.encoder(
- hidden_states,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = encoder_outputs[0]
-
- if self.adapter is not None:
- hidden_states = self.adapter(hidden_states)
-
- if not return_dict:
- return (hidden_states, extract_features) + encoder_outputs[1:]
-
- return Wav2Vec2BaseModelOutput(
- last_hidden_state=hidden_states,
- extract_features=extract_features,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """Data2VecAudio Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).""",
- DATA2VEC_AUDIO_START_DOCSTRING,
-)
-class Data2VecAudioForCTC(Data2VecAudioPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.data2vec_audio = Data2VecAudioModel(config)
- self.dropout = nn.Dropout(config.final_dropout)
-
- if config.vocab_size is None:
- raise ValueError(
- f"You are trying to instantiate {self.__class__} with a configuration that "
- "does not define the vocabulary size of the language model head. Please "
- "instantiate the model as follows: `Data2VecAudioForCTC.from_pretrained(..., vocab_size=vocab_size)`. "
- "or define `vocab_size` of your model's configuration."
- )
- output_hidden_size = (
- config.output_hidden_size if hasattr(config, "add_adapter") and config.add_adapter else config.hidden_size
- )
- self.lm_head = nn.Linear(output_hidden_size, config.vocab_size)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def freeze_feature_extractor(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameter will
- not be updated during training.
- """
- warnings.warn(
- "The method `freeze_feature_extractor` is deprecated and will be removed in Transformers v5. "
- "Please use the equivalent `freeze_feature_encoder` method instead.",
- FutureWarning,
- )
- self.freeze_feature_encoder()
-
- def freeze_feature_encoder(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameter will
- not be updated during training.
- """
- self.data2vec_audio.feature_extractor._freeze_parameters()
-
- @add_start_docstrings_to_model_forward(DATA2VEC_AUDIO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_CTC_EXPECTED_OUTPUT,
- expected_loss=_CTC_EXPECTED_LOSS,
- )
- # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForCTC.forward with wav2vec2->data2vec_audio
- def forward(
- self,
- input_values: Optional[torch.Tensor],
- attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: Optional[torch.Tensor] = None,
- ) -> Union[Tuple, CausalLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, target_length)`, *optional*):
- Labels for connectionist temporal classification. Note that `target_length` has to be smaller or equal to
- the sequence length of the output logits. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`.
- All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ...,
- config.vocab_size - 1]`.
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.data2vec_audio(
- input_values,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- hidden_states = self.dropout(hidden_states)
-
- logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- if labels.max() >= self.config.vocab_size:
- raise ValueError(f"Label values must be <= vocab_size: {self.config.vocab_size}")
-
- # retrieve loss input_lengths from attention_mask
- attention_mask = (
- attention_mask if attention_mask is not None else torch.ones_like(input_values, dtype=torch.long)
- )
- input_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long)
-
- # assuming that padded tokens are filled with -100
- # when not being attended to
- labels_mask = labels >= 0
- target_lengths = labels_mask.sum(-1)
- flattened_targets = labels.masked_select(labels_mask)
-
- # ctc_loss doesn't support fp16
- log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)
-
- with torch.backends.cudnn.flags(enabled=False):
- loss = nn.functional.ctc_loss(
- log_probs,
- flattened_targets,
- input_lengths,
- target_lengths,
- blank=self.config.pad_token_id,
- reduction=self.config.ctc_loss_reduction,
- zero_infinity=self.config.ctc_zero_infinity,
- )
-
- if not return_dict:
- output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutput(
- loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
-
-@add_start_docstrings(
- """
- Data2VecAudio Model with a sequence classification head on top (a linear layer over the pooled output) for tasks
- like SUPERB Keyword Spotting.
- """,
- DATA2VEC_AUDIO_START_DOCSTRING,
-)
-class Data2VecAudioForSequenceClassification(Data2VecAudioPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- if hasattr(config, "add_adapter") and config.add_adapter:
- raise ValueError(
- "Sequence classification does not support the use of Data2VecAudio adapters (config.add_adapter=True)"
- )
- self.data2vec_audio = Data2VecAudioModel(config)
- num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings
- if config.use_weighted_layer_sum:
- self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)
- self.projector = nn.Linear(config.hidden_size, config.classifier_proj_size)
- self.classifier = nn.Linear(config.classifier_proj_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def freeze_feature_extractor(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameters will
- not be updated during training.
- """
- warnings.warn(
- "The method `freeze_feature_extractor` is deprecated and will be removed in Transformers v5. "
- "Please use the equivalent `freeze_feature_encoder` method instead.",
- FutureWarning,
- )
- self.freeze_feature_encoder()
-
- def freeze_feature_encoder(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameter will
- not be updated during training.
- """
- self.data2vec_audio.feature_extractor._freeze_parameters()
-
- def freeze_base_model(self):
- """
- Calling this function will disable the gradient computation for the base model so that its parameters will not
- be updated during training. Only the classification head will be updated.
- """
- for param in self.data2vec_audio.parameters():
- param.requires_grad = False
-
- @add_start_docstrings_to_model_forward(DATA2VEC_AUDIO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="audio",
- )
- # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForSequenceClassification.forward with wav2vec2->data2vec_audio
- def forward(
- self,
- input_values: Optional[torch.Tensor],
- attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: Optional[torch.Tensor] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states
-
- outputs = self.data2vec_audio(
- input_values,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if self.config.use_weighted_layer_sum:
- hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
- hidden_states = torch.stack(hidden_states, dim=1)
- norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)
- hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
- else:
- hidden_states = outputs[0]
-
- hidden_states = self.projector(hidden_states)
- if attention_mask is None:
- pooled_output = hidden_states.mean(dim=1)
- else:
- padding_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask)
- hidden_states[~padding_mask] = 0.0
- pooled_output = hidden_states.sum(dim=1) / padding_mask.sum(dim=1).view(-1, 1)
-
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Data2VecAudio Model with a frame classification head on top for tasks like Speaker Diarization.
- """,
- DATA2VEC_AUDIO_START_DOCSTRING,
-)
-class Data2VecAudioForAudioFrameClassification(Data2VecAudioPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- if hasattr(config, "add_adapter") and config.add_adapter:
- raise ValueError(
- "Audio frame classification does not support the use of Data2VecAudio adapters"
- " (config.add_adapter=True)"
- )
- self.data2vec_audio = Data2VecAudioModel(config)
- num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings
- if config.use_weighted_layer_sum:
- self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
- self.num_labels = config.num_labels
-
- self.init_weights()
-
- def freeze_feature_extractor(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameter will
- not be updated during training.
- """
- warnings.warn(
- "The method `freeze_feature_extractor` is deprecated and will be removed in Transformers v5. "
- "Please use the equivalent `freeze_feature_encoder` method instead.",
- FutureWarning,
- )
- self.freeze_feature_encoder()
-
- def freeze_feature_encoder(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameter will
- not be updated during training.
- """
- self.data2vec_audio.feature_extractor._freeze_parameters()
-
- def freeze_base_model(self):
- """
- Calling this function will disable the gradient computation for the base model so that its parameters will not
- be updated during training. Only the classification head will be updated.
- """
- for param in self.data2vec_audio.parameters():
- param.requires_grad = False
-
- @add_start_docstrings_to_model_forward(DATA2VEC_AUDIO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="audio",
- )
- # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForAudioFrameClassification.forward with wav2vec2->data2vec_audio
- def forward(
- self,
- input_values: Optional[torch.Tensor],
- attention_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states
-
- outputs = self.data2vec_audio(
- input_values,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if self.config.use_weighted_layer_sum:
- hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
- hidden_states = torch.stack(hidden_states, dim=1)
- norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)
- hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
- else:
- hidden_states = outputs[0]
-
- logits = self.classifier(hidden_states)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), torch.argmax(labels.view(-1, self.num_labels), axis=1))
-
- if not return_dict:
- output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
- return output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.AMSoftmaxLoss
-class AMSoftmaxLoss(nn.Module):
- def __init__(self, input_dim, num_labels, scale=30.0, margin=0.4):
- super(AMSoftmaxLoss, self).__init__()
- self.scale = scale
- self.margin = margin
- self.num_labels = num_labels
- self.weight = nn.Parameter(torch.randn(input_dim, num_labels), requires_grad=True)
- self.loss = nn.CrossEntropyLoss()
-
- def forward(self, hidden_states, labels):
- labels = labels.flatten()
- weight = nn.functional.normalize(self.weight, dim=0)
- hidden_states = nn.functional.normalize(hidden_states, dim=1)
- cos_theta = torch.mm(hidden_states, weight)
- psi = cos_theta - self.margin
-
- onehot = nn.functional.one_hot(labels, self.num_labels)
- logits = self.scale * torch.where(onehot.bool(), psi, cos_theta)
- loss = self.loss(logits, labels)
-
- return loss
-
-
-# Copied from transformers.models.wav2vec2.modeling_wav2vec2.TDNNLayer
-class TDNNLayer(nn.Module):
- def __init__(self, config, layer_id=0):
- super().__init__()
- self.in_conv_dim = config.tdnn_dim[layer_id - 1] if layer_id > 0 else config.tdnn_dim[layer_id]
- self.out_conv_dim = config.tdnn_dim[layer_id]
- self.kernel_size = config.tdnn_kernel[layer_id]
- self.dilation = config.tdnn_dilation[layer_id]
-
- self.kernel = nn.Linear(self.in_conv_dim * self.kernel_size, self.out_conv_dim)
- self.activation = nn.ReLU()
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- if is_peft_available():
- from peft.tuners.lora import LoraLayer
-
- if isinstance(self.kernel, LoraLayer):
- warnings.warn(
- "Detected LoRA on TDNNLayer. LoRA weights won't be applied due to optimization. "
- "You should exclude TDNNLayer from LoRA's target modules.",
- )
-
- # for backward compatibility, we keep nn.Linear but call F.conv1d for speed up
- hidden_states = hidden_states.transpose(1, 2)
- weight = self.kernel.weight.view(self.out_conv_dim, self.kernel_size, self.in_conv_dim).transpose(1, 2)
- hidden_states = nn.functional.conv1d(hidden_states, weight, self.kernel.bias, dilation=self.dilation)
- hidden_states = hidden_states.transpose(1, 2)
-
- hidden_states = self.activation(hidden_states)
- return hidden_states
-
-
-@add_start_docstrings(
- """
- Data2VecAudio Model with an XVector feature extraction head on top for tasks like Speaker Verification.
- """,
- DATA2VEC_AUDIO_START_DOCSTRING,
-)
-class Data2VecAudioForXVector(Data2VecAudioPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.data2vec_audio = Data2VecAudioModel(config)
- num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings
- if config.use_weighted_layer_sum:
- self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)
- self.projector = nn.Linear(config.hidden_size, config.tdnn_dim[0])
-
- tdnn_layers = [TDNNLayer(config, i) for i in range(len(config.tdnn_dim))]
- self.tdnn = nn.ModuleList(tdnn_layers)
-
- self.feature_extractor = nn.Linear(config.tdnn_dim[-1] * 2, config.xvector_output_dim)
- self.classifier = nn.Linear(config.xvector_output_dim, config.xvector_output_dim)
-
- self.objective = AMSoftmaxLoss(config.xvector_output_dim, config.num_labels)
-
- self.init_weights()
-
- def freeze_feature_extractor(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameter will
- not be updated during training.
- """
- warnings.warn(
- "The method `freeze_feature_extractor` is deprecated and will be removed in Transformers v5. "
- "Please use the equivalent `freeze_feature_encoder` method instead.",
- FutureWarning,
- )
- self.freeze_feature_encoder()
-
- def freeze_feature_encoder(self):
- """
- Calling this function will disable the gradient computation for the feature encoder so that its parameter will
- not be updated during training.
- """
- self.data2vec_audio.feature_extractor._freeze_parameters()
-
- def freeze_base_model(self):
- """
- Calling this function will disable the gradient computation for the base model so that its parameters will not
- be updated during training. Only the classification head will be updated.
- """
- for param in self.data2vec_audio.parameters():
- param.requires_grad = False
-
- def _get_tdnn_output_lengths(self, input_lengths: Union[torch.LongTensor, int]):
- """
- Computes the output length of the TDNN layers
- """
-
- def _conv_out_length(input_length, kernel_size, stride):
- # 1D convolutional layer output length formula taken
- # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
- return (input_length - kernel_size) // stride + 1
-
- for kernel_size in self.config.tdnn_kernel:
- input_lengths = _conv_out_length(input_lengths, kernel_size, 1)
-
- return input_lengths
-
- @add_start_docstrings_to_model_forward(DATA2VEC_AUDIO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=XVectorOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="audio",
- )
- # Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForXVector.forward with wav2vec2->data2vec_audio
- def forward(
- self,
- input_values: Optional[torch.Tensor],
- attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: Optional[torch.Tensor] = None,
- ) -> Union[Tuple, XVectorOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states
-
- outputs = self.data2vec_audio(
- input_values,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if self.config.use_weighted_layer_sum:
- hidden_states = outputs[_HIDDEN_STATES_START_POSITION]
- hidden_states = torch.stack(hidden_states, dim=1)
- norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)
- hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
- else:
- hidden_states = outputs[0]
-
- hidden_states = self.projector(hidden_states)
-
- for tdnn_layer in self.tdnn:
- hidden_states = tdnn_layer(hidden_states)
-
- # Statistic Pooling
- if attention_mask is None:
- mean_features = hidden_states.mean(dim=1)
- std_features = hidden_states.std(dim=1)
- else:
- feat_extract_output_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(dim=1))
- tdnn_output_lengths = self._get_tdnn_output_lengths(feat_extract_output_lengths)
- mean_features = []
- std_features = []
- for i, length in enumerate(tdnn_output_lengths):
- mean_features.append(hidden_states[i, :length].mean(dim=0))
- std_features.append(hidden_states[i, :length].std(dim=0))
- mean_features = torch.stack(mean_features)
- std_features = torch.stack(std_features)
- statistic_pooling = torch.cat([mean_features, std_features], dim=-1)
-
- output_embeddings = self.feature_extractor(statistic_pooling)
- logits = self.classifier(output_embeddings)
-
- loss = None
- if labels is not None:
- loss = self.objective(logits, labels)
-
- if not return_dict:
- output = (logits, output_embeddings) + outputs[_HIDDEN_STATES_START_POSITION:]
- return ((loss,) + output) if loss is not None else output
-
- return XVectorOutput(
- loss=loss,
- logits=logits,
- embeddings=output_embeddings,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/data2vec/modeling_data2vec_text.py b/transformers/models/data2vec/modeling_data2vec_text.py
deleted file mode 100644
index 7dcc53e2cc15c81bc83088fd03e7a7f4a29bce2b..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/modeling_data2vec_text.py
+++ /dev/null
@@ -1,1557 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch Data2VecText model."""
-
-import math
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN, gelu
-from ...modeling_outputs import (
- BaseModelOutputWithPastAndCrossAttentions,
- BaseModelOutputWithPoolingAndCrossAttentions,
- CausalLMOutputWithCrossAttentions,
- MaskedLMOutput,
- MultipleChoiceModelOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_data2vec_text import Data2VecTextConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-_HIDDEN_STATES_START_POSITION = 2
-
-# General docstring
-_CHECKPOINT_FOR_DOC = "facebook/data2vec-text-base"
-_CONFIG_FOR_DOC = "Data2VecTextConfig"
-
-
-from ..deprecated._archive_maps import DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.roberta.modeling_roberta.RobertaEmbeddings with Roberta->Data2VecText
-class Data2VecTextForTextEmbeddings(nn.Module):
- """
- Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
- """
-
- # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
-
- # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
- # any TensorFlow checkpoint file
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
- self.register_buffer(
- "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
- )
-
- # End copy
- self.padding_idx = config.pad_token_id
- self.position_embeddings = nn.Embedding(
- config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
- )
-
- def forward(
- self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
- ):
- if position_ids is None:
- if input_ids is not None:
- # Create the position ids from the input token ids. Any padded tokens remain padded.
- position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)
- else:
- position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)
-
- if input_ids is not None:
- input_shape = input_ids.size()
- else:
- input_shape = inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
- # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
- # issue #5664
- if token_type_ids is None:
- if hasattr(self, "token_type_ids"):
- buffered_token_type_ids = self.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
- embeddings = inputs_embeds + token_type_embeddings
- if self.position_embedding_type == "absolute":
- position_embeddings = self.position_embeddings(position_ids)
- embeddings += position_embeddings
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
- def create_position_ids_from_inputs_embeds(self, inputs_embeds):
- """
- We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
-
- Args:
- inputs_embeds: torch.Tensor
-
- Returns: torch.Tensor
- """
- input_shape = inputs_embeds.size()[:-1]
- sequence_length = input_shape[1]
-
- position_ids = torch.arange(
- self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
- )
- return position_ids.unsqueeze(0).expand(input_shape)
-
-
-# Copied from transformers.models.roberta.modeling_roberta.RobertaSelfAttention with Roberta->Data2VecText
-class Data2VecTextSelfAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size)
- self.key = nn.Linear(config.hidden_size, self.all_head_size)
- self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
- self.position_embedding_type = position_embedding_type or getattr(
- config, "position_embedding_type", "absolute"
- )
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- self.max_position_embeddings = config.max_position_embeddings
- self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
-
- self.is_decoder = config.is_decoder
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- mixed_query_layer = self.query(hidden_states)
-
- # If this is instantiated as a cross-attention module, the keys
- # and values come from an encoder; the attention mask needs to be
- # such that the encoder's padding tokens are not attended to.
- is_cross_attention = encoder_hidden_states is not None
-
- if is_cross_attention and past_key_value is not None:
- # reuse k,v, cross_attentions
- key_layer = past_key_value[0]
- value_layer = past_key_value[1]
- attention_mask = encoder_attention_mask
- elif is_cross_attention:
- key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
- value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
- attention_mask = encoder_attention_mask
- elif past_key_value is not None:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
- value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
- else:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- use_cache = past_key_value is not None
- if self.is_decoder:
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_layer, value_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- query_length, key_length = query_layer.shape[2], key_layer.shape[2]
- if use_cache:
- position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(
- -1, 1
- )
- else:
- position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
- position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
- distance = position_ids_l - position_ids_r
-
- positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
- positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
-
- if self.position_embedding_type == "relative_key":
- relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores
- elif self.position_embedding_type == "relative_key_query":
- relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in Data2VecTextModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- if self.is_decoder:
- outputs = outputs + (past_key_value,)
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
-class Data2VecTextSelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Data2VecText
-class Data2VecTextAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- self.self = Data2VecTextSelfAttention(config, position_embedding_type=position_embedding_type)
- self.output = Data2VecTextSelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- self_outputs = self.self(
- hidden_states,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertIntermediate
-class Data2VecTextIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOutput
-class Data2VecTextOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Data2VecText
-class Data2VecTextLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = Data2VecTextAttention(config)
- self.is_decoder = config.is_decoder
- self.add_cross_attention = config.add_cross_attention
- if self.add_cross_attention:
- if not self.is_decoder:
- raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
- self.crossattention = Data2VecTextAttention(config, position_embedding_type="absolute")
- self.intermediate = Data2VecTextIntermediate(config)
- self.output = Data2VecTextOutput(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
- self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
- self_attention_outputs = self.attention(
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions=output_attentions,
- past_key_value=self_attn_past_key_value,
- )
- attention_output = self_attention_outputs[0]
-
- # if decoder, the last output is tuple of self-attn cache
- if self.is_decoder:
- outputs = self_attention_outputs[1:-1]
- present_key_value = self_attention_outputs[-1]
- else:
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- cross_attn_present_key_value = None
- if self.is_decoder and encoder_hidden_states is not None:
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
- " by setting `config.add_cross_attention=True`"
- )
-
- # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
- cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
- cross_attention_outputs = self.crossattention(
- attention_output,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- cross_attn_past_key_value,
- output_attentions,
- )
- attention_output = cross_attention_outputs[0]
- outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
-
- # add cross-attn cache to positions 3,4 of present_key_value tuple
- cross_attn_present_key_value = cross_attention_outputs[-1]
- present_key_value = present_key_value + cross_attn_present_key_value
-
- layer_output = apply_chunking_to_forward(
- self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
- )
- outputs = (layer_output,) + outputs
-
- # if decoder, return the attn key/values as the last output
- if self.is_decoder:
- outputs = outputs + (present_key_value,)
-
- return outputs
-
- def feed_forward_chunk(self, attention_output):
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.output(intermediate_output, attention_output)
- return layer_output
-
-
-# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Data2VecText
-class Data2VecTextEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([Data2VecTextLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = False,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- next_decoder_cache = () if use_cache else None
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
- past_key_value = past_key_values[i] if past_key_values is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
-
- hidden_states = layer_outputs[0]
- if use_cache:
- next_decoder_cache += (layer_outputs[-1],)
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- next_decoder_cache,
- all_hidden_states,
- all_self_attentions,
- all_cross_attentions,
- ]
- if v is not None
- )
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=next_decoder_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPooler
-class Data2VecTextPooler(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = nn.Tanh()
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(first_token_tensor)
- pooled_output = self.activation(pooled_output)
- return pooled_output
-
-
-class Data2VecTextPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = Data2VecTextConfig
- base_model_prefix = "data2vec_text"
- supports_gradient_checkpointing = True
- _no_split_modules = ["Data2VecTextForTextEmbeddings", "Data2VecTextLayer"]
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- if hasattr(module, "bias") and module.bias is not None:
- module.bias.data.zero_()
- if hasattr(module, "weight") and module.weight is not None:
- module.weight.data.fill_(1.0)
-
-
-DATA2VECTEXT_START_DOCSTRING = r"""
- Data2VecText was proposed in [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and
- Language](https://arxiv.org/pdf/2202.03555) by Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu and
- Michael Auli.
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`Data2VecTextConfig`]): Model configuration class with all the parameters of the
- model. Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DATA2VECTEXT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Data2VecText Model for text transformer outputting raw hidden-states without any specific head on top.",
- DATA2VECTEXT_START_DOCSTRING,
-)
-class Data2VecTextModel(Data2VecTextPreTrainedModel):
- """
-
- The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
- cross-attention is added between the self-attention layers, following the architecture described in *Attention is
- all you need*_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
- Kaiser and Illia Polosukhin.
-
- To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
- to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
- `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
-
- .. _*Attention is all you need*: https://arxiv.org/abs/1706.03762
-
- """
-
- def __init__(self, config, add_pooling_layer=True):
- super().__init__(config)
- self.config = config
-
- self.embeddings = Data2VecTextForTextEmbeddings(config)
- self.encoder = Data2VecTextEncoder(config)
-
- self.pooler = Data2VecTextPooler(config) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(DATA2VECTEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPoolingAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- # Copied from transformers.models.bert.modeling_bert.BertModel.forward
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
- r"""
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if self.config.is_decoder:
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- else:
- use_cache = False
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- # past_key_values_length
- past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
-
- if attention_mask is None:
- attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
-
- if token_type_ids is None:
- if hasattr(self.embeddings, "token_type_ids"):
- buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
- # ourselves in which case we just need to make it broadcastable to all heads.
- extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
-
- # If a 2D or 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- if self.config.is_decoder and encoder_hidden_states is not None:
- encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
- encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
- if encoder_attention_mask is None:
- encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
- encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
- else:
- encoder_extended_attention_mask = None
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- past_key_values_length=past_key_values_length,
- )
- encoder_outputs = self.encoder(
- embedding_output,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_extended_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- return (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPoolingAndCrossAttentions(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- past_key_values=encoder_outputs.past_key_values,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- cross_attentions=encoder_outputs.cross_attentions,
- )
-
-
-@add_start_docstrings(
- """Data2VecText Model with a `language modeling` head on top for CLM fine-tuning.""", DATA2VECTEXT_START_DOCSTRING
-)
-class Data2VecTextForCausalLM(Data2VecTextPreTrainedModel):
- _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"]
-
- def __init__(self, config):
- super().__init__(config)
-
- if not config.is_decoder:
- logger.warning("If you want to use `Data2VecTextLMHeadModel` as a standalone, add `is_decoder=True.`")
-
- self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)
- self.lm_head = Data2VecTextLMHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.lm_head.decoder
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(DATA2VECTEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
- r"""
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
- `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
- ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, Data2VecTextForCausalLM, Data2VecTextConfig
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("facebook/data2vec-text-base")
- >>> config = Data2VecTextConfig.from_pretrained("facebook/data2vec-text-base")
- >>> config.is_decoder = True
- >>> model = Data2VecTextForCausalLM.from_pretrained("facebook/data2vec-text-base", config=config)
-
- >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> prediction_logits = outputs.logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- if labels is not None:
- use_cache = False
-
- outputs = self.data2vec_text(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- prediction_scores = self.lm_head(sequence_output)
-
- lm_loss = None
- if labels is not None:
- # we are doing next-token prediction; shift prediction scores and input ids by one
- shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
- labels = labels[:, 1:].contiguous()
- loss_fct = CrossEntropyLoss()
-
- labels = labels.to(shifted_prediction_scores.device)
- lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((lm_loss,) + output) if lm_loss is not None else output
-
- return CausalLMOutputWithCrossAttentions(
- loss=lm_loss,
- logits=prediction_scores,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- cross_attentions=outputs.cross_attentions,
- )
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs):
- input_shape = input_ids.shape
- # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
- if attention_mask is None:
- attention_mask = input_ids.new_ones(input_shape)
-
- # cut decoder_input_ids if past_key_values is used
- if past_key_values is not None:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
-
- return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values}
-
- def _reorder_cache(self, past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
- )
- return reordered_past
-
-
-@add_start_docstrings("""data2vec Model with a `language modeling` head on top.""", DATA2VECTEXT_START_DOCSTRING)
-class Data2VecTextForMaskedLM(Data2VecTextPreTrainedModel):
- _tied_weights_keys = ["lm_head.decoder.weight", "lm_head.decoder.bias"]
-
- def __init__(self, config):
- super().__init__(config)
-
- if config.is_decoder:
- logger.warning(
- "If you want to use `Data2VecTextForMaskedLM` make sure `config.is_decoder=False` for "
- "bi-directional self-attention."
- )
-
- self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)
- self.lm_head = Data2VecTextLMHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.lm_head.decoder
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(DATA2VECTEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="",
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- kwargs (`Dict[str, any]`, optional, defaults to *{}*):
- Used to hide legacy arguments that have been deprecated.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.data2vec_text(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = outputs[0]
- prediction_scores = self.lm_head(sequence_output)
-
- masked_lm_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
-
- labels = labels.to(prediction_scores.device)
- masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return MaskedLMOutput(
- loss=masked_lm_loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-# Copied from transformers.models.roberta.modeling_roberta.RobertaLMHead with Roberta->Data2VecText
-class Data2VecTextLMHead(nn.Module):
- """Data2VecText Head for masked language modeling."""
-
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
- self.bias = nn.Parameter(torch.zeros(config.vocab_size))
- self.decoder.bias = self.bias
-
- def forward(self, features, **kwargs):
- x = self.dense(features)
- x = gelu(x)
- x = self.layer_norm(x)
-
- # project back to size of vocabulary with bias
- x = self.decoder(x)
-
- return x
-
- def _tie_weights(self):
- # To tie those two weights if they get disconnected (on TPU or when the bias is resized)
- # For accelerate compatibility and to not break backward compatibility
- if self.decoder.bias.device.type == "meta":
- self.decoder.bias = self.bias
- else:
- self.bias = self.decoder.bias
-
-
-@add_start_docstrings(
- """
- Data2VecText Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- DATA2VECTEXT_START_DOCSTRING,
-)
-class Data2VecTextForSequenceClassification(Data2VecTextPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
-
- self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)
- self.classifier = Data2VecTextClassificationHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DATA2VECTEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.data2vec_text(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = outputs[0]
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
-
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Data2VecText Model with a multiple choice classification head on top (a linear layer on top of the pooled output
- and a softmax) e.g. for RocStories/SWAG tasks.
- """,
- DATA2VECTEXT_START_DOCSTRING,
-)
-class Data2VecTextForMultipleChoice(Data2VecTextPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.data2vec_text = Data2VecTextModel(config)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.classifier = nn.Linear(config.hidden_size, 1)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(
- DATA2VECTEXT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
- )
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
- flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
- flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- flat_inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.data2vec_text(
- flat_input_ids,
- position_ids=flat_position_ids,
- token_type_ids=flat_token_type_ids,
- attention_mask=flat_attention_mask,
- head_mask=head_mask,
- inputs_embeds=flat_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- pooled_output = outputs[1]
-
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
-
- labels = labels.to(reshaped_logits.device)
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Data2VecText Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g.
- for Named-Entity-Recognition (NER) tasks.
- """,
- DATA2VECTEXT_START_DOCSTRING,
-)
-class Data2VecTextForTokenClassification(Data2VecTextPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DATA2VECTEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.data2vec_text(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
-
- labels = labels.to(logits.device)
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-# Copied from transformers.models.roberta.modeling_roberta.RobertaClassificationHead with Roberta->Data2VecText
-class Data2VecTextClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
-
- def forward(self, features, **kwargs):
- x = features[:, 0, :] # take token (equiv. to [CLS])
- x = self.dropout(x)
- x = self.dense(x)
- x = torch.tanh(x)
- x = self.dropout(x)
- x = self.out_proj(x)
- return x
-
-
-@add_start_docstrings(
- """
- Data2VecText Model with a span classification head on top for extractive question-answering tasks like SQuAD (a
- linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- DATA2VECTEXT_START_DOCSTRING,
-)
-class Data2VecTextForQuestionAnswering(Data2VecTextPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.data2vec_text = Data2VecTextModel(config, add_pooling_layer=False)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DATA2VECTEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- start_positions: Optional[torch.LongTensor] = None,
- end_positions: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.data2vec_text(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
- """
- Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
- are ignored. This is modified from fairseq's `utils.make_positions`.
-
- Args:
- x: torch.Tensor x:
-
- Returns: torch.Tensor
- """
- # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
- mask = input_ids.ne(padding_idx).int()
- incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
- return incremental_indices.long() + padding_idx
diff --git a/transformers/models/data2vec/modeling_data2vec_vision.py b/transformers/models/data2vec/modeling_data2vec_vision.py
deleted file mode 100644
index 44088d498f60356890405bb5566fe622b9f86800..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/modeling_data2vec_vision.py
+++ /dev/null
@@ -1,1228 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Data2VecVision model."""
-
-
-import collections.abc
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutput,
- BaseModelOutputWithPooling,
- ImageClassifierOutput,
- SemanticSegmenterOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import find_pruneable_heads_and_indices, meshgrid, prune_linear_layer
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_data2vec_vision import Data2VecVisionConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "Data2VecVisionConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/data2vec-vision-base"
-_EXPECTED_OUTPUT_SHAPE = [1, 197, 768]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "facebook/data2vec-vision-base-ft1k"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "remote control, remote"
-
-
-from ..deprecated._archive_maps import DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-# Copied from transformers.models.beit.modeling_beit.BeitModelOutputWithPooling with Beit->Data2VecVision
-class Data2VecVisionModelOutputWithPooling(BaseModelOutputWithPooling):
- """
- Class for outputs of [`Data2VecVisionModel`].
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):
- Average of the last layer hidden states of the patch tokens (excluding the *[CLS]* token) if
- *config.use_mean_pooling* is set to True. If set to False, then the final hidden state of the *[CLS]* token
- will be returned.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->Data2VecVision
-class Data2VecVisionDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitEmbeddings with Beit->Data2VecVision
-class Data2VecVisionEmbeddings(nn.Module):
- """
- Construct the CLS token, position and patch embeddings. Optionally, also the mask token.
-
- """
-
- def __init__(self, config: Data2VecVisionConfig) -> None:
- super().__init__()
-
- self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
- if config.use_mask_token:
- self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
- else:
- self.mask_token = None
- self.patch_embeddings = Data2VecVisionPatchEmbeddings(config)
- num_patches = self.patch_embeddings.num_patches
- if config.use_absolute_position_embeddings:
- self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))
- else:
- self.position_embeddings = None
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.BoolTensor] = None) -> torch.Tensor:
- embeddings, (patch_height, patch_width) = self.patch_embeddings(
- pixel_values, self.position_embeddings[:, 1:, :] if self.position_embeddings is not None else None
- )
- batch_size, seq_len, _ = embeddings.size()
-
- if bool_masked_pos is not None:
- mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
- # replace the masked visual tokens by mask_tokens
- w = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
- embeddings = embeddings * (1 - w) + mask_tokens * w
-
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- if self.position_embeddings is not None:
- cls_tokens = cls_tokens + self.position_embeddings[:, :1, :]
-
- embeddings = torch.cat((cls_tokens, embeddings), dim=1)
-
- embeddings = self.dropout(embeddings)
-
- return embeddings, (patch_height, patch_width)
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitPatchEmbeddings with Beit->Data2VecVision
-class Data2VecVisionPatchEmbeddings(nn.Module):
- """
- This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
- `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
- Transformer.
- """
-
- def __init__(self, config):
- super().__init__()
- image_size, patch_size = config.image_size, config.patch_size
- num_channels, hidden_size = config.num_channels, config.hidden_size
-
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- patch_shape = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.num_patches = num_patches
- self.patch_shape = patch_shape
-
- self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
-
- def forward(self, pixel_values: torch.Tensor, position_embedding: Optional[torch.Tensor] = None) -> torch.Tensor:
- batch_size, num_channels, height, width = pixel_values.shape
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
-
- embeddings = self.projection(pixel_values)
- patch_height, patch_width = embeddings.shape[2], embeddings.shape[3]
-
- if position_embedding is not None:
- # interpolate the position embedding to the corresponding size
- position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1).permute(
- 0, 3, 1, 2
- )
- position_embedding = nn.functional.interpolate(
- position_embedding, size=(patch_height, patch_width), mode="bicubic"
- )
- embeddings = embeddings + position_embedding
-
- embeddings = embeddings.flatten(2).transpose(1, 2)
-
- return embeddings, (patch_height, patch_width)
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitSelfAttention with Beit->Data2VecVision
-class Data2VecVisionSelfAttention(nn.Module):
- def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None) -> None:
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
- f"heads {config.num_attention_heads}."
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size)
- self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=False)
- self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- if window_size:
- self.relative_position_bias = Data2VecVisionRelativePositionBias(config, window_size=window_size)
- else:
- self.relative_position_bias = None
-
- def transpose_for_scores(self, x):
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(*new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- relative_position_bias: Optional["Data2VecVisionRelativePositionBias"] = None,
- ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
- mixed_query_layer = self.query(hidden_states)
-
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
-
- # Add relative position bias if present.
- if self.relative_position_bias is not None:
- attention_scores = attention_scores + self.relative_position_bias().unsqueeze(0)
-
- # Add shared relative position bias if provided.
- if relative_position_bias is not None:
- attention_scores = attention_scores + relative_position_bias
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(*new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitSelfOutput with Beit->Data2VecVision
-class Data2VecVisionSelfOutput(nn.Module):
- """
- The residual connection is defined in Data2VecVisionLayer instead of here (as is the case with other models), due to the
- layernorm applied before each block.
- """
-
- def __init__(self, config: Data2VecVisionConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor, gamma=None) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitAttention with Beit->Data2VecVision
-class Data2VecVisionAttention(nn.Module):
- def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None) -> None:
- super().__init__()
- self.attention = Data2VecVisionSelfAttention(config, window_size=window_size)
- self.output = Data2VecVisionSelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.attention.query = prune_linear_layer(self.attention.query, index)
- self.attention.key = prune_linear_layer(self.attention.key, index)
- self.attention.value = prune_linear_layer(self.attention.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
- self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- relative_position_bias: Optional["Data2VecVisionRelativePositionBias"] = None,
- ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
- self_outputs = self.attention(hidden_states, head_mask, output_attentions, relative_position_bias)
-
- attention_output = self.output(self_outputs[0], hidden_states)
-
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitIntermediate with Beit->Data2VecVision
-class Data2VecVisionIntermediate(nn.Module):
- def __init__(self, config: Data2VecVisionConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitOutput with Beit->Data2VecVision
-class Data2VecVisionOutput(nn.Module):
- def __init__(self, config: Data2VecVisionConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitLayer with Beit->Data2VecVision,BEiT->Data2VecVision
-class Data2VecVisionLayer(nn.Module):
- """This corresponds to the Block class in the timm implementation."""
-
- def __init__(
- self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None, drop_path_rate: float = 0.0
- ) -> None:
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = Data2VecVisionAttention(config, window_size=window_size)
- self.intermediate = Data2VecVisionIntermediate(config)
- self.output = Data2VecVisionOutput(config)
- self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.drop_path = Data2VecVisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
- self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- init_values = config.layer_scale_init_value
- if init_values > 0:
- self.lambda_1 = nn.Parameter(init_values * torch.ones((config.hidden_size)), requires_grad=True)
- self.lambda_2 = nn.Parameter(init_values * torch.ones((config.hidden_size)), requires_grad=True)
- else:
- self.lambda_1, self.lambda_2 = None, None
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- relative_position_bias: Optional["Data2VecVisionRelativePositionBias"] = None,
- ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
- self_attention_outputs = self.attention(
- self.layernorm_before(hidden_states), # in Data2VecVision, layernorm is applied before self-attention
- head_mask,
- output_attentions=output_attentions,
- relative_position_bias=relative_position_bias,
- )
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- # apply lambda_1 if present
- if self.lambda_1 is not None:
- attention_output = self.lambda_1 * attention_output
-
- # first residual connection
- hidden_states = self.drop_path(attention_output) + hidden_states
-
- # in Data2VecVision, layernorm is also applied after self-attention
- layer_output = self.layernorm_after(hidden_states)
-
- layer_output = self.intermediate(layer_output)
- layer_output = self.output(layer_output)
-
- if self.lambda_2 is not None:
- layer_output = self.lambda_2 * layer_output
-
- # second residual connection
- layer_output = self.drop_path(layer_output) + hidden_states
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitRelativePositionBias with Beit->Data2VecVision
-class Data2VecVisionRelativePositionBias(nn.Module):
- def __init__(self, config: Data2VecVisionConfig, window_size: tuple) -> None:
- super().__init__()
- self.window_size = window_size
- self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
- self.relative_position_bias_table = nn.Parameter(
- torch.zeros(self.num_relative_distance, config.num_attention_heads)
- ) # 2*Wh-1 * 2*Ww-1, nH
- # cls to token & token 2 cls & cls to cls
-
- # get pair-wise relative position index for each token inside the window
- coords_h = torch.arange(window_size[0])
- coords_w = torch.arange(window_size[1])
- coords = torch.stack(meshgrid([coords_h, coords_w], indexing="ij")) # 2, Wh, Ww
- coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
- relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
- relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
- relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0
- relative_coords[:, :, 1] += window_size[1] - 1
- relative_coords[:, :, 0] *= 2 * window_size[1] - 1
- relative_position_index = torch.zeros(
- size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype
- )
- relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
- relative_position_index[0, 0:] = self.num_relative_distance - 3
- relative_position_index[0:, 0] = self.num_relative_distance - 2
- relative_position_index[0, 0] = self.num_relative_distance - 1
-
- self.register_buffer("relative_position_index", relative_position_index, persistent=False)
-
- def forward(self) -> torch.Tensor:
- relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
- self.window_size[0] * self.window_size[1] + 1, self.window_size[0] * self.window_size[1] + 1, -1
- ) # Wh*Ww,Wh*Ww,nH
-
- return relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitEncoder with Beit->Data2VecVision
-class Data2VecVisionEncoder(nn.Module):
- def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None) -> None:
- super().__init__()
- self.config = config
- if config.use_shared_relative_position_bias:
- self.relative_position_bias = Data2VecVisionRelativePositionBias(config, window_size=window_size)
- else:
- self.relative_position_bias = None
-
- # stochastic depth decay rule
- dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]
- self.layer = nn.ModuleList(
- [
- Data2VecVisionLayer(
- config,
- window_size=window_size if config.use_relative_position_bias else None,
- drop_path_rate=dpr[i],
- )
- for i in range(config.num_hidden_layers)
- ]
- )
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[tuple, BaseModelOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- layer_head_mask,
- output_attentions,
- )
- else:
- relative_position_bias = (
- self.relative_position_bias() if self.relative_position_bias is not None else None
- )
- layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions, relative_position_bias)
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitPreTrainedModel with Beit->Data2VecVision,beit->data2vec_vision
-class Data2VecVisionPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = Data2VecVisionConfig
- base_model_prefix = "data2vec_vision"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-DATA2VEC_VISION_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`Data2VecVisionConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DATA2VEC_VISION_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`BeitImageProcessor.__call__`] for details.
-
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Data2VecVision Model transformer outputting raw hidden-states without any specific head on top.",
- DATA2VEC_VISION_START_DOCSTRING,
-)
-# Copied from transformers.models.beit.modeling_beit.BeitModel with BEIT->DATA2VEC_VISION,Beit->Data2VecVision,True->False
-class Data2VecVisionModel(Data2VecVisionPreTrainedModel):
- def __init__(self, config: Data2VecVisionConfig, add_pooling_layer: bool = False) -> None:
- super().__init__(config)
- self.config = config
-
- self.embeddings = Data2VecVisionEmbeddings(config)
- self.encoder = Data2VecVisionEncoder(config, window_size=self.embeddings.patch_embeddings.patch_shape)
-
- self.layernorm = (
- nn.Identity() if config.use_mean_pooling else nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- )
- self.pooler = Data2VecVisionPooler(config) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.patch_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=Data2VecVisionModelOutputWithPooling,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, Data2VecVisionModelOutputWithPooling]:
- r"""
- bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- embedding_output, (patch_height, patch_width) = self.embeddings(pixel_values, bool_masked_pos)
-
- encoder_outputs = self.encoder(
- embedding_output,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,)
- return head_outputs + encoder_outputs[1:]
-
- return Data2VecVisionModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitPooler with Beit->Data2VecVision
-class Data2VecVisionPooler(nn.Module):
- def __init__(self, config: Data2VecVisionConfig) -> None:
- super().__init__()
- self.layernorm = (
- nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) if config.use_mean_pooling else None
- )
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- if self.layernorm is not None:
- # Mean pool the final hidden states of the patch tokens
- patch_tokens = hidden_states[:, 1:, :]
- pooled_output = self.layernorm(patch_tokens.mean(1))
- else:
- # Pool by simply taking the final hidden state of the [CLS] token
- pooled_output = hidden_states[:, 0]
-
- return pooled_output
-
-
-@add_start_docstrings(
- """
- Data2VecVision Model transformer with an image classification head on top (a linear layer on top of the average of
- the final hidden states of the patch tokens) e.g. for ImageNet.
- """,
- DATA2VEC_VISION_START_DOCSTRING,
-)
-# Copied from transformers.models.beit.modeling_beit.BeitForImageClassification with BEIT->DATA2VEC_VISION,Beit->Data2VecVision,beit->data2vec_vision
-class Data2VecVisionForImageClassification(Data2VecVisionPreTrainedModel):
- def __init__(self, config: Data2VecVisionConfig) -> None:
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.data2vec_vision = Data2VecVisionModel(config, add_pooling_layer=True)
-
- # Classifier head
- self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=ImageClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, ImageClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- outputs = self.data2vec_vision(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs.pooler_output if return_dict else outputs[1]
-
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitConvModule with Beit->Data2VecVision
-class Data2VecVisionConvModule(nn.Module):
- """
- A convolutional block that bundles conv/norm/activation layers. This block simplifies the usage of convolution
- layers, which are commonly used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
-
- Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.
- """
-
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- kernel_size: Union[int, Tuple[int, int]],
- padding: Union[int, Tuple[int, int], str] = 0,
- bias: bool = False,
- dilation: Union[int, Tuple[int, int]] = 1,
- ) -> None:
- super().__init__()
- self.conv = nn.Conv2d(
- in_channels=in_channels,
- out_channels=out_channels,
- kernel_size=kernel_size,
- padding=padding,
- bias=bias,
- dilation=dilation,
- )
- self.bn = nn.BatchNorm2d(out_channels)
- self.activation = nn.ReLU()
-
- def forward(self, input: torch.Tensor) -> torch.Tensor:
- output = self.conv(input)
- output = self.bn(output)
- output = self.activation(output)
-
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitPyramidPoolingBlock with Beit->Data2VecVision
-class Data2VecVisionPyramidPoolingBlock(nn.Module):
- def __init__(self, pool_scale: int, in_channels: int, channels: int) -> None:
- super().__init__()
- self.layers = [
- nn.AdaptiveAvgPool2d(pool_scale),
- Data2VecVisionConvModule(in_channels, channels, kernel_size=1),
- ]
- for i, layer in enumerate(self.layers):
- self.add_module(str(i), layer)
-
- def forward(self, input: torch.Tensor) -> torch.Tensor:
- hidden_state = input
- for layer in self.layers:
- hidden_state = layer(hidden_state)
- return hidden_state
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitPyramidPoolingModule with Beit->Data2VecVision
-class Data2VecVisionPyramidPoolingModule(nn.Module):
- """
- Pyramid Pooling Module (PPM) used in PSPNet.
-
- Args:
- pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid
- Module.
- in_channels (int): Input channels.
- channels (int): Channels after modules, before conv_seg.
- align_corners (bool): align_corners argument of F.interpolate.
-
- Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.
- """
-
- def __init__(self, pool_scales: Tuple[int, ...], in_channels: int, channels: int, align_corners: bool) -> None:
- super().__init__()
- self.pool_scales = pool_scales
- self.align_corners = align_corners
- self.in_channels = in_channels
- self.channels = channels
- self.blocks = []
- for i, pool_scale in enumerate(pool_scales):
- block = Data2VecVisionPyramidPoolingBlock(
- pool_scale=pool_scale, in_channels=in_channels, channels=channels
- )
- self.blocks.append(block)
- self.add_module(str(i), block)
-
- def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
- ppm_outs = []
- for ppm in self.blocks:
- ppm_out = ppm(x)
- upsampled_ppm_out = nn.functional.interpolate(
- ppm_out, size=x.size()[2:], mode="bilinear", align_corners=self.align_corners
- )
- ppm_outs.append(upsampled_ppm_out)
- return ppm_outs
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitUperHead with Beit->Data2VecVision
-class Data2VecVisionUperHead(nn.Module):
- """
- Unified Perceptual Parsing for Scene Understanding. This head is the implementation of
- [UPerNet](https://arxiv.org/abs/1807.10221).
-
- Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.
- """
-
- def __init__(self, config: Data2VecVisionConfig) -> None:
- super().__init__()
-
- self.pool_scales = config.pool_scales # e.g. (1, 2, 3, 6)
- self.in_channels = [config.hidden_size] * 4 # e.g. [768, 768, 768, 768]
- self.channels = config.hidden_size
- self.align_corners = False
- self.classifier = nn.Conv2d(self.channels, config.num_labels, kernel_size=1)
-
- # PSP Module
- self.psp_modules = Data2VecVisionPyramidPoolingModule(
- self.pool_scales,
- self.in_channels[-1],
- self.channels,
- align_corners=self.align_corners,
- )
- self.bottleneck = Data2VecVisionConvModule(
- self.in_channels[-1] + len(self.pool_scales) * self.channels,
- self.channels,
- kernel_size=3,
- padding=1,
- )
- # FPN Module
- self.lateral_convs = nn.ModuleList()
- self.fpn_convs = nn.ModuleList()
- for in_channels in self.in_channels[:-1]: # skip the top layer
- l_conv = Data2VecVisionConvModule(in_channels, self.channels, kernel_size=1)
- fpn_conv = Data2VecVisionConvModule(self.channels, self.channels, kernel_size=3, padding=1)
- self.lateral_convs.append(l_conv)
- self.fpn_convs.append(fpn_conv)
-
- self.fpn_bottleneck = Data2VecVisionConvModule(
- len(self.in_channels) * self.channels,
- self.channels,
- kernel_size=3,
- padding=1,
- )
-
- def psp_forward(self, inputs):
- x = inputs[-1]
- psp_outs = [x]
- psp_outs.extend(self.psp_modules(x))
- psp_outs = torch.cat(psp_outs, dim=1)
- output = self.bottleneck(psp_outs)
-
- return output
-
- def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
- # build laterals
- laterals = [lateral_conv(encoder_hidden_states[i]) for i, lateral_conv in enumerate(self.lateral_convs)]
-
- laterals.append(self.psp_forward(encoder_hidden_states))
-
- # build top-down path
- used_backbone_levels = len(laterals)
- for i in range(used_backbone_levels - 1, 0, -1):
- prev_shape = laterals[i - 1].shape[2:]
- laterals[i - 1] = laterals[i - 1] + nn.functional.interpolate(
- laterals[i], size=prev_shape, mode="bilinear", align_corners=self.align_corners
- )
-
- # build outputs
- fpn_outs = [self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels - 1)]
- # append psp feature
- fpn_outs.append(laterals[-1])
-
- for i in range(used_backbone_levels - 1, 0, -1):
- fpn_outs[i] = nn.functional.interpolate(
- fpn_outs[i], size=fpn_outs[0].shape[2:], mode="bilinear", align_corners=self.align_corners
- )
- fpn_outs = torch.cat(fpn_outs, dim=1)
- output = self.fpn_bottleneck(fpn_outs)
- output = self.classifier(output)
-
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitFCNHead with Beit->Data2VecVision
-class Data2VecVisionFCNHead(nn.Module):
- """
- Fully Convolution Networks for Semantic Segmentation. This head is implemented of
- [FCNNet](https://arxiv.org/abs/1411.4038>).
-
- Args:
- config (Data2VecVisionConfig): Configuration.
- in_channels
- kernel_size (int): The kernel size for convs in the head. Default: 3.
- dilation (int): The dilation rate for convs in the head. Default: 1.
-
-
- Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.
- """
-
- def __init__(
- self,
- config: Data2VecVisionConfig,
- in_index: int = 2,
- kernel_size: int = 3,
- dilation: Union[int, Tuple[int, int]] = 1,
- ) -> None:
- super().__init__()
- self.in_channels = config.hidden_size
- self.channels = config.auxiliary_channels
- self.num_convs = config.auxiliary_num_convs
- self.concat_input = config.auxiliary_concat_input
- self.in_index = in_index
-
- conv_padding = (kernel_size // 2) * dilation
- convs = []
- convs.append(
- Data2VecVisionConvModule(
- self.in_channels, self.channels, kernel_size=kernel_size, padding=conv_padding, dilation=dilation
- )
- )
- for i in range(self.num_convs - 1):
- convs.append(
- Data2VecVisionConvModule(
- self.channels, self.channels, kernel_size=kernel_size, padding=conv_padding, dilation=dilation
- )
- )
- if self.num_convs == 0:
- self.convs = nn.Identity()
- else:
- self.convs = nn.Sequential(*convs)
- if self.concat_input:
- self.conv_cat = Data2VecVisionConvModule(
- self.in_channels + self.channels, self.channels, kernel_size=kernel_size, padding=kernel_size // 2
- )
-
- self.classifier = nn.Conv2d(self.channels, config.num_labels, kernel_size=1)
-
- def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
- # just take the relevant feature maps
- hidden_states = encoder_hidden_states[self.in_index]
- output = self.convs(hidden_states)
- if self.concat_input:
- output = self.conv_cat(torch.cat([hidden_states, output], dim=1))
- output = self.classifier(output)
- return output
-
-
-@add_start_docstrings(
- """
- Data2VecVision Model transformer with a semantic segmentation head on top e.g. for ADE20k, CityScapes.
- """,
- DATA2VEC_VISION_START_DOCSTRING,
-)
-# Copied from transformers.models.beit.modeling_beit.BeitForSemanticSegmentation with BEIT->DATA2VEC_VISION,Beit->Data2VecVision,microsoft/beit-base-finetuned-ade-640-640->facebook/data2vec-vision-base,beit->data2vec_vision
-class Data2VecVisionForSemanticSegmentation(Data2VecVisionPreTrainedModel):
- def __init__(self, config: Data2VecVisionConfig) -> None:
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.data2vec_vision = Data2VecVisionModel(config, add_pooling_layer=False)
-
- # FPNs
- if len(self.config.out_indices) != 4:
- raise ValueError(
- "Data2VecVisionForSemanticSegmentation requires config.out_indices to be a list of 4 integers, "
- "specifying which features to use from the backbone. One can use [3, 5, 7, 11] in case of "
- "a base-sized architecture."
- )
- self.fpn1 = nn.Sequential(
- nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2),
- nn.BatchNorm2d(config.hidden_size),
- nn.GELU(),
- nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2),
- )
- self.fpn2 = nn.Sequential(
- nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2),
- )
- self.fpn3 = nn.Identity()
- self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)
-
- # Semantic segmentation head(s)
- self.decode_head = Data2VecVisionUperHead(config)
- self.auxiliary_head = Data2VecVisionFCNHead(config) if config.use_auxiliary_head else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def compute_loss(self, logits, auxiliary_logits, labels):
- # upsample logits to the images' original size
- upsampled_logits = nn.functional.interpolate(
- logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
- )
- if auxiliary_logits is not None:
- upsampled_auxiliary_logits = nn.functional.interpolate(
- auxiliary_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
- )
- # compute weighted loss
- loss_fct = CrossEntropyLoss(ignore_index=self.config.semantic_loss_ignore_index)
- main_loss = loss_fct(upsampled_logits, labels)
- loss = main_loss
- if auxiliary_logits is not None:
- auxiliary_loss = loss_fct(upsampled_auxiliary_logits, labels)
- loss += self.config.auxiliary_loss_weight * auxiliary_loss
-
- return loss
-
- @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=SemanticSegmenterOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, SemanticSegmenterOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
- Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy).
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, Data2VecVisionForSemanticSegmentation
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/data2vec-vision-base")
- >>> model = Data2VecVisionForSemanticSegmentation.from_pretrained("facebook/data2vec-vision-base")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
- >>> outputs = model(**inputs)
- >>> # logits are of shape (batch_size, num_labels, height, width)
- >>> logits = outputs.logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- outputs = self.data2vec_vision(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=True, # we need the intermediate hidden states
- return_dict=return_dict,
- )
-
- encoder_hidden_states = outputs.hidden_states if return_dict else outputs[1]
-
- # only keep certain features, and reshape
- # note that we do +1 as the encoder_hidden_states also includes the initial embeddings
- features = [feature for idx, feature in enumerate(encoder_hidden_states) if idx + 1 in self.config.out_indices]
- batch_size = pixel_values.shape[0]
- patch_resolution = self.config.image_size // self.config.patch_size
- features = [
- x[:, 1:, :].permute(0, 2, 1).reshape(batch_size, -1, patch_resolution, patch_resolution) for x in features
- ]
-
- # apply FPNs
- ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
- for i in range(len(features)):
- features[i] = ops[i](features[i])
-
- logits = self.decode_head(features)
-
- auxiliary_logits = None
- if self.auxiliary_head is not None:
- auxiliary_logits = self.auxiliary_head(features)
-
- loss = None
- if labels is not None:
- if self.config.num_labels == 1:
- raise ValueError("The number of labels should be greater than one")
- else:
- loss = self.compute_loss(logits, auxiliary_logits, labels)
-
- if not return_dict:
- if output_hidden_states:
- output = (logits,) + outputs[1:]
- else:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SemanticSegmenterOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/data2vec/modeling_tf_data2vec_vision.py b/transformers/models/data2vec/modeling_tf_data2vec_vision.py
deleted file mode 100644
index e65a61fae5f881df1028bb8b23936ff86bb171d0..0000000000000000000000000000000000000000
--- a/transformers/models/data2vec/modeling_tf_data2vec_vision.py
+++ /dev/null
@@ -1,1717 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 Data2Vec Vision model."""
-
-
-from __future__ import annotations
-
-import collections.abc
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFBaseModelOutputWithPooling,
- TFSemanticSegmenterOutput,
- TFSequenceClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFModelInputType,
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import shape_list, stable_softmax
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_data2vec_vision import Data2VecVisionConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "Data2VecVisionConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/data2vec-vision-base"
-_EXPECTED_OUTPUT_SHAPE = [1, 197, 768]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "facebook/data2vec-vision-base-ft1k"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "remote control, remote"
-
-
-@dataclass
-class TFData2VecVisionModelOutputWithPooling(TFBaseModelOutputWithPooling):
- """
- Class for outputs of [`TFData2VecVisionModel`].
-
- Args:
- last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- pooler_output (`tf.Tensor` of shape `(batch_size, hidden_size)`):
- Average of the last layer hidden states of the patch tokens (excluding the *[CLS]* token) if
- *config.use_mean_pooling* is set to True. If set to False, then the final hidden state of the *[CLS]* token
- will be returned.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- last_hidden_state: tf.Tensor = None
- pooler_output: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-class TFData2VecVisionDropPath(keras.layers.Layer):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
- References:
- (1) github.com:rwightman/pytorch-image-models
- """
-
- def __init__(self, drop_path, **kwargs):
- super().__init__(**kwargs)
- self.drop_path = drop_path
-
- def call(self, x, training=None):
- if training:
- keep_prob = 1 - self.drop_path
- shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)
- random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)
- random_tensor = tf.floor(random_tensor)
- return (x / keep_prob) * random_tensor
- return x
-
-
-class TFData2VecVisionEmbeddings(keras.layers.Layer):
- """
- Construct the CLS token, position and patch embeddings. Optionally, also the mask token.
-
- """
-
- def __init__(self, config: Data2VecVisionConfig, **kwargs):
- super().__init__(**kwargs)
- self.config = config
-
- self.patch_embeddings = TFData2VecVisionPatchEmbeddings(config, name="patch_embeddings")
- self.num_patches = self.patch_embeddings.num_patches
- self.config = config
-
- self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
-
- def build(self, input_shape=None):
- self.cls_token = self.add_weight(
- shape=(1, 1, self.config.hidden_size),
- initializer=tf.random_normal_initializer(stddev=self.config.initializer_range),
- trainable=True,
- name="cls_token",
- )
- if self.config.use_mask_token:
- self.mask_token = self.add_weight(
- shape=(1, 1, self.config.hidden_size),
- initializer=tf.random_normal_initializer(stddev=self.config.initializer_range),
- trainable=True,
- name="mask_token",
- )
- else:
- self.mask_token = None
-
- if self.config.use_absolute_position_embeddings:
- self.position_embeddings = self.add_weight(
- shape=(1, self.num_patches + 1, self.config.hidden_size),
- initializer=tf.random_normal_initializer(stddev=self.config.initializer_range),
- trainable=True,
- name="position_embeddings",
- )
- else:
- self.position_embeddings = None
-
- if self.built:
- return
- self.built = True
- if getattr(self, "patch_embeddings", None) is not None:
- with tf.name_scope(self.patch_embeddings.name):
- self.patch_embeddings.build(None)
-
- def call(self, pixel_values: tf.Tensor, bool_masked_pos: tf.Tensor | None = None) -> tf.Tensor:
- embeddings = self.patch_embeddings(pixel_values)
- batch_size, seq_len, projection_dim = shape_list(embeddings)
-
- cls_tokens = tf.tile(self.cls_token, (batch_size, 1, 1))
-
- if bool_masked_pos is not None:
- mask_tokens = tf.broadcast_to(self.mask_token, (batch_size, seq_len, projection_dim))
- # replace the masked visual tokens by mask_tokens
- w = bool_masked_pos[..., None]
- w = tf.cast(w, mask_tokens.dtype)
- # since TF doesn't support eager tensor assignment
- embeddings = embeddings * (1 - w) + mask_tokens * w
-
- embeddings = tf.concat([cls_tokens, embeddings], axis=1)
- if self.position_embeddings is not None:
- embeddings = embeddings + self.position_embeddings
- embeddings = self.dropout(embeddings)
-
- return embeddings
-
-
-class TFData2VecVisionPatchEmbeddings(keras.layers.Layer):
- """
- Image to Patch Embedding.
- """
-
- def __init__(self, config: Data2VecVisionConfig, **kwargs):
- super().__init__(**kwargs)
- self.config = config
-
- image_size, patch_size = config.image_size, config.patch_size
- num_channels, hidden_size = config.num_channels, config.hidden_size
-
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- patch_shape = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_patches = num_patches
- self.patch_shape = patch_shape
- self.num_channels = num_channels
-
- self.projection = keras.layers.Conv2D(
- filters=hidden_size,
- kernel_size=patch_size,
- strides=patch_size,
- padding="valid",
- data_format="channels_last",
- kernel_initializer="glorot_uniform", # following torch.nn.Linear
- bias_initializer="zeros",
- name="projection",
- )
-
- def call(self, pixel_values: tf.Tensor, training: bool = False) -> tf.Tensor:
- batch_size, num_channels, height, width = shape_list(pixel_values)
- if tf.executing_eagerly():
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the"
- " configuration."
- )
- if height != self.image_size[0] or width != self.image_size[1]:
- raise ValueError(
- f"Input image size ({height}*{width}) doesn't match model"
- f" ({self.image_size[0]}*{self.image_size[1]})."
- )
-
- # When running on CPU, `keras.layers.Conv2D` doesn't support `NCHW` format.
- # So change the input format from `NCHW` to `NHWC`.
- # shape = (batch_size, in_height, in_width, in_channels=num_channels)
- pixel_values = tf.transpose(pixel_values, perm=(0, 2, 3, 1))
-
- projection = self.projection(pixel_values)
-
- # Change the 2D spatial dimensions to a single temporal dimension.
- # shape = (batch_size, num_patches, out_channels=embed_dim)
- num_patches = (width // self.patch_size[1]) * (height // self.patch_size[0])
-
- return tf.reshape(tensor=projection, shape=(batch_size, num_patches, -1))
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "projection", None) is not None:
- with tf.name_scope(self.projection.name):
- self.projection.build([None, None, None, self.num_channels])
-
-
-class TFData2VecVisionSelfAttention(keras.layers.Layer):
- def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None, **kwargs):
- super().__init__(**kwargs)
-
- if config.hidden_size % config.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number "
- f"of attention heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.sqrt_att_head_size = math.sqrt(self.attention_head_size)
-
- self.query = keras.layers.Dense(
- units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"
- )
- self.key = keras.layers.Dense(
- units=self.all_head_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="key",
- use_bias=False,
- )
- self.value = keras.layers.Dense(
- units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"
- )
- self.dropout = keras.layers.Dropout(rate=config.attention_probs_dropout_prob)
-
- if window_size:
- self.relative_position_bias = TFData2VecVisionRelativePositionBias(
- config, window_size=window_size, name="relative_position_bias"
- )
- else:
- self.relative_position_bias = None
- self.config = config
-
- def transpose_for_scores(self, tensor: tf.Tensor, batch_size: int) -> tf.Tensor:
- # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]
- tensor = tf.reshape(tensor=tensor, shape=(batch_size, -1, self.num_attention_heads, self.attention_head_size))
-
- # Transpose the tensor from [batch_size, seq_length, num_attention_heads, attention_head_size] to [batch_size, num_attention_heads, seq_length, attention_head_size]
- return tf.transpose(tensor, perm=[0, 2, 1, 3])
-
- def call(
- self,
- hidden_states: tf.Tensor,
- head_mask: tf.Tensor,
- output_attentions: bool,
- relative_position_bias: Optional["TFData2VecVisionRelativePositionBias"] = None,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- batch_size = shape_list(hidden_states)[0]
- mixed_query_layer = self.query(inputs=hidden_states)
- mixed_key_layer = self.key(inputs=hidden_states)
- mixed_value_layer = self.value(inputs=hidden_states)
- query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
- key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
- value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- # (batch size, num_heads, seq_len_q, seq_len_k)
- attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
- attention_scores = attention_scores / self.sqrt_att_head_size
-
- # Add relative position bias if present.
- if self.relative_position_bias is not None:
- # Passing `0.0` to the `relative_position_bias()` layer because otherwise Keras
- # might complain about `Layer.call()` not being invoked properly. In this case this input
- # i.e., 0.0 is not going to be used in any calculations so we're safe.
- attention_scores = attention_scores + self.relative_position_bias(0.0)[None, ...]
-
- # Add shared relative position bias if provided.
- if relative_position_bias is not None:
- attention_scores = attention_scores + relative_position_bias
-
- # Normalize the attention scores to probabilities.
- attention_probs = stable_softmax(logits=attention_scores, axis=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(inputs=attention_probs, training=training)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = tf.multiply(attention_probs, head_mask)
-
- attention_output = tf.matmul(attention_probs, value_layer)
- attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])
-
- # (batch_size, seq_len_q, all_head_size)
- attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
- outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "query", None) is not None:
- with tf.name_scope(self.query.name):
- self.query.build([None, None, self.config.hidden_size])
- if getattr(self, "key", None) is not None:
- with tf.name_scope(self.key.name):
- self.key.build([None, None, self.config.hidden_size])
- if getattr(self, "value", None) is not None:
- with tf.name_scope(self.value.name):
- self.value.build([None, None, self.config.hidden_size])
- if getattr(self, "relative_position_bias", None) is not None:
- with tf.name_scope(self.relative_position_bias.name):
- self.relative_position_bias.build(None)
-
-
-class TFData2VecVisionSelfOutput(keras.layers.Layer):
- """
- The residual connection is defined in TFData2VecVisionLayer instead of here (as is the case with other models), due
- to the layernorm applied before each block.
- """
-
- def __init__(self, config: Data2VecVisionConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, gamma=None, training: bool = False) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.dropout(inputs=hidden_states, training=training)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFData2VecVisionAttention(keras.layers.Layer):
- def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None, **kwargs):
- super().__init__(**kwargs)
-
- self.attention = TFData2VecVisionSelfAttention(config, window_size=window_size, name="attention")
- self.dense_output = TFData2VecVisionSelfOutput(config, name="output")
-
- def prune_heads(self, heads):
- raise NotImplementedError
-
- def call(
- self,
- input_tensor: tf.Tensor,
- head_mask: tf.Tensor,
- output_attentions: bool,
- relative_position_bias: Optional["TFData2VecVisionRelativePositionBias"] = None,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- self_outputs = self.attention(
- hidden_states=input_tensor,
- head_mask=head_mask,
- output_attentions=output_attentions,
- relative_position_bias=relative_position_bias,
- training=training,
- )
- attention_output = self.dense_output(
- hidden_states=self_outputs[0], input_tensor=input_tensor, training=training
- )
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "dense_output", None) is not None:
- with tf.name_scope(self.dense_output.name):
- self.dense_output.build(None)
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTIntermediate with ViT->Data2VecVision
-class TFData2VecVisionIntermediate(keras.layers.Layer):
- def __init__(self, config: Data2VecVisionConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
-
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.intermediate_act_fn = config.hidden_act
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFData2VecVisionOutput(keras.layers.Layer):
- def __init__(self, config: Data2VecVisionConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.dropout(inputs=hidden_states, training=training)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.intermediate_size])
-
-
-class TFData2VecVisionLayer(keras.layers.Layer):
- """This corresponds to the Block class in the timm implementation."""
-
- def __init__(
- self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None, drop_path_rate: float = 0.0, **kwargs
- ):
- super().__init__(**kwargs)
- self.config = config
-
- self.attention = TFData2VecVisionAttention(config, window_size=window_size, name="attention")
- self.intermediate = TFData2VecVisionIntermediate(config, name="intermediate")
- self.data2vec_output = TFData2VecVisionOutput(config, name="output")
-
- self.layernorm_before = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm_before")
- self.layernorm_after = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm_after")
- # Using `layers.Activation` instead of `tf.identity` to better control `training`
- # behaviour.
- self.drop_path = (
- TFData2VecVisionDropPath(drop_path_rate, name="drop_path")
- if drop_path_rate > 0.0
- else keras.layers.Activation("linear", name="drop_path")
- )
- self.init_values = config.layer_scale_init_value
-
- def build(self, input_shape: tf.TensorShape = None):
- if self.init_values > 0:
- self.lambda_1 = self.add_weight(
- shape=(self.config.hidden_size),
- initializer="ones",
- trainable=True,
- name="lambda_1",
- )
- self.lambda_2 = self.add_weight(
- shape=(self.config.hidden_size),
- initializer="ones",
- trainable=True,
- name="lambda_2",
- )
- self.lambda_1.assign(self.init_values * tf.ones((self.config.hidden_size)))
- self.lambda_2.assign(self.init_values * tf.ones((self.config.hidden_size)))
- else:
- self.lambda_1, self.lambda_2 = None, None
-
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "intermediate", None) is not None:
- with tf.name_scope(self.intermediate.name):
- self.intermediate.build(None)
- if getattr(self, "data2vec_output", None) is not None:
- with tf.name_scope(self.data2vec_output.name):
- self.data2vec_output.build(None)
- if getattr(self, "layernorm_before", None) is not None:
- with tf.name_scope(self.layernorm_before.name):
- self.layernorm_before.build([None, None, self.config.hidden_size])
- if getattr(self, "layernorm_after", None) is not None:
- with tf.name_scope(self.layernorm_after.name):
- self.layernorm_after.build([None, None, self.config.hidden_size])
- if getattr(self, "drop_path", None) is not None:
- with tf.name_scope(self.drop_path.name):
- self.drop_path.build(None)
-
- def call(
- self,
- hidden_states: tf.Tensor,
- head_mask: tf.Tensor,
- output_attentions: bool,
- relative_position_bias: Optional["TFData2VecVisionRelativePositionBias"] = None,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- self_attention_outputs = self.attention(
- # in Data2VecVision, layernorm is applied before self-attention
- input_tensor=self.layernorm_before(inputs=hidden_states),
- head_mask=head_mask,
- output_attentions=output_attentions,
- relative_position_bias=relative_position_bias,
- training=training,
- )
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- # apply lambda_1 if present
- if self.lambda_1 is not None:
- attention_output = self.lambda_1 * attention_output
-
- # first residual connection
- hidden_states = self.drop_path(attention_output) + hidden_states
-
- # in Data2VecVision, layernorm is also applied after self-attention
- layer_output = self.layernorm_after(hidden_states)
-
- layer_output = self.intermediate(layer_output)
- layer_output = self.data2vec_output(layer_output)
-
- if self.lambda_2 is not None:
- layer_output = self.lambda_2 * layer_output
-
- # second residual connection
- layer_output = self.drop_path(layer_output) + hidden_states
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-# Taken and modified from here:
-# https://github.com/leondgarse/keras_cv_attention_models/blob/main/keras_cv_attention_models/beit/beit.py#L28
-class TFData2VecVisionRelativePositionBias(keras.layers.Layer):
- def __init__(self, config: Data2VecVisionConfig, window_size: tuple, **kwargs) -> None:
- super().__init__(**kwargs)
- self.config = config
-
- self.window_size = window_size
- # +3 for cls_token_pos_len
- # window_size can be something like (14, 14)
- self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
-
- self.relative_position_index = self.get_position_index()
-
- def build(self, input_shape):
- self.relative_position_bias_table = self.add_weight(
- shape=(self.num_relative_distance, self.config.num_attention_heads),
- initializer="zeros",
- trainable=True,
- name="relative_position_bias_table",
- ) # [2*Wh-1 * 2*Ww-1, nH]
- # cls to token & token 2 cls & cls to cls
-
- super().build(input_shape)
-
- def get_position_index(self):
- # get pair-wise relative position index for each token inside the window
- xx, yy = tf.meshgrid(range(self.window_size[0]), range(self.window_size[1]))
- coords = tf.stack([yy, xx], axis=0) # [2, Wh, Ww]
- coords_flatten = tf.reshape(coords, [2, -1]) # [2, Wh*Ww]
-
- relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # [2, Wh*Ww, Wh*Ww]
- relative_coords = tf.transpose(relative_coords, perm=[1, 2, 0]) # [Wh*Ww, Wh*Ww, 2]
-
- xx = (relative_coords[:, :, 0] + self.window_size[0] - 1) * (2 * self.window_size[1] - 1)
- yy = relative_coords[:, :, 1] + self.window_size[1] - 1
- relative_coords = tf.stack([xx, yy], axis=-1)
-
- relative_position_index = tf.reduce_sum(relative_coords, axis=-1) # [Wh*Ww, Wh*Ww]
-
- top = tf.ones((1, relative_position_index.shape[1]), dtype=relative_position_index.dtype) * (
- self.num_relative_distance - 3
- )
- left = tf.ones((relative_position_index.shape[0], 1), dtype=relative_position_index.dtype) * (
- self.num_relative_distance - 2
- )
- corner = tf.ones((1, 1), dtype=relative_position_index.dtype) * (self.num_relative_distance - 1)
-
- left_corner = tf.concat([corner, left], axis=0)
- relative_position_index = tf.concat([top, relative_position_index], axis=0)
- relative_position_index = tf.concat([left_corner, relative_position_index], axis=1) # [Wh*Ww + 1, Wh*Ww + 1]
- return relative_position_index
-
- def call(self, inputs=None) -> tf.Tensor:
- relative_position_bias = tf.gather(self.relative_position_bias_table, self.relative_position_index, axis=0)
- return tf.transpose(relative_position_bias, [2, 0, 1])
-
-
-class TFData2VecVisionEncoder(keras.layers.Layer):
- def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None, **kwargs):
- super().__init__(**kwargs)
- self.config = config
- if config.use_shared_relative_position_bias:
- self.relative_position_bias = TFData2VecVisionRelativePositionBias(
- config, window_size=window_size, name="relative_position_bias"
- )
- else:
- self.relative_position_bias = None
-
- # stochastic depth decay rule
- dpr = list(tf.linspace(0.0, config.drop_path_rate, config.num_hidden_layers))
- self.layer = [
- TFData2VecVisionLayer(
- config,
- window_size=window_size if config.use_relative_position_bias else None,
- drop_path_rate=dpr[i],
- name=f"layer_._{i}",
- )
- for i in range(config.num_hidden_layers)
- ]
-
- def call(
- self,
- hidden_states: tf.Tensor,
- head_mask: tf.Tensor | None = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[tuple, TFBaseModelOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
- # Passing `0.0` to the `relative_position_bias()` layer because otherwise Keras
- # might complain about `Layer.call()` not being invoked properly. In this case this input
- # i.e., 0.0 is not going to be used in any calculations so we're safe.
- relative_position_bias = (
- self.relative_position_bias(0.0) if self.relative_position_bias is not None else None
- )
- layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions, relative_position_bias)
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
-
- return TFBaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "relative_position_bias", None) is not None:
- with tf.name_scope(self.relative_position_bias.name):
- self.relative_position_bias.build(None)
- if getattr(self, "layer", None) is not None:
- for layer in self.layer:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-@keras_serializable
-class TFData2VecVisionMainLayer(keras.layers.Layer):
- config_class = Data2VecVisionConfig
-
- def __init__(self, config: Data2VecVisionConfig, add_pooling_layer: bool = True, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.add_pooling_layer = add_pooling_layer
-
- self.embeddings = TFData2VecVisionEmbeddings(config, name="embeddings")
- self.encoder = TFData2VecVisionEncoder(
- config, window_size=self.embeddings.patch_embeddings.patch_shape, name="encoder"
- )
- self.layernorm = (
- tf.identity
- if config.use_mean_pooling
- else keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
- )
-
- # We are setting the `data_format` like so because from here on we will revert to the
- # NCHW output format
- self.pooler = TFData2VecVisionPooler(config, name="pooler") if add_pooling_layer else None
-
- def get_input_embeddings(self) -> keras.layers.Layer:
- return self.embeddings.patch_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- raise NotImplementedError
-
- @unpack_inputs
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- bool_masked_pos: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[tuple, TFData2VecVisionModelOutputWithPooling]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.config.num_hidden_layers
-
- embedding_output = self.embeddings(pixel_values, bool_masked_pos, training=training)
-
- encoder_outputs = self.encoder(
- embedding_output,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,)
- return head_outputs + encoder_outputs[1:]
-
- return TFData2VecVisionModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "layernorm", None) is not None:
- if hasattr(self.layernorm, "name"):
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build((None, self.config.hidden_size))
- if getattr(self, "pooler", None) is not None:
- with tf.name_scope(self.pooler.name):
- self.pooler.build(None)
-
-
-class TFData2VecVisionPooler(keras.layers.Layer):
- def __init__(self, config: Data2VecVisionConfig, **kwargs):
- super().__init__(**kwargs)
- self.layernorm = (
- keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
- if config.use_mean_pooling
- else None
- )
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- if self.layernorm is not None:
- # Mean pool the final hidden states of the patch tokens
- patch_tokens = hidden_states[:, 1:, :]
- pooled_output = self.layernorm(tf.reduce_mean(patch_tokens, axis=1))
- else:
- # Pool by simply taking the final hidden state of the [CLS] token
- pooled_output = hidden_states[:, 0]
-
- return pooled_output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "layernorm", None) is not None:
- if hasattr(self.layernorm, "name"):
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build((None, self.config.hidden_size))
-
-
-class TFData2VecVisionPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = Data2VecVisionConfig
- base_model_prefix = "data2vec_vision"
- main_input_name = "pixel_values"
- _keys_to_ignore_on_load_unexpected = [r"relative_position_index"]
-
-
-DATA2VEC_VISION_START_DOCSTRING = r"""
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.).
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `pixel_values` only and nothing else: `model(pixel_values)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([pixel_values, attention_mask])` or `model([pixel_values, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"pixel_values": pixel_values, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Args:
- config ([`Data2VecVisionConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DATA2VEC_VISION_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`np.ndarray`, `tf.Tensor`, `List[tf.Tensor]` `Dict[str, tf.Tensor]` or `Dict[str, np.ndarray]` and each example must have the shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`BeitImageProcessor.__call__`] for details.
-
- head_mask (`np.ndarray` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
-
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
-
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. This argument can be used
- in eager mode, in graph mode the value will always be set to True.
-
- training (`bool`, *optional*, defaults to `False``):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare Data2VecVision Model transformer outputting raw hidden-states without any specific head on top.",
- DATA2VEC_VISION_START_DOCSTRING,
-)
-class TFData2VecVisionModel(TFData2VecVisionPreTrainedModel):
- def __init__(self, config: Data2VecVisionConfig, add_pooling_layer: bool = False, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.config = config
-
- self.data2vec_vision = TFData2VecVisionMainLayer(
- config, add_pooling_layer=add_pooling_layer, name="data2vec_vision"
- )
-
- def get_input_embeddings(self):
- return self.data2vec_vision.get_input_embeddings()
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFData2VecVisionModelOutputWithPooling,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- bool_masked_pos: tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[tuple, TFData2VecVisionModelOutputWithPooling]:
- r"""
- bool_masked_pos (`tf.Tensor` of shape `(batch_size, num_patches)`, *optional*):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
- """
- outputs = self.data2vec_vision(
- pixel_values=pixel_values,
- bool_masked_pos=bool_masked_pos,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "data2vec_vision", None) is not None:
- with tf.name_scope(self.data2vec_vision.name):
- self.data2vec_vision.build(None)
-
-
-@add_start_docstrings(
- """
- Data2VecVision Model transformer with an image classification head on top (a linear layer on top of the average of
- the final hidden states of the patch tokens) e.g. for ImageNet.
- """,
- DATA2VEC_VISION_START_DOCSTRING,
-)
-class TFData2VecVisionForImageClassification(TFData2VecVisionPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: Data2VecVisionConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
- self.data2vec_vision = TFData2VecVisionMainLayer(config, add_pooling_layer=True, name="data2vec_vision")
-
- # Classifier head
- self.classifier = keras.layers.Dense(
- units=config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- name="classifier",
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def call(
- self,
- pixel_values: TFModelInputType | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFSequenceClassifierOutput, tuple]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.data2vec_vision(
- pixel_values=pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- pooled_output = outputs.pooler_output if return_dict else outputs[1]
- logits = self.classifier(pooled_output)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "data2vec_vision", None) is not None:
- with tf.name_scope(self.data2vec_vision.name):
- self.data2vec_vision.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-class TFData2VecVisionConvModule(keras.layers.Layer):
- """
- A convolutional block that bundles conv/norm/activation layers. This block simplifies the usage of convolution
- layers, which are commonly used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
-
- Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.
- """
-
- def __init__(
- self,
- in_channels: int,
- out_channels: int,
- kernel_size: Union[int, Tuple[int, int]],
- padding: str = "valid",
- bias: bool = False,
- dilation: Union[int, Tuple[int, int]] = 1,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.conv = keras.layers.Conv2D(
- filters=out_channels,
- kernel_size=kernel_size,
- padding=padding,
- use_bias=bias,
- dilation_rate=dilation,
- name="conv",
- )
- self.bn = keras.layers.BatchNormalization(name="bn", momentum=0.9, epsilon=1e-5)
- self.activation = tf.nn.relu
- self.in_channels = in_channels
- self.out_channels = out_channels
-
- def call(self, input: tf.Tensor) -> tf.Tensor:
- output = self.conv(input)
- output = self.bn(output)
- output = self.activation(output)
- return output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "conv", None) is not None:
- with tf.name_scope(self.conv.name):
- self.conv.build([None, None, None, self.in_channels])
- if getattr(self, "bn", None) is not None:
- with tf.name_scope(self.bn.name):
- self.bn.build((None, None, None, self.out_channels))
-
-
-class TFAdaptiveAvgPool2D(keras.layers.Layer):
- def __init__(self, output_dims: Tuple[int, int], input_ordering: str = "NHWC", **kwargs):
- super().__init__(**kwargs)
- self.output_dims = output_dims
- self.input_ordering = input_ordering
- if input_ordering not in ("NCHW", "NHWC"):
- raise ValueError("Unrecognized input_ordering, should be 'NCHW' or 'NHWC'!")
- self.h_axis = input_ordering.index("H")
- self.w_axis = input_ordering.index("W")
-
- def pseudo_1d_pool(self, inputs: tf.Tensor, h_pooling: bool):
- # Figure out which axis we're pooling on
- if h_pooling:
- axis = self.h_axis
- output_dim = self.output_dims[0]
- else:
- axis = self.w_axis
- output_dim = self.output_dims[1]
- input_dim = inputs.shape[axis]
-
- # Figure out the potential pooling windows
- # This is the key idea - the torch op always uses only two
- # consecutive pooling window sizes, like 3 and 4. Therefore,
- # if we pool with both possible sizes, we simply need to gather
- # the 'correct' pool at each position to reimplement the torch op.
- small_window = math.ceil(input_dim / output_dim)
- big_window = small_window + 1
- if h_pooling:
- output_dim = self.output_dims[0]
- small_window_shape = (small_window, 1)
- big_window_shape = (big_window, 1)
- else:
- output_dim = self.output_dims[1]
- small_window_shape = (1, small_window)
- big_window_shape = (1, big_window)
-
- # For resizes to 1, or integer resizes, we can take quick shortcuts
- if output_dim == input_dim:
- return inputs
- elif output_dim == 1:
- return tf.reduce_mean(inputs, axis=axis, keepdims=True)
- elif input_dim % output_dim == 0:
- return tf.nn.avg_pool2d(
- inputs,
- ksize=small_window_shape,
- strides=small_window_shape,
- padding="VALID",
- data_format=self.input_ordering,
- )
- # When upscaling by an integer factor we can also take a quick shortcut
- elif output_dim > input_dim and output_dim % input_dim == 0:
- return tf.repeat(inputs, repeats=output_dim // input_dim, axis=axis)
-
- # For non-integer resizes, we pool with both possible window sizes and concatenate them
- if output_dim < input_dim:
- small_pool = tf.nn.avg_pool2d(
- inputs, ksize=small_window_shape, strides=1, padding="VALID", data_format=self.input_ordering
- )
- big_pool = tf.nn.avg_pool2d(
- inputs, ksize=big_window_shape, strides=1, padding="VALID", data_format=self.input_ordering
- )
- both_pool = tf.concat([small_pool, big_pool], axis=axis)
- else:
- # When we're actually upscaling instead, then we build the pools a bit differently
- small_pool = inputs
- big_pool = tf.nn.avg_pool2d(
- inputs, ksize=big_window_shape, strides=1, padding="VALID", data_format=self.input_ordering
- )
- both_pool = tf.concat([small_pool, big_pool], axis=axis)
-
- # We compute vectors of the start and end positions for each pooling window
- # Each (start, end) pair here corresponds to a single output position
- window_starts = tf.math.floor((tf.range(output_dim, dtype=tf.float32) * input_dim) / output_dim)
- window_starts = tf.cast(window_starts, tf.int64)
- window_ends = tf.math.ceil((tf.range(1, output_dim + 1, dtype=tf.float32) * input_dim) / output_dim)
- window_ends = tf.cast(window_ends, tf.int64)
-
- # pool_selector is a boolean array of shape (output_dim,) where 1 indicates that output position
- # has a big receptive field and 0 indicates that that output position has a small receptive field
- pool_selector = tf.cast(window_ends - window_starts - small_window, tf.bool)
-
- # Since we concatenated the small and big pools, we need to do a bit of
- # pointer arithmetic to get the indices of the big pools
- small_indices = window_starts
- big_indices = window_starts + small_pool.shape[axis]
-
- # Finally, we use the pool_selector to generate a list of indices, one per output position
- gather_indices = tf.where(pool_selector, big_indices, small_indices)
-
- # Gathering from those indices yields the final, correct pooling
- return tf.gather(both_pool, gather_indices, axis=axis)
-
- def call(self, inputs: tf.Tensor):
- if self.input_ordering == "NHWC":
- input_shape = inputs.shape[1:3]
- else:
- input_shape = inputs.shape[2:]
-
- # We break the task down into each possible case
- # Firstly, if we're resizing down to 1, it's just tf.reduce_mean
- if self.output_dims[0] == self.output_dims[1] == 1:
- if self.input_ordering == "NHWC":
- reduce_dims = [1, 2]
- else:
- reduce_dims = [2, 3]
- return tf.reduce_mean(inputs, axis=reduce_dims, keepdims=True)
- # Secondly, if we're resizing by an integer factor on both dimensions, we can take a quick shortcut
- elif input_shape[0] % self.output_dims[0] == 0 and input_shape[1] % self.output_dims[1] == 0:
- h_resize = int(input_shape[0] // self.output_dims[0])
- w_resize = int(input_shape[1] // self.output_dims[1])
- return tf.nn.avg_pool2d(
- inputs,
- ksize=(h_resize, w_resize),
- strides=(h_resize, w_resize),
- padding="VALID",
- data_format=self.input_ordering,
- )
- else:
- # Finally, if we can't take the shortcut, we do a 1D pool on each axis. pseudo_1d_pool will take a shortcut
- # for dimensions where an integer resize is possible. It can also handle upscaling.
- h_pooled = self.pseudo_1d_pool(inputs, h_pooling=True)
- return self.pseudo_1d_pool(h_pooled, h_pooling=False)
-
-
-class TFData2VecVisionPyramidPoolingModule(keras.layers.Layer):
- """
- Pyramid Pooling Module (PPM) used in PSPNet.
-
- Args:
- pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid
- Module.
- channels (int): Channels after modules, before conv_seg.
-
- Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.
- """
-
- def __init__(self, pool_scales: Tuple[int, ...], in_channels: int, out_channels: int, **kwargs) -> None:
- super().__init__(**kwargs)
- self.pool_scales = pool_scales
- self.in_channels = in_channels
- self.out_channels = out_channels
-
- self.layer_list = []
- for idx, pool_scale in enumerate(pool_scales):
- pool_scale = pool_scale if isinstance(pool_scale, collections.abc.Iterable) else (pool_scale, pool_scale)
- self.layer_list.append(
- [
- TFAdaptiveAvgPool2D(output_dims=pool_scale),
- TFData2VecVisionConvModule(
- in_channels=in_channels, out_channels=self.out_channels, kernel_size=1, name=f"{idx}.1"
- ),
- ]
- )
-
- def call(self, x: tf.Tensor) -> List[tf.Tensor]:
- ppm_outs = []
- inputs = x
-
- for ppm in self.layer_list:
- for layer_module in ppm:
- ppm_out = layer_module(x)
- x = ppm_out
-
- upsampled_ppm_out = tf.image.resize(ppm_out, size=shape_list(inputs)[1:-1], method="bilinear")
- ppm_outs.append(upsampled_ppm_out)
- return ppm_outs
-
- def build(self, input_shape=None):
- for layer in self.layer_list:
- for layer_module in layer:
- with tf.name_scope(layer_module.name):
- layer_module.build(None)
-
-
-class TFData2VecVisionUperHead(keras.layers.Layer):
- """
- Unified Perceptual Parsing for Scene Understanding. This head is the implementation of
- [UPerNet](https://arxiv.org/abs/1807.10221).
-
- Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.
- """
-
- def __init__(self, config: Data2VecVisionConfig, **kwargs) -> None:
- super().__init__(**kwargs)
-
- self.pool_scales = config.pool_scales # e.g. (1, 2, 3, 6)
- self.in_channels = [config.hidden_size] * 4 # e.g. [768, 768, 768, 768]
- self.channels = config.hidden_size
- self.classifier = keras.layers.Conv2D(config.num_labels, kernel_size=1, name="classifier")
-
- # PSP Module
- self.psp_modules = TFData2VecVisionPyramidPoolingModule(
- self.pool_scales, self.in_channels[-1], self.channels, name="psp_modules"
- )
- self.bottleneck = TFData2VecVisionConvModule(
- self.in_channels[-1] + len(self.pool_scales) * self.channels,
- self.channels,
- kernel_size=3,
- padding="same",
- name="bottleneck",
- )
- # FPN Module
- self.lateral_convs = []
- self.fpn_convs = []
- for idx, in_channels in enumerate(self.in_channels[:-1]): # skip the top layer
- l_conv = TFData2VecVisionConvModule(
- in_channels, out_channels=self.channels, kernel_size=1, name=f"lateral_convs.{idx}"
- )
- fpn_conv = TFData2VecVisionConvModule(
- in_channels=self.channels,
- out_channels=self.channels,
- kernel_size=3,
- padding="same",
- name=f"fpn_convs.{idx}",
- )
- self.lateral_convs.append(l_conv)
- self.fpn_convs.append(fpn_conv)
-
- self.fpn_bottleneck = TFData2VecVisionConvModule(
- in_channels=len(self.in_channels) * self.channels,
- out_channels=self.channels,
- kernel_size=3,
- padding="same",
- name="fpn_bottleneck",
- )
-
- def psp_forward(self, inputs):
- x = inputs[-1]
- psp_outs = [x]
- psp_outs.extend(self.psp_modules(x))
- psp_outs = tf.concat(psp_outs, axis=-1)
- output = self.bottleneck(psp_outs)
-
- return output
-
- def call(self, encoder_hidden_states: tf.Tensor) -> tf.Tensor:
- # build laterals
- laterals = [lateral_conv(encoder_hidden_states[i]) for i, lateral_conv in enumerate(self.lateral_convs)]
-
- laterals.append(self.psp_forward(encoder_hidden_states))
-
- # build top-down path
- used_backbone_levels = len(laterals)
- for i in range(used_backbone_levels - 1, 0, -1):
- prev_shape = shape_list(laterals[i - 1])[1:-1]
- laterals[i - 1] = laterals[i - 1] + tf.image.resize(laterals[i], size=prev_shape, method="bilinear")
-
- # build outputs
- fpn_outs = [self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels - 1)]
- # append psp feature
- fpn_outs.append(laterals[-1])
-
- for i in range(used_backbone_levels - 1, 0, -1):
- fpn_outs[i] = tf.image.resize(fpn_outs[i], size=shape_list(fpn_outs[0])[1:-1], method="bilinear")
- fpn_outs = tf.concat(fpn_outs, axis=-1)
- output = self.fpn_bottleneck(fpn_outs)
- output = self.classifier(output)
-
- return output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, None, self.channels])
- if getattr(self, "psp_modules", None) is not None:
- with tf.name_scope(self.psp_modules.name):
- self.psp_modules.build(None)
- if getattr(self, "bottleneck", None) is not None:
- with tf.name_scope(self.bottleneck.name):
- self.bottleneck.build(None)
- if getattr(self, "fpn_bottleneck", None) is not None:
- with tf.name_scope(self.fpn_bottleneck.name):
- self.fpn_bottleneck.build(None)
- for layer in self.lateral_convs:
- with tf.name_scope(layer.name):
- layer.build(None)
- for layer in self.fpn_convs:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-class TFData2VecVisionFCNHead(keras.layers.Layer):
- """
- Fully Convolution Networks for Semantic Segmentation. This head is implemented from
- [FCNNet](https://arxiv.org/abs/1411.4038).
-
- Args:
- config (Data2VecVisionConfig): Configuration.
- kernel_size (int): The kernel size for convs in the head. Default: 3.
- dilation (int): The dilation rate for convs in the head. Default: 1.
-
-
- Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.
- """
-
- def __init__(
- self,
- config: Data2VecVisionConfig,
- in_index: int = 2,
- kernel_size: int = 3,
- dilation: Union[int, Tuple[int, int]] = 1,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- self.in_channels = config.hidden_size
- self.channels = config.auxiliary_channels
- self.num_convs = config.auxiliary_num_convs
- self.concat_input = config.auxiliary_concat_input
- self.in_index = in_index
-
- convs = []
- convs.append(
- TFData2VecVisionConvModule(
- in_channels=self.in_channels,
- out_channels=self.channels,
- kernel_size=kernel_size,
- padding="same",
- dilation=dilation,
- name="convs.0",
- )
- )
- for i in range(self.num_convs - 1):
- convs.append(
- TFData2VecVisionConvModule(
- in_channels=self.channels,
- out_channels=self.channels,
- kernel_size=kernel_size,
- padding="same",
- dilation=dilation,
- name=f"conv_module_{i+2}",
- )
- )
- if self.num_convs == 0:
- self.convs = [tf.identity]
- else:
- self.convs = convs
- if self.concat_input:
- self.conv_cat = TFData2VecVisionConvModule(
- self.in_channels + self.channels,
- out_channels=self.channels,
- kernel_size=kernel_size,
- padding="same",
- name="conv_cat",
- )
-
- self.classifier = keras.layers.Conv2D(config.num_labels, kernel_size=1, name="classifier")
-
- def call(self, encoder_hidden_states: tf.Tensor) -> tf.Tensor:
- # just take the relevant feature maps
- hidden_states = encoder_hidden_states[self.in_index]
- output = hidden_states
- for layer_module in self.convs:
- output = layer_module(output)
- if self.concat_input:
- output = self.conv_cat(tf.concat([hidden_states, output], axis=-1))
- output = self.classifier(output)
- return output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, None, self.channels])
- if getattr(self, "conv_cat", None) is not None:
- with tf.name_scope(self.conv_cat.name):
- self.conv_cat.build(None)
-
-
-@add_start_docstrings(
- """
- Data2VecVision Model transformer with a semantic segmentation head on top e.g. for ADE20k, CityScapes.
- """,
- DATA2VEC_VISION_START_DOCSTRING,
-)
-class TFData2VecVisionForSemanticSegmentation(TFData2VecVisionPreTrainedModel):
- def __init__(self, config: Data2VecVisionConfig, *inputs, **kwargs) -> None:
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
- self.data2vec_vision = TFData2VecVisionMainLayer(config, add_pooling_layer=False, name="data2vec_vision")
-
- # FPNs
- self.fpn1 = [
- keras.layers.Conv2DTranspose(config.hidden_size, kernel_size=2, strides=2, name="fpn1.0"),
- keras.layers.BatchNormalization(name="fpn1.1", momentum=0.9, epsilon=1e-5),
- keras.layers.Activation("gelu"),
- keras.layers.Conv2DTranspose(config.hidden_size, kernel_size=2, strides=2, name="fpn1.3"),
- ]
- self.fpn2 = [keras.layers.Conv2DTranspose(config.hidden_size, kernel_size=2, strides=2, name="fpn2.0")]
-
- self.fpn3 = tf.identity
- self.fpn4 = keras.layers.MaxPool2D(pool_size=2, strides=2)
-
- # Semantic segmentation head(s)
- self.decode_head = TFData2VecVisionUperHead(config, name="decode_head")
- self.auxiliary_head = (
- TFData2VecVisionFCNHead(config, name="auxiliary_head") if config.use_auxiliary_head else None
- )
-
- def compute_loss(self, logits, auxiliary_logits, labels):
- # upsample logits to the images' original size
- if len(shape_list(labels)) > 3:
- label_interp_shape = shape_list(labels)[1:-1]
- else:
- label_interp_shape = shape_list(labels)[-2:]
-
- upsampled_logits = tf.image.resize(logits, size=label_interp_shape, method="bilinear")
- if auxiliary_logits is not None:
- upsampled_auxiliary_logits = tf.image.resize(auxiliary_logits, size=label_interp_shape, method="bilinear")
- # compute weighted loss
- loss_fct = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction="none")
-
- # Copied from https://www.tensorflow.org/text/tutorials/transformer#loss_and_metrics.
- # Utility to mask the index to ignore during computing the loss.
- def masked_loss(real, pred):
- mask = tf.math.logical_not(tf.math.equal(real, self.config.semantic_loss_ignore_index))
- loss_ = loss_fct(real, pred)
- mask = tf.cast(mask, dtype=loss_.dtype)
- loss_ *= mask
- reduced_masked_loss = tf.reduce_sum(loss_) / tf.reduce_sum(mask)
- return tf.reshape(reduced_masked_loss, (1,))
-
- main_loss = masked_loss(labels, upsampled_logits)
- auxiliary_loss = masked_loss(labels, upsampled_auxiliary_logits)
- loss = main_loss + self.config.auxiliary_loss_weight * auxiliary_loss
-
- return loss
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFSemanticSegmenterOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- labels: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, TFSemanticSegmenterOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, height, width)`, *optional*):
- Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy).
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, TFData2VecVisionForSemanticSegmentation
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/data2vec-vision-base")
- >>> model = TFData2VecVisionForSemanticSegmentation.from_pretrained("facebook/data2vec-vision-base")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
- >>> outputs = model(**inputs)
- >>> # logits are of shape (batch_size, num_labels, height, width)
- >>> logits = outputs.logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- outputs = self.data2vec_vision(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=True, # we need the intermediate hidden states
- return_dict=return_dict,
- )
- encoder_hidden_states = outputs.hidden_states if return_dict else outputs[1]
-
- # only keep certain features, and reshape
- # note that we do +1 as the encoder_hidden_states also includes the initial embeddings
- features = [feature for idx, feature in enumerate(encoder_hidden_states) if idx + 1 in self.config.out_indices]
- patch_resolution = self.config.image_size // self.config.patch_size
-
- def reshape_features(x):
- # We do it this way so TF can always infer the non-batch dims at compile time
- x = tf.reshape(x, (-1, patch_resolution, patch_resolution, self.config.hidden_size))
- return x
-
- features = [reshape_features(x[:, 1:, :]) for x in features]
-
- # apply FPNs
- ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
- for module in ops[0]:
- features[0] = module(features[0])
- features[1] = ops[1][0](features[1])
- for i in range(len(features[2:])):
- features[i + 2] = ops[i + 2](features[i + 2])
-
- logits = self.decode_head(features)
- # Tranpose the logits to maintain consistency in the output formats.
- transposed_logits = tf.transpose(logits, perm=[0, 3, 1, 2])
-
- auxiliary_logits = None
- if self.auxiliary_head is not None:
- auxiliary_logits = self.auxiliary_head(features)
-
- loss = None
- if labels is not None:
- if self.config.num_labels == 1:
- raise ValueError("The number of labels should be greater than one")
- else:
- loss = self.compute_loss(logits, auxiliary_logits, labels)
-
- if not return_dict:
- if output_hidden_states:
- output = (logits,) + outputs[1:]
- else:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSemanticSegmenterOutput(
- loss=loss,
- logits=transposed_logits,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "data2vec_vision", None) is not None:
- with tf.name_scope(self.data2vec_vision.name):
- self.data2vec_vision.build(None)
- if getattr(self, "decode_head", None) is not None:
- with tf.name_scope(self.decode_head.name):
- self.decode_head.build(None)
- if getattr(self, "auxiliary_head", None) is not None:
- with tf.name_scope(self.auxiliary_head.name):
- self.auxiliary_head.build(None)
- if getattr(self, "fpn1", None) is not None:
- with tf.name_scope(self.fpn1[0].name):
- self.fpn1[0].build([None, None, None, self.config.hidden_size])
- with tf.name_scope(self.fpn1[1].name):
- self.fpn1[1].build((None, None, None, self.config.hidden_size))
- with tf.name_scope(self.fpn1[3].name):
- self.fpn1[3].build([None, None, None, self.config.hidden_size])
- if getattr(self, "fpn2", None) is not None:
- with tf.name_scope(self.fpn2[0].name):
- self.fpn2[0].build([None, None, None, self.config.hidden_size])
diff --git a/transformers/models/dbrx/__init__.py b/transformers/models/dbrx/__init__.py
deleted file mode 100644
index 693a544c4b3d3fe238a6ebd106a3235ee32e4fea..0000000000000000000000000000000000000000
--- a/transformers/models/dbrx/__init__.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {
- "configuration_dbrx": ["DbrxConfig"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_dbrx"] = [
- "DbrxForCausalLM",
- "DbrxModel",
- "DbrxPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_dbrx import DbrxConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_dbrx import DbrxForCausalLM, DbrxModel, DbrxPreTrainedModel
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/dbrx/__pycache__/__init__.cpython-310.pyc b/transformers/models/dbrx/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index cb958accc97722c9eb800f8a42865427f4446f56..0000000000000000000000000000000000000000
Binary files a/transformers/models/dbrx/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dbrx/__pycache__/configuration_dbrx.cpython-310.pyc b/transformers/models/dbrx/__pycache__/configuration_dbrx.cpython-310.pyc
deleted file mode 100644
index 2266ceb43b4911835d008aa197bd9d2d6b613bc2..0000000000000000000000000000000000000000
Binary files a/transformers/models/dbrx/__pycache__/configuration_dbrx.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dbrx/__pycache__/modeling_dbrx.cpython-310.pyc b/transformers/models/dbrx/__pycache__/modeling_dbrx.cpython-310.pyc
deleted file mode 100644
index 3d3afdd4eb241768b361d1e3974ca4192630d9e3..0000000000000000000000000000000000000000
Binary files a/transformers/models/dbrx/__pycache__/modeling_dbrx.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dbrx/configuration_dbrx.py b/transformers/models/dbrx/configuration_dbrx.py
deleted file mode 100644
index b03d2c17b09e0787fc09ce3fbe0d1d54b44801a3..0000000000000000000000000000000000000000
--- a/transformers/models/dbrx/configuration_dbrx.py
+++ /dev/null
@@ -1,257 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Databricks Mosaic Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DBRX model configuration """
-
-from typing import Any, Optional
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class DbrxAttentionConfig(PretrainedConfig):
- """Configuration class for Dbrx Attention.
-
- [`DbrxAttention`] class. It is used to instantiate attention layers
- according to the specified arguments, defining the layers architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- attn_pdrop (`float`, *optional*, defaults to 0.0):
- The dropout probability for the attention layers.
- clip_qkv (`float`, *optional*):
- If set, clip the queries, keys, and values in the attention layer to this value.
- kv_n_heads (`Optional[int]`, defaults to 1): For grouped_query_attention only, allow user to specify number of kv heads.
- rope_theta (`float`, defaults to 10000.0): The base frequency for rope.
- """
-
- def __init__(
- self,
- attn_pdrop: float = 0.0,
- clip_qkv: Optional[float] = None,
- kv_n_heads: int = 1,
- rope_theta: float = 10000.0,
- **kwargs: Any,
- ):
- super().__init__(**kwargs)
- self.attn_pdrop = attn_pdrop
- self.clip_qkv = clip_qkv
- self.kv_n_heads = kv_n_heads
- self.rope_theta = rope_theta
-
- for k in ["model_type"]:
- if k in kwargs:
- kwargs.pop(k)
- if len(kwargs) != 0:
- raise ValueError(f"Found unknown {kwargs=}")
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs: Any) -> "PretrainedConfig":
- cls._set_token_in_kwargs(kwargs)
-
- config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
-
- if config_dict.get("model_type") == "dbrx":
- config_dict = config_dict["attn_config"]
-
- if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
- logger.warning(
- f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
- + f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
- )
-
- return cls.from_dict(config_dict, **kwargs)
-
-
-class DbrxFFNConfig(PretrainedConfig):
- """Configuration class for Dbrx FFN.
-
- [`DbrxFFN`] class. It is used to instantiate feedforward layers according to
- the specified arguments, defining the layers architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- ffn_act_fn (`dict`, *optional*, defaults to `None`): A dict specifying activation function for the FFN.
- The dict should have a key 'name' with the value being the name of the activation function along with
- any additional keyword arguments. If `None`, then set to `{"name": "silu"}`.
- ffn_hidden_size (`int`, defaults to 3584): The hidden size of the feedforward network.
- moe_num_experts (`int`, defaults to 4): The number of experts in the mixture of experts layer.
- moe_top_k (`int`, defaults to 1): The number of experts to use in the mixture of experts layer.
- moe_jitter_eps (`float`, *optional*, defaults to `None`): If not `None`, the jitter epsilon for the mixture of experts layer.
- moe_loss_weight (`float`, defaults to 0.01): The loss weight for the mixture of experts layer.
- moe_normalize_expert_weights (`float`, *optional*, defaults to 1.0): The normalization factor for the expert weights.
- """
-
- def __init__(
- self,
- ffn_act_fn: dict = None,
- ffn_hidden_size: int = 3584,
- moe_num_experts: int = 4,
- moe_top_k: int = 1,
- moe_jitter_eps: Optional[float] = None,
- moe_loss_weight: float = 0.01,
- moe_normalize_expert_weights: Optional[float] = 1.0,
- **kwargs: Any,
- ):
- super().__init__()
- if ffn_act_fn is None:
- ffn_act_fn = {"name": "silu"}
- self.ffn_act_fn = ffn_act_fn
- self.ffn_hidden_size = ffn_hidden_size
- self.moe_num_experts = moe_num_experts
- self.moe_top_k = moe_top_k
- self.moe_jitter_eps = moe_jitter_eps
- self.moe_loss_weight = moe_loss_weight
- self.moe_normalize_expert_weights = moe_normalize_expert_weights
-
- for k in ["model_type"]:
- if k in kwargs:
- kwargs.pop(k)
- if len(kwargs) != 0:
- raise ValueError(f"Found unknown {kwargs=}")
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs: Any) -> "PretrainedConfig":
- cls._set_token_in_kwargs(kwargs)
-
- config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
-
- if config_dict.get("model_type") == "dbrx":
- config_dict = config_dict["ffn_config"]
-
- if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
- logger.warning(
- f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
- + f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
- )
-
- return cls.from_dict(config_dict, **kwargs)
-
-
-class DbrxConfig(PretrainedConfig):
- r"""
-
- This is the configuration class to store the configuration of a [`DbrxModel`]. It is used to instantiate a Dbrx model according to the
- specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a different configuration to that of the [databricks/dbrx-instruct](https://huggingface.co/databricks/dbrx-instruct) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- d_model (`int`, *optional*, defaults to 2048):
- Dimensionality of the embeddings and hidden states.
- n_heads (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- n_layers (`int`, *optional*, defaults to 24):
- Number of hidden layers in the Transformer encoder.
- max_seq_len (`int`, *optional*, defaults to 2048):
- The maximum sequence length of the model.
- vocab_size (`int`, *optional*, defaults to 32000):
- Vocabulary size of the Dbrx model. Defines the maximum number of different tokens that can be represented by
- the `inputs_ids` passed when calling [`DbrxModel`].
- resid_pdrop (`float`, *optional*, defaults to 0.0):
- The dropout probability applied to the attention output before combining with residual.
- emb_pdrop (`float`, *optional*, defaults to 0.0):
- The dropout probability for the embedding layer.
- attn_config (`dict`, *optional*):
- A dictionary used to configure the model's attention module.
- ffn_config (`dict`, *optional*):
- A dictionary used to configure the model's FFN module.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- output_router_logits (`bool`, *optional*, defaults to `False`):
- Whether or not the router logits should be returned by the model. Enabling this will also
- allow the model to output the auxiliary loss. See [here]() for more details.
-
-
- Example:
- ```python
- >>> from transformers import DbrxConfig, DbrxModel
-
- >>> # Initializing a Dbrx configuration
- >>> configuration = DbrxConfig(n_layers=2, d_model=256, n_heads=8, vocab_size=128)
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = DbrxModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```
- """
-
- model_type = "dbrx"
- attribute_map = {
- "num_attention_heads": "n_heads",
- "hidden_size": "d_model",
- "num_hidden_layers": "n_layers",
- "max_position_embeddings": "max_seq_len",
- }
-
- def __init__(
- self,
- d_model: int = 2048,
- n_heads: int = 16,
- n_layers: int = 24,
- max_seq_len: int = 2048,
- vocab_size: int = 32000,
- resid_pdrop: float = 0.0,
- emb_pdrop: float = 0.0,
- attn_config: Optional[DbrxAttentionConfig] = None,
- ffn_config: Optional[DbrxFFNConfig] = None,
- use_cache: bool = True,
- initializer_range: float = 0.02,
- output_router_logits: bool = False,
- **kwargs: Any,
- ):
- if attn_config is None:
- self.attn_config = DbrxAttentionConfig()
- elif isinstance(attn_config, dict):
- self.attn_config = DbrxAttentionConfig(**attn_config)
- else:
- self.attn_config = attn_config
-
- if ffn_config is None:
- self.ffn_config = DbrxFFNConfig()
- elif isinstance(ffn_config, dict):
- self.ffn_config = DbrxFFNConfig(**ffn_config)
- else:
- self.ffn_config = ffn_config
-
- self.d_model = d_model
- self.n_heads = n_heads
- self.n_layers = n_layers
- self.max_seq_len = max_seq_len
- self.vocab_size = vocab_size
- self.resid_pdrop = resid_pdrop
- self.emb_pdrop = emb_pdrop
- self.use_cache = use_cache
- self.initializer_range = initializer_range
- self.output_router_logits = output_router_logits
-
- tie_word_embeddings = kwargs.pop("tie_word_embeddings", False)
- if tie_word_embeddings:
- raise ValueError("tie_word_embeddings is not supported for DBRX models.")
-
- super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
diff --git a/transformers/models/dbrx/modeling_dbrx.py b/transformers/models/dbrx/modeling_dbrx.py
deleted file mode 100644
index 99b865c773f81da489d546863832febae0388ab6..0000000000000000000000000000000000000000
--- a/transformers/models/dbrx/modeling_dbrx.py
+++ /dev/null
@@ -1,1523 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Databricks Mosaic Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DBRX model. """
-
-import math
-from typing import Any, Dict, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch import nn
-
-from ...activations import ACT2FN
-from ...cache_utils import Cache, DynamicCache, StaticCache
-from ...modeling_attn_mask_utils import AttentionMaskConverter
-from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- logging,
- replace_return_docstrings,
-)
-from .configuration_dbrx import DbrxConfig
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "DbrxConfig"
-
-
-# Copied from transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding with Gemma->Dbrx
-class DbrxRotaryEmbedding(nn.Module):
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
- super().__init__()
-
- self.dim = dim
- self.max_position_embeddings = max_position_embeddings
- self.base = base
- self.register_buffer("inv_freq", None, persistent=False)
-
- @torch.no_grad()
- def forward(self, x, position_ids, seq_len=None):
- # x: [bs, num_attention_heads, seq_len, head_size]
- if self.inv_freq is None:
- self.inv_freq = 1.0 / (
- self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim)
- )
- inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
- position_ids_expanded = position_ids[:, None, :].float()
- # Force float32 since bfloat16 loses precision on long contexts
- # See https://github.com/huggingface/transformers/pull/29285
- device_type = x.device.type
- device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
- with torch.autocast(device_type=device_type, enabled=False):
- freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
- emb = torch.cat((freqs, freqs), dim=-1)
- cos = emb.cos()
- sin = emb.sin()
- return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-# Copied from transformers.models.llama.modeling_llama.rotate_half
-def rotate_half(x):
- """Rotates half the hidden dims of the input."""
- x1 = x[..., : x.shape[-1] // 2]
- x2 = x[..., x.shape[-1] // 2 :]
- return torch.cat((-x2, x1), dim=-1)
-
-
-# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
- """Applies Rotary Position Embedding to the query and key tensors.
-
- Args:
- q (`torch.Tensor`): The query tensor.
- k (`torch.Tensor`): The key tensor.
- cos (`torch.Tensor`): The cosine part of the rotary embedding.
- sin (`torch.Tensor`): The sine part of the rotary embedding.
- position_ids (`torch.Tensor`, *optional*):
- Deprecated and unused.
- unsqueeze_dim (`int`, *optional*, defaults to 1):
- The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
- sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
- that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
- k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
- cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
- the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
- Returns:
- `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
- """
- cos = cos.unsqueeze(unsqueeze_dim)
- sin = sin.unsqueeze(unsqueeze_dim)
- q_embed = (q * cos) + (rotate_half(q) * sin)
- k_embed = (k * cos) + (rotate_half(k) * sin)
- return q_embed, k_embed
-
-
-# Copied from transformers.models.llama.modeling_llama.repeat_kv
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
- """
- This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
- num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
- """
- batch, num_key_value_heads, slen, head_dim = hidden_states.shape
- if n_rep == 1:
- return hidden_states
- hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
- return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-def load_balancing_loss_func(
- gate_logits: torch.Tensor,
- num_experts: int,
- top_k: int,
- attention_mask: Optional[torch.Tensor],
-) -> torch.Tensor:
- r"""Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
-
- See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
- function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
- experts is too unbalanced.
-
- Args:
- gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
- Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
- shape [batch_size X sequence_length, num_experts].
- num_experts (`int`):
- Number of experts.
- top_k (`int`):
- The number of experts each token is routed to.
- attention_mask (`torch.Tensor`, None):
- The attention_mask used in forward function
- shape [batch_size X sequence_length] if not None.
-
- Returns:
- The auxiliary loss.
- """
- if gate_logits is None or not isinstance(gate_logits, tuple):
- return torch.tensor(0.0)
-
- if isinstance(gate_logits, tuple):
- compute_device = gate_logits[0].device
- concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
-
- routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
-
- _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
-
- expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
-
- if attention_mask is None:
- # Compute the percentage of tokens routed to each experts
- tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
-
- # Compute the average probability of routing to these experts
- router_prob_per_expert = torch.mean(routing_weights, dim=0)
- else:
- batch_size, sequence_length = attention_mask.shape
- num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
-
- # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
- expert_attention_mask = (
- attention_mask[None, :, :, None, None]
- .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
- .reshape(-1, top_k, num_experts)
- .to(compute_device)
- )
-
- # Compute the percentage of tokens routed to each experts
- tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
- expert_attention_mask, dim=0
- )
-
- # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
- router_per_expert_attention_mask = (
- attention_mask[None, :, :, None]
- .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
- .reshape(-1, num_experts)
- .to(compute_device)
- )
-
- # Compute the average probability of routing to these experts
- router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
- router_per_expert_attention_mask, dim=0
- )
-
- overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
- return overall_loss * num_experts
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-class DbrxAttention(nn.Module):
- """Multi-head self attention."""
-
- def __init__(self, config: DbrxConfig, block_idx: Optional[int] = None):
- super().__init__()
- self.config = config
- self.hidden_size = config.d_model
- self.num_heads = config.n_heads
- self.head_dim = self.hidden_size // self.num_heads
- self.max_position_embeddings = config.max_seq_len
- self.block_idx = block_idx
- if block_idx is None:
- logger.warning_once(
- f"Instantiating {self.__class__.__name__} without passing a `block_idx` is not recommended and will "
- + "lead to errors during the forward call if caching is used. Please make sure to provide a `block_idx` "
- + "when creating this class."
- )
-
- attn_config = config.attn_config
- self.attn_pdrop = attn_config.attn_pdrop
- self.clip_qkv = attn_config.clip_qkv
- self.num_key_value_heads = attn_config.kv_n_heads
- self.num_key_value_groups = self.num_heads // self.num_key_value_heads
- self.rope_theta = attn_config.rope_theta
- self.is_causal = True
-
- self.Wqkv = nn.Linear(
- self.hidden_size, self.hidden_size + 2 * self.num_key_value_heads * self.head_dim, bias=False
- )
- self.out_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
- self.rotary_emb = DbrxRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- base=self.rope_theta,
- )
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- position_ids: torch.LongTensor,
- attention_mask: Optional[torch.Tensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs: Any,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
- bsz, q_len, _ = hidden_states.size()
-
- qkv_states = self.Wqkv(hidden_states)
- min_val = -self.clip_qkv if self.clip_qkv is not None else None
- max_val = self.clip_qkv
- qkv_states = qkv_states.clamp(min=min_val, max=max_val)
-
- query_states, key_states, value_states = qkv_states.split(
- [
- self.hidden_size,
- self.num_key_value_heads * self.head_dim,
- self.num_key_value_heads * self.head_dim,
- ],
- dim=2,
- )
-
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- past_key_value = getattr(self, "past_key_value", past_key_value)
- cos, sin = self.rotary_emb(value_states, position_ids)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; position_ids needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.block_idx, cache_kwargs)
-
- key_states = repeat_kv(key_states, self.num_key_value_groups)
- value_states = repeat_kv(value_states, self.num_key_value_groups)
-
- attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-
- if attention_mask is not None: # no matter the length, we just slice it
- causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
- attn_weights = attn_weights + causal_mask
-
- # upcast attention to fp32
- attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
- attn_weights = nn.functional.dropout(attn_weights, p=self.attn_pdrop, training=self.training)
- attn_output = torch.matmul(attn_weights, value_states)
-
- if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
- + f" {attn_output.size()}"
- )
-
- attn_output = attn_output.transpose(1, 2).contiguous()
- attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
- attn_output = self.out_proj(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights, past_key_value
-
-
-class DbrxFlashAttention2(DbrxAttention):
- """Dbrx flash attention module.
-
- This module inherits from `DbrxAttention` as the weights of the module stays
- untouched. The only required change would be on the forward pass where it
- calls the public API of flash attention.
- """
-
- def __init__(self, *args: Any, **kwargs: Any):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- # From: https://github.com/huggingface/transformers/blob/3b8e2932ce743008f63585aae1e1b8b30dc8b3ac/src/transformers/models/gemma/modeling_gemma.py#L318
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs: Any,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- logger.info("Implicitly setting `output_attentions` to False as it is not supported in Flash Attention.")
- output_attentions = False
-
- bsz, q_len, _ = hidden_states.size()
-
- qkv_states = self.Wqkv(hidden_states)
- if self.clip_qkv is not None:
- qkv_states = qkv_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
-
- query_states, key_states, value_states = qkv_states.split(
- [
- self.hidden_size,
- self.num_key_value_heads * self.head_dim,
- self.num_key_value_heads * self.head_dim,
- ],
- dim=2,
- )
-
- # Flash attention requires the input to have the shape
- # batch_size x seq_length x head_dim x hidden_dim
- # therefore we just need to keep the original shape
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- cos, sin = self.rotary_emb(value_states, position_ids)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
- past_key_value = getattr(self, "past_key_value", past_key_value)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; cache_position needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.block_idx, cache_kwargs)
-
- # TODO: These transpose are quite inefficient but Flash Attention requires the layout
- # [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
- # to be able to avoid many of these transpose/reshape/view.
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.transpose(1, 2)
-
- dropout_rate = self.attn_pdrop if self.training else 0.0
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in the correct dtype just to be sure everything works as expected.
- # This might slowdown training & inference so it is recommended to not cast the LayerNorms
- # in fp32. (LlamaRMSNorm handles it correctly)
- input_dtype = query_states.dtype
- if input_dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = query_states.dtype
-
- logger.warning_once(
- "The input hidden states seems to be silently casted in float32, this might be "
- + "related to the fact you have upcasted embedding or layer norm layers in "
- + f"float32. We will cast back the input in {target_dtype}."
- )
-
- query_states = query_states.to(target_dtype)
- key_states = key_states.to(target_dtype)
- value_states = value_states.to(target_dtype)
-
- attn_output = self._flash_attention_forward(
- query_states,
- key_states,
- value_states,
- attention_mask,
- q_len,
- dropout=dropout_rate,
- )
-
- attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
- attn_output = self.out_proj(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights, past_key_value
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-class DbrxSdpaAttention(DbrxAttention):
- """
- Dbrx attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
- `DbrxAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
- SDPA API.
- """
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- if output_attentions:
- # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
- logger.warning_once(
- "DbrxModel is using DbrxSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
- 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
- )
- return super().forward(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- )
-
- bsz, q_len, _ = hidden_states.size()
-
- qkv_states = self.Wqkv(hidden_states)
- if self.clip_qkv is not None:
- qkv_states = qkv_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
-
- query_states, key_states, value_states = qkv_states.split(
- [
- self.hidden_size,
- self.num_key_value_heads * self.head_dim,
- self.num_key_value_heads * self.head_dim,
- ],
- dim=2,
- )
-
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- cos, sin = self.rotary_emb(value_states, position_ids, seq_len=None)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, None)
-
- past_key_value = getattr(self, "past_key_value", past_key_value)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; cache_position needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.block_idx, cache_kwargs)
-
- key_states = repeat_kv(key_states, self.num_key_value_groups)
- value_states = repeat_kv(value_states, self.num_key_value_groups)
-
- causal_mask = attention_mask
- if attention_mask is not None:
- causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
-
- # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
- # Reference: https://github.com/pytorch/pytorch/issues/112577.
- if query_states.device.type == "cuda" and causal_mask is not None:
- query_states = query_states.contiguous()
- key_states = key_states.contiguous()
- value_states = value_states.contiguous()
-
- attn_output = torch.nn.functional.scaled_dot_product_attention(
- query_states,
- key_states,
- value_states,
- attn_mask=causal_mask,
- dropout_p=self.attn_pdrop if self.training else 0.0,
- )
-
- attn_output = attn_output.transpose(1, 2).contiguous()
- attn_output = attn_output.view(bsz, q_len, -1)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, None, past_key_value
-
-
-DBRX_ATTENTION_CLASSES = {
- "eager": DbrxAttention,
- "flash_attention_2": DbrxFlashAttention2,
- "sdpa": DbrxSdpaAttention,
-}
-
-
-class DbrxNormAttentionNorm(nn.Module):
- def __init__(self, config: DbrxConfig, block_idx: Optional[int] = None):
- super().__init__()
- self.block_idx = block_idx
- self.resid_pdrop = config.resid_pdrop
- self.norm_1 = nn.LayerNorm(config.d_model, bias=False)
- self.attn = DBRX_ATTENTION_CLASSES[config._attn_implementation](
- config=config,
- block_idx=block_idx,
- )
- self.norm_2 = nn.LayerNorm(config.d_model, bias=False)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- position_ids: torch.LongTensor,
- attention_mask: Optional[torch.Tensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs: Any,
- ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
- residual_states = hidden_states
- hidden_states = self.norm_1(hidden_states).to(hidden_states.dtype)
-
- hidden_states, attn_weights, past_key_value = self.attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- **kwargs,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.resid_pdrop, training=self.training)
- hidden_states = hidden_states + residual_states
-
- residual_states = hidden_states
- hidden_states = self.norm_2(hidden_states).to(hidden_states.dtype)
-
- return residual_states, hidden_states, attn_weights, past_key_value
-
-
-class DbrxRouter(nn.Module):
- def __init__(
- self,
- hidden_size: int,
- moe_num_experts: int,
- moe_top_k: int,
- moe_jitter_eps: Optional[float],
- moe_normalize_expert_weights: Optional[float],
- ):
- super().__init__()
- self.hidden_size = hidden_size
- self.moe_num_experts = moe_num_experts
- self.moe_top_k = moe_top_k
- self.moe_jitter_eps = moe_jitter_eps
- self.moe_normalize_expert_weights = moe_normalize_expert_weights
-
- self.layer = nn.Linear(self.hidden_size, self.moe_num_experts, bias=False)
-
- def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
- if self.training and self.moe_jitter_eps is not None:
- hidden_states *= torch.empty_like(hidden_states).uniform_(
- 1.0 - self.moe_jitter_eps, 1.0 + self.moe_jitter_eps
- )
- hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
- weights = self.layer(hidden_states).softmax(dim=-1, dtype=torch.float32)
- top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)
-
- top_weights_scale = (
- torch.norm(top_weights, p=self.moe_normalize_expert_weights, dim=-1, keepdim=True)
- if self.moe_normalize_expert_weights is not None
- else 1.0
- )
- top_weights = top_weights / top_weights_scale
-
- weights = weights.to(hidden_states.dtype)
- top_weights = top_weights.to(hidden_states.dtype)
- return weights, top_weights, top_experts
-
-
-class DbrxExpertGLU(nn.Module):
- def __init__(self, hidden_size: int, ffn_hidden_size: int, moe_num_experts: int, ffn_act_fn: dict):
- super().__init__()
- self.hidden_size = hidden_size
- self.ffn_hidden_size = ffn_hidden_size
- self.moe_num_experts = moe_num_experts
-
- self.w1 = nn.Parameter(torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))
- self.v1 = nn.Parameter(torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))
- self.w2 = nn.Parameter(torch.empty(moe_num_experts * ffn_hidden_size, hidden_size))
-
- act_fn_name = ffn_act_fn.get("name", "silu")
- self.activation_fn = ACT2FN[act_fn_name]
-
- def forward(
- self, x: torch.Tensor, expert_w1: torch.Tensor, expert_v1: torch.Tensor, expert_w2: torch.Tensor
- ) -> torch.Tensor:
- gate_proj = x.matmul(expert_w1.t())
- up_proj = x.matmul(expert_v1.t())
- gate_proj = self.activation_fn(gate_proj)
- intermediate_states = gate_proj * up_proj
- down_proj = intermediate_states.matmul(expert_w2)
- return down_proj
-
-
-class DbrxExperts(nn.Module):
- def __init__(self, hidden_size: int, ffn_hidden_size: int, moe_num_experts: int, ffn_act_fn: dict):
- super().__init__()
- self.moe_num_experts = moe_num_experts
- self.mlp = DbrxExpertGLU(
- hidden_size=hidden_size,
- ffn_hidden_size=ffn_hidden_size,
- moe_num_experts=moe_num_experts,
- ffn_act_fn=ffn_act_fn,
- )
-
- def forward(
- self, x: torch.Tensor, weights: torch.Tensor, top_weights: torch.Tensor, top_experts: torch.LongTensor
- ) -> torch.Tensor:
- bsz, q_len, hidden_size = x.shape
- x = x.view(-1, hidden_size)
- out = torch.zeros_like(x)
-
- expert_mask = nn.functional.one_hot(top_experts, num_classes=self.moe_num_experts).permute(2, 1, 0)
- # Chunk experts at once to avoid storing full parameter multiple times in autograd
- w1_chunked = self.mlp.w1.view(self.mlp.moe_num_experts, self.mlp.ffn_hidden_size, self.mlp.hidden_size).chunk(
- self.moe_num_experts, dim=0
- )
- v1_chunked = self.mlp.v1.view(self.mlp.moe_num_experts, self.mlp.ffn_hidden_size, self.mlp.hidden_size).chunk(
- self.moe_num_experts, dim=0
- )
- w2_chunked = self.mlp.w2.view(self.mlp.moe_num_experts, self.mlp.ffn_hidden_size, self.mlp.hidden_size).chunk(
- self.moe_num_experts, dim=0
- )
- w1_chunked = [w1.squeeze(dim=0) for w1 in w1_chunked]
- v1_chunked = [v1.squeeze(dim=0) for v1 in v1_chunked]
- w2_chunked = [w2.squeeze(dim=0) for w2 in w2_chunked]
- for expert_idx in range(0, self.moe_num_experts):
- topk_idx, token_idx = torch.where(expert_mask[expert_idx])
- if token_idx.shape[0] == 0:
- continue
-
- token_list = token_idx
- topk_list = topk_idx
-
- expert_tokens = x[None, token_list].reshape(-1, hidden_size)
- expert_out = (
- self.mlp(expert_tokens, w1_chunked[expert_idx], v1_chunked[expert_idx], w2_chunked[expert_idx])
- * top_weights[token_list, topk_list, None]
- )
-
- out.index_add_(0, token_idx, expert_out)
-
- out = out.reshape(bsz, q_len, hidden_size)
- return out
-
-
-class DbrxFFN(nn.Module):
- def __init__(self, config: DbrxConfig):
- super().__init__()
-
- ffn_config = config.ffn_config
- self.router = DbrxRouter(
- hidden_size=config.d_model,
- moe_num_experts=ffn_config.moe_num_experts,
- moe_top_k=ffn_config.moe_top_k,
- moe_jitter_eps=ffn_config.moe_jitter_eps,
- moe_normalize_expert_weights=ffn_config.moe_normalize_expert_weights,
- )
-
- self.experts = DbrxExperts(
- hidden_size=config.d_model,
- ffn_hidden_size=ffn_config.ffn_hidden_size,
- moe_num_experts=ffn_config.moe_num_experts,
- ffn_act_fn=ffn_config.ffn_act_fn,
- )
-
- def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
- weights, top_weights, top_experts = self.router(x)
- out = self.experts(x, weights, top_weights, top_experts)
- return out, weights
-
-
-class DbrxBlock(nn.Module):
- def __init__(self, config: DbrxConfig, block_idx: int):
- super().__init__()
- self.hidden_size = config.d_model
- self.resid_pdrop = config.resid_pdrop
- self.block_idx = block_idx
- self.norm_attn_norm = DbrxNormAttentionNorm(
- config=config,
- block_idx=block_idx,
- )
- self.ffn = DbrxFFN(config=config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: torch.LongTensor = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: Optional[bool] = False,
- output_router_logits: Optional[bool] = False,
- use_cache: Optional[bool] = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs: Any,
- ) -> Union[
- Tuple[torch.Tensor],
- Tuple[torch.Tensor, Optional[torch.Tensor]],
- Tuple[torch.Tensor, Optional[Cache]],
- Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]],
- Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]],
- Tuple[torch.Tensor, Optional[Cache], Optional[torch.Tensor]],
- Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache], Optional[torch.Tensor]],
- ]:
- """Forward function for DbrxBlock.
-
- Args:
- hidden_states (`torch.Tensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- position_ids (`torch.LongTensor`): position ids of shape `(batch, seq_len)`
- attention_mask (`torch.Tensor`, optional): attention mask of size (batch_size, sequence_length)
- if flash attention is used or (batch_size, 1, query_sequence_length, key_sequence_length)
- if default attention is used.
- past_key_value (`Tuple(torch.Tensor)`, optional): cached past key and value projection states
- output_attentions (`bool`, optional): Whether or not to return the attentions tensors of all
- attention layers. See `attentions` under returned tensors for more detail.
- output_router_logits (`bool`, optional): Whether or not to return the router logits.
- use_cache (`bool`, optional): If set to `True`, `past_key_values` key value states are
- returned and can be used to speed up decoding (see `past_key_values`).
- cache_position (`torch.LongTensor`, optional): position ids of the cache
- """
-
- # Norm + Attention + Norm
- resid_states, hidden_states, self_attn_weights, present_key_value = self.norm_attn_norm(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- **kwargs,
- )
-
- # Fully Connected
- hidden_states, router_logits = self.ffn(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.resid_pdrop, training=self.training)
- hidden_states = resid_states + hidden_states
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights,)
-
- if use_cache:
- outputs += (present_key_value,)
-
- if output_router_logits:
- outputs += (router_logits,)
-
- return outputs
-
-
-DBRX_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`DbrxConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-@add_start_docstrings(
- "The bare DBRX Model outputting raw hidden-states without any specific head on top.",
- DBRX_START_DOCSTRING,
-)
-class DbrxPreTrainedModel(PreTrainedModel):
- config_class = DbrxConfig
- base_model_prefix = "transformer"
- supports_gradient_checkpointing = True
- _no_split_modules = ["DbrxBlock"]
- _skip_keys_device_placement = ["past_key_values"]
- _supports_flash_attn_2 = True
- _supports_sdpa = True
- _supports_cache_class = True
-
- def _init_weights(self, module: nn.Module):
- std = self.config.initializer_range
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, DbrxExpertGLU):
- module.w1.data.normal_(mean=0.0, std=std)
- module.v1.data.normal_(mean=0.0, std=std)
- module.w2.data.normal_(mean=0.0, std=std)
-
- def _setup_cache(self, cache_cls: Any, max_batch_size: int, max_cache_len: int):
- if self.config._attn_implementation == "flash_attention_2" and cache_cls == StaticCache:
- raise ValueError(
- "`static` cache implementation is not compatible with "
- + "`attn_implementation==flash_attention_2`. Make sure to use "
- + "`spda` in the mean time and open an issue at https://github.com/huggingface/transformers."
- )
-
- for block in self.transformer.blocks:
- device = block.norm_attn_norm.norm_1.weight.device
- if hasattr(self.config, "_pre_quantization_dtype"):
- dtype = self.config._pre_quantization_dtype
- else:
- dtype = block.norm_attn_norm.attn.out_proj.weight.dtype
- block.norm_attn_norm.attn.past_key_value = cache_cls(
- self.config, max_batch_size, max_cache_len, device=device, dtype=dtype
- )
-
- def _reset_cache(self):
- for block in self.transformer.blocks:
- block.norm_attn_norm.attn.past_key_value = None
-
-
-DBRX_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
- it.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
- `past_key_values`).
-
- If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
- and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
- information on the default strategy.
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
- Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
- blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
- returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
-
- Two formats are allowed:
- - a [`~cache_utils.Cache`] instance;
- - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
- shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
- cache format.
-
- The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
- legacy cache format will be returned.
-
- If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
- have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
- of shape `(batch_size, sequence_length)`.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- output_router_logits (`bool`, *optional*):
- Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
- should not be returned during inference.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
- Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
- this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
- the complete sequence length.
-"""
-
-
-@add_start_docstrings(
- "The bare DBRX Model outputting raw hidden-states without any specific head on top.",
- DBRX_START_DOCSTRING,
-)
-class DbrxModel(DbrxPreTrainedModel):
- """Transformer decoder consisting of *config.num_hidden_layers*. Each layer is a [`DbrxBlock`] layer.
-
- Args:
- config ([`DbrxConfig`]): Model configuration class with all parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
- """
-
- def __init__(self, config: DbrxConfig):
- super().__init__(config)
- self.padding_idx = config.pad_token_id
- self.vocab_size = config.vocab_size
- self.emb_pdrop = config.emb_pdrop
-
- self.wte = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx)
- self.blocks = nn.ModuleList([DbrxBlock(config, block_idx) for block_idx in range(config.n_layers)])
- self.norm_f = nn.LayerNorm(config.d_model, bias=False)
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self) -> nn.Embedding:
- return self.wte
-
- def set_input_embeddings(self, value: nn.Embedding):
- self.wte = value
-
- @add_start_docstrings_to_model_forward(DBRX_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Cache] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- output_router_logits: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Union[Tuple, MoeModelOutputWithPast]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- output_router_logits = (
- output_router_logits if output_router_logits is not None else self.config.output_router_logits
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if (input_ids is None) ^ (inputs_embeds is not None):
- raise ValueError(
- "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
- )
-
- if self.gradient_checkpointing and self.training and use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
- )
- use_cache = False
-
- if inputs_embeds is None:
- inputs_embeds = self.wte(input_ids)
-
- inputs_embeds = nn.functional.dropout(inputs_embeds, p=self.emb_pdrop, training=self.training)
-
- past_seen_tokens = 0
- if use_cache: # kept for BC (cache positions)
- if not isinstance(past_key_values, StaticCache):
- past_key_values = DynamicCache.from_legacy_cache(past_key_values)
- past_seen_tokens = past_key_values.get_seq_length()
-
- if cache_position is None:
- if isinstance(past_key_values, StaticCache):
- raise ValueError("cache_position is a required argument when using StaticCache.")
- cache_position = torch.arange(
- past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
- )
-
- if position_ids is None:
- position_ids = cache_position.unsqueeze(0)
- causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)
-
- # embed positions
- hidden_states = inputs_embeds
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- all_router_logits = () if output_router_logits else None
- next_decoder_cache = None
-
- for block in self.blocks:
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- block_outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- causal_mask,
- position_ids,
- past_key_values,
- output_attentions,
- output_router_logits,
- use_cache,
- cache_position,
- )
- else:
- block_outputs = block(
- hidden_states,
- attention_mask=causal_mask,
- position_ids=position_ids,
- past_key_value=past_key_values,
- output_attentions=output_attentions,
- output_router_logits=output_router_logits,
- use_cache=use_cache,
- cache_position=cache_position,
- )
-
- hidden_states = block_outputs[0]
-
- if use_cache:
- next_decoder_cache = block_outputs[2 if output_attentions else 1]
-
- if output_attentions:
- all_self_attns += (block_outputs[1],)
-
- if output_router_logits:
- all_router_logits += (block_outputs[-1],)
-
- hidden_states = self.norm_f(hidden_states)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- next_cache = None
- if use_cache:
- next_cache = (
- next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache
- )
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
- if v is not None
- )
- return MoeModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=next_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- router_logits=all_router_logits,
- )
-
- # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static
- # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.
- # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using
- # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114
- def _update_causal_mask(
- self, attention_mask: Optional[torch.Tensor], input_tensor: torch.Tensor, cache_position: torch.Tensor
- ) -> Optional[torch.Tensor]:
- if self.config._attn_implementation == "flash_attention_2":
- if attention_mask is not None and 0.0 in attention_mask:
- return attention_mask
- return None
-
- dtype, device = input_tensor.dtype, input_tensor.device
- min_dtype = torch.finfo(dtype).min
- sequence_length = input_tensor.shape[1]
- if hasattr(self.blocks[0].norm_attn_norm.attn, "past_key_value"): # static cache
- target_length = self.config.max_position_embeddings
- else: # dynamic cache
- target_length = (
- attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else cache_position[-1] + 1
- )
- target_length = int(target_length)
-
- causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
- if sequence_length != 1:
- causal_mask = torch.triu(causal_mask, diagonal=1)
- causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
- causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
- if attention_mask is not None:
- causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
- if attention_mask.dim() == 2:
- mask_length = attention_mask.shape[-1]
- padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
- causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype)
- elif attention_mask.dim() == 4:
- # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with
- # cache. In that case, the 4D attention mask attends to the newest tokens only.
- if attention_mask.shape[-2] < cache_position[0] + sequence_length:
- offset = cache_position[0]
- else:
- offset = 0
- mask_shape = attention_mask.shape
- mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype
- causal_mask[
- : mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]
- ] = mask_slice
-
- if (
- self.config._attn_implementation == "sdpa"
- and attention_mask is not None
- and attention_mask.device.type == "cuda"
- ):
- # TODO: For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400).
- is_tracing = (
- torch.jit.is_tracing()
- or isinstance(input_tensor, torch.fx.Proxy)
- or (hasattr(torch, "_dynamo") and torch._dynamo.is_compiling())
- )
- if not is_tracing and torch.any(attention_mask != 1):
- # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
- # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
- # Details: https://github.com/pytorch/pytorch/issues/110213
- causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
-
- return causal_mask
-
-
-@add_start_docstrings("The DBRX Model transformer for causal language modeling.", DBRX_START_DOCSTRING)
-class DbrxForCausalLM(DbrxPreTrainedModel):
- def __init__(self, config: DbrxConfig):
- super().__init__(config)
- self.transformer = DbrxModel(config)
- self.vocab_size = config.vocab_size
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
- self.moe_loss_weight = config.ffn_config.moe_loss_weight
- self.num_experts = config.ffn_config.moe_num_experts
- self.num_experts_per_tok = config.ffn_config.moe_top_k
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self) -> nn.Embedding:
- return self.transformer.get_input_embeddings()
-
- def set_input_embeddings(self, value: nn.Embedding):
- self.transformer.set_input_embeddings(value)
-
- def get_output_embeddings(self) -> nn.Linear:
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings: nn.Linear):
- self.lm_head = new_embeddings
-
- def set_decoder(self, decoder: DbrxModel):
- self.transformer = decoder
-
- def get_decoder(self) -> DbrxModel:
- return self.transformer
-
- @add_start_docstrings_to_model_forward(DBRX_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Cache] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- output_router_logits: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
- r"""Forward function for causal language modeling.
-
- Args:
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- Returns:
-
- Example:
-
- ```python
- >> from transformers import AutoTokenizer, DbrxForCausalLM
-
- >> model = DbrxForCausalLM.from_pretrained("databricks/dbrx-instruct")
- >> tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-instruct")
-
- >> prompt = "Hey, are you conscious? Can you talk to me?"
- >> inputs = tokenizer(prompt, return_tensors="pt")
-
- >> # Generate
- >> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
- ```
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- output_router_logits = (
- output_router_logits if output_router_logits is not None else self.config.output_router_logits
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.transformer(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- output_router_logits=output_router_logits,
- return_dict=return_dict,
- cache_position=cache_position,
- )
-
- hidden_states = outputs[0]
- logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = nn.CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- loss = loss_fct(shift_logits, shift_labels)
-
- aux_loss = None
- if output_router_logits:
- aux_loss = load_balancing_loss_func(
- outputs.router_logits if return_dict else outputs[-1],
- self.num_experts,
- self.num_experts_per_tok,
- attention_mask,
- )
- if labels is not None and loss is not None:
- loss += self.moe_loss_weight * aux_loss.to(loss.device) # make sure to reside in the same device
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- if output_router_logits:
- output = (aux_loss,) + output
- return (loss,) + output if loss is not None else output
-
- return MoeCausalLMOutputWithPast(
- loss=loss,
- aux_loss=aux_loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- router_logits=outputs.router_logits,
- )
-
- def prepare_inputs_for_generation(
- self,
- input_ids: torch.Tensor,
- past_key_values: Optional[Cache] = None,
- attention_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- **kwargs: Any,
- ) -> Dict[str, Any]:
- past_length = 0
- if past_key_values is not None:
- if isinstance(past_key_values, Cache):
- cache_length = past_key_values.get_seq_length()
- past_length = past_key_values.seen_tokens
- max_cache_length = past_key_values.get_max_length()
- else:
- cache_length = past_length = past_key_values[0][0].shape[2]
- max_cache_length = None
-
- # Keep only the unprocessed tokens:
- # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
- # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
- # input)
- if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
- input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
- # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
- # input_ids based on the past_length.
- elif past_length < input_ids.shape[1]:
- input_ids = input_ids[:, past_length:]
- # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
-
- # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
- if (
- max_cache_length is not None
- and attention_mask is not None
- and cache_length + input_ids.shape[1] > max_cache_length
- ):
- attention_mask = attention_mask[:, -max_cache_length:]
-
- position_ids = kwargs.get("position_ids", None)
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- if self.generation_config.cache_implementation == "static":
- # generation with static cache
- cache_position = kwargs.get("cache_position", None)
- if cache_position is None:
- past_length = 0
- else:
- past_length = cache_position[-1] + 1
- input_ids = input_ids[:, past_length:]
- position_ids = position_ids[:, past_length:] if position_ids is not None else None
-
- # TODO @gante we should only keep a `cache_position` in generate, and do +=1.
- # same goes for position ids. Could also help with continued generation.
- input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
- cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
- position_ids = position_ids.contiguous() if position_ids is not None else None
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
- # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
- # TODO: use `next_tokens` directly instead.
- model_inputs = {"input_ids": input_ids.contiguous()}
-
- model_inputs.update(
- {
- "position_ids": position_ids,
- "cache_position": cache_position,
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "attention_mask": attention_mask,
- }
- )
- return model_inputs
-
- @staticmethod
- def _reorder_cache(past_key_values: Cache, beam_idx: torch.LongTensor):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
- )
- return reordered_past
diff --git a/transformers/models/deberta/__init__.py b/transformers/models/deberta/__init__.py
deleted file mode 100644
index 87806dd60d60c5247554c9458de8fd8ca3f45f0f..0000000000000000000000000000000000000000
--- a/transformers/models/deberta/__init__.py
+++ /dev/null
@@ -1,120 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_tf_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_deberta": ["DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP", "DebertaConfig", "DebertaOnnxConfig"],
- "tokenization_deberta": ["DebertaTokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_deberta_fast"] = ["DebertaTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_deberta"] = [
- "DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DebertaForMaskedLM",
- "DebertaForQuestionAnswering",
- "DebertaForSequenceClassification",
- "DebertaForTokenClassification",
- "DebertaModel",
- "DebertaPreTrainedModel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_deberta"] = [
- "TF_DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFDebertaForMaskedLM",
- "TFDebertaForQuestionAnswering",
- "TFDebertaForSequenceClassification",
- "TFDebertaForTokenClassification",
- "TFDebertaModel",
- "TFDebertaPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_deberta import DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, DebertaConfig, DebertaOnnxConfig
- from .tokenization_deberta import DebertaTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_deberta_fast import DebertaTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_deberta import (
- DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
- DebertaForMaskedLM,
- DebertaForQuestionAnswering,
- DebertaForSequenceClassification,
- DebertaForTokenClassification,
- DebertaModel,
- DebertaPreTrainedModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_deberta import (
- TF_DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFDebertaForMaskedLM,
- TFDebertaForQuestionAnswering,
- TFDebertaForSequenceClassification,
- TFDebertaForTokenClassification,
- TFDebertaModel,
- TFDebertaPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deberta/__pycache__/__init__.cpython-310.pyc b/transformers/models/deberta/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index dc718652273cbf6ead07e27107d915735e0d14eb..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta/__pycache__/configuration_deberta.cpython-310.pyc b/transformers/models/deberta/__pycache__/configuration_deberta.cpython-310.pyc
deleted file mode 100644
index cfcbf9520d63f4a1292b93cc731d7a484d45f636..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta/__pycache__/configuration_deberta.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta/__pycache__/modeling_deberta.cpython-310.pyc b/transformers/models/deberta/__pycache__/modeling_deberta.cpython-310.pyc
deleted file mode 100644
index 15c4629b4ef28fcbebb1a3ad5ee4954b96b95022..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta/__pycache__/modeling_deberta.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta/__pycache__/modeling_tf_deberta.cpython-310.pyc b/transformers/models/deberta/__pycache__/modeling_tf_deberta.cpython-310.pyc
deleted file mode 100644
index fd1b18eabad3d4b0c38010f7f6376ddf93e78258..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta/__pycache__/modeling_tf_deberta.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta/__pycache__/tokenization_deberta.cpython-310.pyc b/transformers/models/deberta/__pycache__/tokenization_deberta.cpython-310.pyc
deleted file mode 100644
index 80f36730544baf83dd9c5619f90b9c2b3887cca9..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta/__pycache__/tokenization_deberta.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta/__pycache__/tokenization_deberta_fast.cpython-310.pyc b/transformers/models/deberta/__pycache__/tokenization_deberta_fast.cpython-310.pyc
deleted file mode 100644
index e0b55ca2188abb457976e1ab87faa32286619a9c..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta/__pycache__/tokenization_deberta_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta/configuration_deberta.py b/transformers/models/deberta/configuration_deberta.py
deleted file mode 100644
index 5907f0869d6821c49d6cbc417593b4b0a81e4768..0000000000000000000000000000000000000000
--- a/transformers/models/deberta/configuration_deberta.py
+++ /dev/null
@@ -1,193 +0,0 @@
-# coding=utf-8
-# Copyright 2020, Microsoft and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DeBERTa model configuration"""
-from collections import OrderedDict
-from typing import TYPE_CHECKING, Any, Mapping, Optional, Union
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-if TYPE_CHECKING:
- from ... import FeatureExtractionMixin, PreTrainedTokenizerBase, TensorType
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DebertaConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DebertaModel`] or a [`TFDebertaModel`]. It is
- used to instantiate a DeBERTa model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the DeBERTa
- [microsoft/deberta-base](https://huggingface.co/microsoft/deberta-base) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Arguments:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the DeBERTa model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`DebertaModel`] or [`TFDebertaModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"`, `"gelu"`, `"tanh"`, `"gelu_fast"`, `"mish"`, `"linear"`, `"sigmoid"` and `"gelu_new"`
- are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 2):
- The vocabulary size of the `token_type_ids` passed when calling [`DebertaModel`] or [`TFDebertaModel`].
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- relative_attention (`bool`, *optional*, defaults to `False`):
- Whether use relative position encoding.
- max_relative_positions (`int`, *optional*, defaults to 1):
- The range of relative positions `[-max_position_embeddings, max_position_embeddings]`. Use the same value
- as `max_position_embeddings`.
- pad_token_id (`int`, *optional*, defaults to 0):
- The value used to pad input_ids.
- position_biased_input (`bool`, *optional*, defaults to `True`):
- Whether add absolute position embedding to content embedding.
- pos_att_type (`List[str]`, *optional*):
- The type of relative position attention, it can be a combination of `["p2c", "c2p"]`, e.g. `["p2c"]`,
- `["p2c", "c2p"]`.
- layer_norm_eps (`float`, optional, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
-
- Example:
-
- ```python
- >>> from transformers import DebertaConfig, DebertaModel
-
- >>> # Initializing a DeBERTa microsoft/deberta-base style configuration
- >>> configuration = DebertaConfig()
-
- >>> # Initializing a model (with random weights) from the microsoft/deberta-base style configuration
- >>> model = DebertaModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "deberta"
-
- def __init__(
- self,
- vocab_size=50265,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=0,
- initializer_range=0.02,
- layer_norm_eps=1e-7,
- relative_attention=False,
- max_relative_positions=-1,
- pad_token_id=0,
- position_biased_input=True,
- pos_att_type=None,
- pooler_dropout=0,
- pooler_hidden_act="gelu",
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.initializer_range = initializer_range
- self.relative_attention = relative_attention
- self.max_relative_positions = max_relative_positions
- self.pad_token_id = pad_token_id
- self.position_biased_input = position_biased_input
-
- # Backwards compatibility
- if isinstance(pos_att_type, str):
- pos_att_type = [x.strip() for x in pos_att_type.lower().split("|")]
-
- self.pos_att_type = pos_att_type
- self.vocab_size = vocab_size
- self.layer_norm_eps = layer_norm_eps
-
- self.pooler_hidden_size = kwargs.get("pooler_hidden_size", hidden_size)
- self.pooler_dropout = pooler_dropout
- self.pooler_hidden_act = pooler_hidden_act
-
-
-# Copied from transformers.models.deberta_v2.configuration_deberta_v2.DebertaV2OnnxConfig
-class DebertaOnnxConfig(OnnxConfig):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- if self.task == "multiple-choice":
- dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
- else:
- dynamic_axis = {0: "batch", 1: "sequence"}
- if self._config.type_vocab_size > 0:
- return OrderedDict(
- [("input_ids", dynamic_axis), ("attention_mask", dynamic_axis), ("token_type_ids", dynamic_axis)]
- )
- else:
- return OrderedDict([("input_ids", dynamic_axis), ("attention_mask", dynamic_axis)])
-
- @property
- def default_onnx_opset(self) -> int:
- return 12
-
- def generate_dummy_inputs(
- self,
- preprocessor: Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"],
- batch_size: int = -1,
- seq_length: int = -1,
- num_choices: int = -1,
- is_pair: bool = False,
- framework: Optional["TensorType"] = None,
- num_channels: int = 3,
- image_width: int = 40,
- image_height: int = 40,
- tokenizer: "PreTrainedTokenizerBase" = None,
- ) -> Mapping[str, Any]:
- dummy_inputs = super().generate_dummy_inputs(preprocessor=preprocessor, framework=framework)
- if self._config.type_vocab_size == 0 and "token_type_ids" in dummy_inputs:
- del dummy_inputs["token_type_ids"]
- return dummy_inputs
diff --git a/transformers/models/deberta/modeling_deberta.py b/transformers/models/deberta/modeling_deberta.py
deleted file mode 100644
index 42dae5c80894a8ee8265fbb096fa579905aa1bb2..0000000000000000000000000000000000000000
--- a/transformers/models/deberta/modeling_deberta.py
+++ /dev/null
@@ -1,1426 +0,0 @@
-# coding=utf-8
-# Copyright 2020 Microsoft and the Hugging Face Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DeBERTa model."""
-
-from collections.abc import Sequence
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutput,
- MaskedLMOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import softmax_backward_data
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_deberta import DebertaConfig
-
-
-logger = logging.get_logger(__name__)
-_CONFIG_FOR_DOC = "DebertaConfig"
-_CHECKPOINT_FOR_DOC = "microsoft/deberta-base"
-
-# Masked LM docstring
-_CHECKPOINT_FOR_MASKED_LM = "lsanochkin/deberta-large-feedback"
-_MASKED_LM_EXPECTED_OUTPUT = "' Paris'"
-_MASKED_LM_EXPECTED_LOSS = "0.54"
-
-# QuestionAnswering docstring
-_CHECKPOINT_FOR_QA = "Palak/microsoft_deberta-large_squad"
-_QA_EXPECTED_OUTPUT = "' a nice puppet'"
-_QA_EXPECTED_LOSS = 0.14
-_QA_TARGET_START_INDEX = 12
-_QA_TARGET_END_INDEX = 14
-
-
-from ..deprecated._archive_maps import DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class ContextPooler(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.pooler_hidden_size, config.pooler_hidden_size)
- self.dropout = StableDropout(config.pooler_dropout)
- self.config = config
-
- def forward(self, hidden_states):
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
-
- context_token = hidden_states[:, 0]
- context_token = self.dropout(context_token)
- pooled_output = self.dense(context_token)
- pooled_output = ACT2FN[self.config.pooler_hidden_act](pooled_output)
- return pooled_output
-
- @property
- def output_dim(self):
- return self.config.hidden_size
-
-
-class XSoftmax(torch.autograd.Function):
- """
- Masked Softmax which is optimized for saving memory
-
- Args:
- input (`torch.tensor`): The input tensor that will apply softmax.
- mask (`torch.IntTensor`):
- The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
- dim (int): The dimension that will apply softmax
-
- Example:
-
- ```python
- >>> import torch
- >>> from transformers.models.deberta.modeling_deberta import XSoftmax
-
- >>> # Make a tensor
- >>> x = torch.randn([4, 20, 100])
-
- >>> # Create a mask
- >>> mask = (x > 0).int()
-
- >>> # Specify the dimension to apply softmax
- >>> dim = -1
-
- >>> y = XSoftmax.apply(x, mask, dim)
- ```"""
-
- @staticmethod
- def forward(self, input, mask, dim):
- self.dim = dim
- rmask = ~(mask.to(torch.bool))
-
- output = input.masked_fill(rmask, torch.tensor(torch.finfo(input.dtype).min))
- output = torch.softmax(output, self.dim)
- output.masked_fill_(rmask, 0)
- self.save_for_backward(output)
- return output
-
- @staticmethod
- def backward(self, grad_output):
- (output,) = self.saved_tensors
- inputGrad = softmax_backward_data(self, grad_output, output, self.dim, output)
- return inputGrad, None, None
-
- @staticmethod
- def symbolic(g, self, mask, dim):
- import torch.onnx.symbolic_helper as sym_help
- from torch.onnx.symbolic_opset9 import masked_fill, softmax
-
- mask_cast_value = g.op("Cast", mask, to_i=sym_help.cast_pytorch_to_onnx["Long"])
- r_mask = g.op(
- "Cast",
- g.op("Sub", g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)), mask_cast_value),
- to_i=sym_help.cast_pytorch_to_onnx["Bool"],
- )
- output = masked_fill(
- g, self, r_mask, g.op("Constant", value_t=torch.tensor(torch.finfo(self.type().dtype()).min))
- )
- output = softmax(g, output, dim)
- return masked_fill(g, output, r_mask, g.op("Constant", value_t=torch.tensor(0, dtype=torch.bool)))
-
-
-class DropoutContext(object):
- def __init__(self):
- self.dropout = 0
- self.mask = None
- self.scale = 1
- self.reuse_mask = True
-
-
-def get_mask(input, local_context):
- if not isinstance(local_context, DropoutContext):
- dropout = local_context
- mask = None
- else:
- dropout = local_context.dropout
- dropout *= local_context.scale
- mask = local_context.mask if local_context.reuse_mask else None
-
- if dropout > 0 and mask is None:
- mask = (1 - torch.empty_like(input).bernoulli_(1 - dropout)).to(torch.bool)
-
- if isinstance(local_context, DropoutContext):
- if local_context.mask is None:
- local_context.mask = mask
-
- return mask, dropout
-
-
-class XDropout(torch.autograd.Function):
- """Optimized dropout function to save computation and memory by using mask operation instead of multiplication."""
-
- @staticmethod
- def forward(ctx, input, local_ctx):
- mask, dropout = get_mask(input, local_ctx)
- ctx.scale = 1.0 / (1 - dropout)
- if dropout > 0:
- ctx.save_for_backward(mask)
- return input.masked_fill(mask, 0) * ctx.scale
- else:
- return input
-
- @staticmethod
- def backward(ctx, grad_output):
- if ctx.scale > 1:
- (mask,) = ctx.saved_tensors
- return grad_output.masked_fill(mask, 0) * ctx.scale, None
- else:
- return grad_output, None
-
- @staticmethod
- def symbolic(g: torch._C.Graph, input: torch._C.Value, local_ctx: Union[float, DropoutContext]) -> torch._C.Value:
- from torch.onnx import symbolic_opset12
-
- dropout_p = local_ctx
- if isinstance(local_ctx, DropoutContext):
- dropout_p = local_ctx.dropout
- # StableDropout only calls this function when training.
- train = True
- # TODO: We should check if the opset_version being used to export
- # is > 12 here, but there's no good way to do that. As-is, if the
- # opset_version < 12, export will fail with a CheckerError.
- # Once https://github.com/pytorch/pytorch/issues/78391 is fixed, do something like:
- # if opset_version < 12:
- # return torch.onnx.symbolic_opset9.dropout(g, input, dropout_p, train)
- return symbolic_opset12.dropout(g, input, dropout_p, train)
-
-
-class StableDropout(nn.Module):
- """
- Optimized dropout module for stabilizing the training
-
- Args:
- drop_prob (float): the dropout probabilities
- """
-
- def __init__(self, drop_prob):
- super().__init__()
- self.drop_prob = drop_prob
- self.count = 0
- self.context_stack = None
-
- def forward(self, x):
- """
- Call the module
-
- Args:
- x (`torch.tensor`): The input tensor to apply dropout
- """
- if self.training and self.drop_prob > 0:
- return XDropout.apply(x, self.get_context())
- return x
-
- def clear_context(self):
- self.count = 0
- self.context_stack = None
-
- def init_context(self, reuse_mask=True, scale=1):
- if self.context_stack is None:
- self.context_stack = []
- self.count = 0
- for c in self.context_stack:
- c.reuse_mask = reuse_mask
- c.scale = scale
-
- def get_context(self):
- if self.context_stack is not None:
- if self.count >= len(self.context_stack):
- self.context_stack.append(DropoutContext())
- ctx = self.context_stack[self.count]
- ctx.dropout = self.drop_prob
- self.count += 1
- return ctx
- else:
- return self.drop_prob
-
-
-class DebertaLayerNorm(nn.Module):
- """LayerNorm module in the TF style (epsilon inside the square root)."""
-
- def __init__(self, size, eps=1e-12):
- super().__init__()
- self.weight = nn.Parameter(torch.ones(size))
- self.bias = nn.Parameter(torch.zeros(size))
- self.variance_epsilon = eps
-
- def forward(self, hidden_states):
- input_type = hidden_states.dtype
- hidden_states = hidden_states.float()
- mean = hidden_states.mean(-1, keepdim=True)
- variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True)
- hidden_states = (hidden_states - mean) / torch.sqrt(variance + self.variance_epsilon)
- hidden_states = hidden_states.to(input_type)
- y = self.weight * hidden_states + self.bias
- return y
-
-
-class DebertaSelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.LayerNorm = DebertaLayerNorm(config.hidden_size, config.layer_norm_eps)
- self.dropout = StableDropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class DebertaAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.self = DisentangledSelfAttention(config)
- self.output = DebertaSelfOutput(config)
- self.config = config
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- output_attentions=False,
- query_states=None,
- relative_pos=None,
- rel_embeddings=None,
- ):
- self_output = self.self(
- hidden_states,
- attention_mask,
- output_attentions,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- )
- if output_attentions:
- self_output, att_matrix = self_output
- if query_states is None:
- query_states = hidden_states
- attention_output = self.output(self_output, query_states)
-
- if output_attentions:
- return (attention_output, att_matrix)
- else:
- return attention_output
-
-
-# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->Deberta
-class DebertaIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-class DebertaOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.LayerNorm = DebertaLayerNorm(config.hidden_size, config.layer_norm_eps)
- self.dropout = StableDropout(config.hidden_dropout_prob)
- self.config = config
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class DebertaLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.attention = DebertaAttention(config)
- self.intermediate = DebertaIntermediate(config)
- self.output = DebertaOutput(config)
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- query_states=None,
- relative_pos=None,
- rel_embeddings=None,
- output_attentions=False,
- ):
- attention_output = self.attention(
- hidden_states,
- attention_mask,
- output_attentions=output_attentions,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- )
- if output_attentions:
- attention_output, att_matrix = attention_output
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.output(intermediate_output, attention_output)
- if output_attentions:
- return (layer_output, att_matrix)
- else:
- return layer_output
-
-
-class DebertaEncoder(nn.Module):
- """Modified BertEncoder with relative position bias support"""
-
- def __init__(self, config):
- super().__init__()
- self.layer = nn.ModuleList([DebertaLayer(config) for _ in range(config.num_hidden_layers)])
- self.relative_attention = getattr(config, "relative_attention", False)
- if self.relative_attention:
- self.max_relative_positions = getattr(config, "max_relative_positions", -1)
- if self.max_relative_positions < 1:
- self.max_relative_positions = config.max_position_embeddings
- self.rel_embeddings = nn.Embedding(self.max_relative_positions * 2, config.hidden_size)
- self.gradient_checkpointing = False
-
- def get_rel_embedding(self):
- rel_embeddings = self.rel_embeddings.weight if self.relative_attention else None
- return rel_embeddings
-
- def get_attention_mask(self, attention_mask):
- if attention_mask.dim() <= 2:
- extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
- attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1)
- elif attention_mask.dim() == 3:
- attention_mask = attention_mask.unsqueeze(1)
-
- return attention_mask
-
- def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
- if self.relative_attention and relative_pos is None:
- q = query_states.size(-2) if query_states is not None else hidden_states.size(-2)
- relative_pos = build_relative_position(q, hidden_states.size(-2), hidden_states.device)
- return relative_pos
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- output_hidden_states=True,
- output_attentions=False,
- query_states=None,
- relative_pos=None,
- return_dict=True,
- ):
- attention_mask = self.get_attention_mask(attention_mask)
- relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
-
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- if isinstance(hidden_states, Sequence):
- next_kv = hidden_states[0]
- else:
- next_kv = hidden_states
- rel_embeddings = self.get_rel_embedding()
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- hidden_states = self._gradient_checkpointing_func(
- layer_module.__call__,
- next_kv,
- attention_mask,
- query_states,
- relative_pos,
- rel_embeddings,
- output_attentions,
- )
- else:
- hidden_states = layer_module(
- next_kv,
- attention_mask,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- output_attentions=output_attentions,
- )
-
- if output_attentions:
- hidden_states, att_m = hidden_states
-
- if query_states is not None:
- query_states = hidden_states
- if isinstance(hidden_states, Sequence):
- next_kv = hidden_states[i + 1] if i + 1 < len(self.layer) else None
- else:
- next_kv = hidden_states
-
- if output_attentions:
- all_attentions = all_attentions + (att_m,)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
-
-def build_relative_position(query_size, key_size, device):
- """
- Build relative position according to the query and key
-
- We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
- \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
- P_k\\)
-
- Args:
- query_size (int): the length of query
- key_size (int): the length of key
-
- Return:
- `torch.LongTensor`: A tensor with shape [1, query_size, key_size]
-
- """
-
- q_ids = torch.arange(query_size, dtype=torch.long, device=device)
- k_ids = torch.arange(key_size, dtype=torch.long, device=device)
- rel_pos_ids = q_ids[:, None] - k_ids.view(1, -1).repeat(query_size, 1)
- rel_pos_ids = rel_pos_ids[:query_size, :]
- rel_pos_ids = rel_pos_ids.unsqueeze(0)
- return rel_pos_ids
-
-
-@torch.jit.script
-def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos):
- return c2p_pos.expand([query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)])
-
-
-@torch.jit.script
-def p2c_dynamic_expand(c2p_pos, query_layer, key_layer):
- return c2p_pos.expand([query_layer.size(0), query_layer.size(1), key_layer.size(-2), key_layer.size(-2)])
-
-
-@torch.jit.script
-def pos_dynamic_expand(pos_index, p2c_att, key_layer):
- return pos_index.expand(p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2)))
-
-
-class DisentangledSelfAttention(nn.Module):
- """
- Disentangled self-attention module
-
- Parameters:
- config (`str`):
- A model config class instance with the configuration to build a new model. The schema is similar to
- *BertConfig*, for more details, please refer [`DebertaConfig`]
-
- """
-
- def __init__(self, config):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.in_proj = nn.Linear(config.hidden_size, self.all_head_size * 3, bias=False)
- self.q_bias = nn.Parameter(torch.zeros((self.all_head_size), dtype=torch.float))
- self.v_bias = nn.Parameter(torch.zeros((self.all_head_size), dtype=torch.float))
- self.pos_att_type = config.pos_att_type if config.pos_att_type is not None else []
-
- self.relative_attention = getattr(config, "relative_attention", False)
- self.talking_head = getattr(config, "talking_head", False)
-
- if self.talking_head:
- self.head_logits_proj = nn.Linear(config.num_attention_heads, config.num_attention_heads, bias=False)
- self.head_weights_proj = nn.Linear(config.num_attention_heads, config.num_attention_heads, bias=False)
-
- if self.relative_attention:
- self.max_relative_positions = getattr(config, "max_relative_positions", -1)
- if self.max_relative_positions < 1:
- self.max_relative_positions = config.max_position_embeddings
- self.pos_dropout = StableDropout(config.hidden_dropout_prob)
-
- if "c2p" in self.pos_att_type:
- self.pos_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=False)
- if "p2c" in self.pos_att_type:
- self.pos_q_proj = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = StableDropout(config.attention_probs_dropout_prob)
-
- def transpose_for_scores(self, x):
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, -1)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- output_attentions=False,
- query_states=None,
- relative_pos=None,
- rel_embeddings=None,
- ):
- """
- Call the module
-
- Args:
- hidden_states (`torch.FloatTensor`):
- Input states to the module usually the output from previous layer, it will be the Q,K and V in
- *Attention(Q,K,V)*
-
- attention_mask (`torch.BoolTensor`):
- An attention mask matrix of shape [*B*, *N*, *N*] where *B* is the batch size, *N* is the maximum
- sequence length in which element [i,j] = *1* means the *i* th token in the input can attend to the *j*
- th token.
-
- output_attentions (`bool`, optional):
- Whether return the attention matrix.
-
- query_states (`torch.FloatTensor`, optional):
- The *Q* state in *Attention(Q,K,V)*.
-
- relative_pos (`torch.LongTensor`):
- The relative position encoding between the tokens in the sequence. It's of shape [*B*, *N*, *N*] with
- values ranging in [*-max_relative_positions*, *max_relative_positions*].
-
- rel_embeddings (`torch.FloatTensor`):
- The embedding of relative distances. It's a tensor of shape [\\(2 \\times
- \\text{max_relative_positions}\\), *hidden_size*].
-
-
- """
- if query_states is None:
- qp = self.in_proj(hidden_states) # .split(self.all_head_size, dim=-1)
- query_layer, key_layer, value_layer = self.transpose_for_scores(qp).chunk(3, dim=-1)
- else:
-
- def linear(w, b, x):
- if b is not None:
- return torch.matmul(x, w.t()) + b.t()
- else:
- return torch.matmul(x, w.t()) # + b.t()
-
- ws = self.in_proj.weight.chunk(self.num_attention_heads * 3, dim=0)
- qkvw = [torch.cat([ws[i * 3 + k] for i in range(self.num_attention_heads)], dim=0) for k in range(3)]
- qkvb = [None] * 3
-
- q = linear(qkvw[0], qkvb[0], query_states.to(dtype=qkvw[0].dtype))
- k, v = [linear(qkvw[i], qkvb[i], hidden_states.to(dtype=qkvw[i].dtype)) for i in range(1, 3)]
- query_layer, key_layer, value_layer = [self.transpose_for_scores(x) for x in [q, k, v]]
-
- query_layer = query_layer + self.transpose_for_scores(self.q_bias[None, None, :])
- value_layer = value_layer + self.transpose_for_scores(self.v_bias[None, None, :])
-
- rel_att = None
- # Take the dot product between "query" and "key" to get the raw attention scores.
- scale_factor = 1 + len(self.pos_att_type)
- scale = torch.sqrt(torch.tensor(query_layer.size(-1), dtype=torch.float) * scale_factor)
- query_layer = query_layer / scale.to(dtype=query_layer.dtype)
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
- if self.relative_attention:
- rel_embeddings = self.pos_dropout(rel_embeddings)
- rel_att = self.disentangled_att_bias(query_layer, key_layer, relative_pos, rel_embeddings, scale_factor)
-
- if rel_att is not None:
- attention_scores = attention_scores + rel_att
-
- # bxhxlxd
- if self.talking_head:
- attention_scores = self.head_logits_proj(attention_scores.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-
- attention_probs = XSoftmax.apply(attention_scores, attention_mask, -1)
- attention_probs = self.dropout(attention_probs)
- if self.talking_head:
- attention_probs = self.head_weights_proj(attention_probs.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-
- context_layer = torch.matmul(attention_probs, value_layer)
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (-1,)
- context_layer = context_layer.view(new_context_layer_shape)
- if output_attentions:
- return (context_layer, attention_probs)
- else:
- return context_layer
-
- def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
- if relative_pos is None:
- q = query_layer.size(-2)
- relative_pos = build_relative_position(q, key_layer.size(-2), query_layer.device)
- if relative_pos.dim() == 2:
- relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
- elif relative_pos.dim() == 3:
- relative_pos = relative_pos.unsqueeze(1)
- # bxhxqxk
- elif relative_pos.dim() != 4:
- raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {relative_pos.dim()}")
-
- att_span = min(max(query_layer.size(-2), key_layer.size(-2)), self.max_relative_positions)
- relative_pos = relative_pos.long().to(query_layer.device)
- rel_embeddings = rel_embeddings[
- self.max_relative_positions - att_span : self.max_relative_positions + att_span, :
- ].unsqueeze(0)
-
- score = 0
-
- # content->position
- if "c2p" in self.pos_att_type:
- pos_key_layer = self.pos_proj(rel_embeddings)
- pos_key_layer = self.transpose_for_scores(pos_key_layer)
- c2p_att = torch.matmul(query_layer, pos_key_layer.transpose(-1, -2))
- c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
- c2p_att = torch.gather(c2p_att, dim=-1, index=c2p_dynamic_expand(c2p_pos, query_layer, relative_pos))
- score += c2p_att
-
- # position->content
- if "p2c" in self.pos_att_type:
- pos_query_layer = self.pos_q_proj(rel_embeddings)
- pos_query_layer = self.transpose_for_scores(pos_query_layer)
- pos_query_layer /= torch.sqrt(torch.tensor(pos_query_layer.size(-1), dtype=torch.float) * scale_factor)
- if query_layer.size(-2) != key_layer.size(-2):
- r_pos = build_relative_position(key_layer.size(-2), key_layer.size(-2), query_layer.device)
- else:
- r_pos = relative_pos
- p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
- p2c_att = torch.matmul(key_layer, pos_query_layer.transpose(-1, -2).to(dtype=key_layer.dtype))
- p2c_att = torch.gather(
- p2c_att, dim=-1, index=p2c_dynamic_expand(p2c_pos, query_layer, key_layer)
- ).transpose(-1, -2)
-
- if query_layer.size(-2) != key_layer.size(-2):
- pos_index = relative_pos[:, :, :, 0].unsqueeze(-1)
- p2c_att = torch.gather(p2c_att, dim=-2, index=pos_dynamic_expand(pos_index, p2c_att, key_layer))
- score += p2c_att
-
- return score
-
-
-class DebertaEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config):
- super().__init__()
- pad_token_id = getattr(config, "pad_token_id", 0)
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
- self.word_embeddings = nn.Embedding(config.vocab_size, self.embedding_size, padding_idx=pad_token_id)
-
- self.position_biased_input = getattr(config, "position_biased_input", True)
- if not self.position_biased_input:
- self.position_embeddings = None
- else:
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.embedding_size)
-
- if config.type_vocab_size > 0:
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, self.embedding_size)
-
- if self.embedding_size != config.hidden_size:
- self.embed_proj = nn.Linear(self.embedding_size, config.hidden_size, bias=False)
- self.LayerNorm = DebertaLayerNorm(config.hidden_size, config.layer_norm_eps)
- self.dropout = StableDropout(config.hidden_dropout_prob)
- self.config = config
-
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
-
- def forward(self, input_ids=None, token_type_ids=None, position_ids=None, mask=None, inputs_embeds=None):
- if input_ids is not None:
- input_shape = input_ids.size()
- else:
- input_shape = inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, :seq_length]
-
- if token_type_ids is None:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
-
- if self.position_embeddings is not None:
- position_embeddings = self.position_embeddings(position_ids.long())
- else:
- position_embeddings = torch.zeros_like(inputs_embeds)
-
- embeddings = inputs_embeds
- if self.position_biased_input:
- embeddings += position_embeddings
- if self.config.type_vocab_size > 0:
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
- embeddings += token_type_embeddings
-
- if self.embedding_size != self.config.hidden_size:
- embeddings = self.embed_proj(embeddings)
-
- embeddings = self.LayerNorm(embeddings)
-
- if mask is not None:
- if mask.dim() != embeddings.dim():
- if mask.dim() == 4:
- mask = mask.squeeze(1).squeeze(1)
- mask = mask.unsqueeze(2)
- mask = mask.to(embeddings.dtype)
-
- embeddings = embeddings * mask
-
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-class DebertaPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DebertaConfig
- base_model_prefix = "deberta"
- _keys_to_ignore_on_load_unexpected = ["position_embeddings"]
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
-
-DEBERTA_START_DOCSTRING = r"""
- The DeBERTa model was proposed in [DeBERTa: Decoding-enhanced BERT with Disentangled
- Attention](https://arxiv.org/abs/2006.03654) by Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen. It's build
- on top of BERT/RoBERTa with two improvements, i.e. disentangled attention and enhanced mask decoder. With those two
- improvements, it out perform BERT/RoBERTa on a majority of tasks with 80GB pretraining data.
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
-
- Parameters:
- config ([`DebertaConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DEBERTA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DeBERTa Model transformer outputting raw hidden-states without any specific head on top.",
- DEBERTA_START_DOCSTRING,
-)
-class DebertaModel(DebertaPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.embeddings = DebertaEmbeddings(config)
- self.encoder = DebertaEncoder(config)
- self.z_steps = 0
- self.config = config
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, new_embeddings):
- self.embeddings.word_embeddings = new_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- raise NotImplementedError("The prune function is not implemented in DeBERTa model.")
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device)
- if token_type_ids is None:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- mask=attention_mask,
- inputs_embeds=inputs_embeds,
- )
-
- encoder_outputs = self.encoder(
- embedding_output,
- attention_mask,
- output_hidden_states=True,
- output_attentions=output_attentions,
- return_dict=return_dict,
- )
- encoded_layers = encoder_outputs[1]
-
- if self.z_steps > 1:
- hidden_states = encoded_layers[-2]
- layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
- query_states = encoded_layers[-1]
- rel_embeddings = self.encoder.get_rel_embedding()
- attention_mask = self.encoder.get_attention_mask(attention_mask)
- rel_pos = self.encoder.get_rel_pos(embedding_output)
- for layer in layers[1:]:
- query_states = layer(
- hidden_states,
- attention_mask,
- output_attentions=False,
- query_states=query_states,
- relative_pos=rel_pos,
- rel_embeddings=rel_embeddings,
- )
- encoded_layers.append(query_states)
-
- sequence_output = encoded_layers[-1]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[(1 if output_hidden_states else 2) :]
-
- return BaseModelOutput(
- last_hidden_state=sequence_output,
- hidden_states=encoder_outputs.hidden_states if output_hidden_states else None,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings("""DeBERTa Model with a `language modeling` head on top.""", DEBERTA_START_DOCSTRING)
-class DebertaForMaskedLM(DebertaPreTrainedModel):
- _tied_weights_keys = ["cls.predictions.decoder.weight", "cls.predictions.decoder.bias"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.deberta = DebertaModel(config)
- self.cls = DebertaOnlyMLMHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.cls.predictions.decoder
-
- def set_output_embeddings(self, new_embeddings):
- self.cls.predictions.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_MASKED_LM,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="[MASK]",
- expected_output=_MASKED_LM_EXPECTED_OUTPUT,
- expected_loss=_MASKED_LM_EXPECTED_LOSS,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deberta(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- prediction_scores = self.cls(sequence_output)
-
- masked_lm_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss() # -100 index = padding token
- masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[1:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return MaskedLMOutput(
- loss=masked_lm_loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-class DebertaPredictionHeadTransform(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
-
- self.dense = nn.Linear(config.hidden_size, self.embedding_size)
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = ACT2FN[config.hidden_act]
- else:
- self.transform_act_fn = config.hidden_act
- self.LayerNorm = nn.LayerNorm(self.embedding_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
- return hidden_states
-
-
-class DebertaLMPredictionHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.transform = DebertaPredictionHeadTransform(config)
-
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
- # The output weights are the same as the input embeddings, but there is
- # an output-only bias for each token.
- self.decoder = nn.Linear(self.embedding_size, config.vocab_size, bias=False)
-
- self.bias = nn.Parameter(torch.zeros(config.vocab_size))
-
- # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
- self.decoder.bias = self.bias
-
- def forward(self, hidden_states):
- hidden_states = self.transform(hidden_states)
- hidden_states = self.decoder(hidden_states)
- return hidden_states
-
-
-# copied from transformers.models.bert.BertOnlyMLMHead with bert -> deberta
-class DebertaOnlyMLMHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.predictions = DebertaLMPredictionHead(config)
-
- def forward(self, sequence_output):
- prediction_scores = self.predictions(sequence_output)
- return prediction_scores
-
-
-@add_start_docstrings(
- """
- DeBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-class DebertaForSequenceClassification(DebertaPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- num_labels = getattr(config, "num_labels", 2)
- self.num_labels = num_labels
-
- self.deberta = DebertaModel(config)
- self.pooler = ContextPooler(config)
- output_dim = self.pooler.output_dim
-
- self.classifier = nn.Linear(output_dim, num_labels)
- drop_out = getattr(config, "cls_dropout", None)
- drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
- self.dropout = StableDropout(drop_out)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.deberta.get_input_embeddings()
-
- def set_input_embeddings(self, new_embeddings):
- self.deberta.set_input_embeddings(new_embeddings)
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deberta(
- input_ids,
- token_type_ids=token_type_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- encoder_layer = outputs[0]
- pooled_output = self.pooler(encoder_layer)
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- # regression task
- loss_fn = nn.MSELoss()
- logits = logits.view(-1).to(labels.dtype)
- loss = loss_fn(logits, labels.view(-1))
- elif labels.dim() == 1 or labels.size(-1) == 1:
- label_index = (labels >= 0).nonzero()
- labels = labels.long()
- if label_index.size(0) > 0:
- labeled_logits = torch.gather(
- logits, 0, label_index.expand(label_index.size(0), logits.size(1))
- )
- labels = torch.gather(labels, 0, label_index.view(-1))
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(labeled_logits.view(-1, self.num_labels).float(), labels.view(-1))
- else:
- loss = torch.tensor(0).to(logits)
- else:
- log_softmax = nn.LogSoftmax(-1)
- loss = -((log_softmax(logits) * labels).sum(-1)).mean()
- elif self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-class DebertaForTokenClassification(DebertaPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.deberta = DebertaModel(config)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deberta(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- DEBERTA_START_DOCSTRING,
-)
-class DebertaForQuestionAnswering(DebertaPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.deberta = DebertaModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_QA,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_QA_EXPECTED_OUTPUT,
- expected_loss=_QA_EXPECTED_LOSS,
- qa_target_start_index=_QA_TARGET_START_INDEX,
- qa_target_end_index=_QA_TARGET_END_INDEX,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deberta(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[1:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/deberta/modeling_tf_deberta.py b/transformers/models/deberta/modeling_tf_deberta.py
deleted file mode 100644
index 3cef6a50c873f438cd894b8231d890858ada44c2..0000000000000000000000000000000000000000
--- a/transformers/models/deberta/modeling_tf_deberta.py
+++ /dev/null
@@ -1,1644 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Microsoft and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 DeBERTa model."""
-
-
-from __future__ import annotations
-
-import math
-from typing import Dict, Optional, Sequence, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFMaskedLMOutput,
- TFQuestionAnsweringModelOutput,
- TFSequenceClassifierOutput,
- TFTokenClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFMaskedLanguageModelingLoss,
- TFModelInputType,
- TFPreTrainedModel,
- TFQuestionAnsweringLoss,
- TFSequenceClassificationLoss,
- TFTokenClassificationLoss,
- get_initializer,
- keras,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_deberta import DebertaConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-_CONFIG_FOR_DOC = "DebertaConfig"
-_CHECKPOINT_FOR_DOC = "kamalkraj/deberta-base"
-
-
-from ..deprecated._archive_maps import TF_DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class TFDebertaContextPooler(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
- self.dense = keras.layers.Dense(config.pooler_hidden_size, name="dense")
- self.dropout = TFDebertaStableDropout(config.pooler_dropout, name="dropout")
- self.config = config
-
- def call(self, hidden_states, training: bool = False):
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- context_token = hidden_states[:, 0]
- context_token = self.dropout(context_token, training=training)
- pooled_output = self.dense(context_token)
- pooled_output = get_tf_activation(self.config.pooler_hidden_act)(pooled_output)
- return pooled_output
-
- @property
- def output_dim(self) -> int:
- return self.config.hidden_size
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.pooler_hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
-
-
-class TFDebertaXSoftmax(keras.layers.Layer):
- """
- Masked Softmax which is optimized for saving memory
-
- Args:
- input (`tf.Tensor`): The input tensor that will apply softmax.
- mask (`tf.Tensor`): The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
- dim (int): The dimension that will apply softmax
- """
-
- def __init__(self, axis=-1, **kwargs):
- super().__init__(**kwargs)
- self.axis = axis
-
- def call(self, inputs: tf.Tensor, mask: tf.Tensor):
- rmask = tf.logical_not(tf.cast(mask, tf.bool))
- output = tf.where(rmask, float("-inf"), inputs)
- output = stable_softmax(output, self.axis)
- output = tf.where(rmask, 0.0, output)
- return output
-
-
-class TFDebertaStableDropout(keras.layers.Layer):
- """
- Optimized dropout module for stabilizing the training
-
- Args:
- drop_prob (float): the dropout probabilities
- """
-
- def __init__(self, drop_prob, **kwargs):
- super().__init__(**kwargs)
- self.drop_prob = drop_prob
-
- @tf.custom_gradient
- def xdropout(self, inputs):
- """
- Applies dropout to the inputs, as vanilla dropout, but also scales the remaining elements up by 1/drop_prob.
- """
- mask = tf.cast(
- 1
- - tf.compat.v1.distributions.Bernoulli(probs=1.0 - self.drop_prob).sample(sample_shape=shape_list(inputs)),
- tf.bool,
- )
- scale = tf.convert_to_tensor(1.0 / (1 - self.drop_prob), dtype=tf.float32)
- if self.drop_prob > 0:
- inputs = tf.where(mask, 0.0, inputs) * scale
-
- def grad(upstream):
- if self.drop_prob > 0:
- return tf.where(mask, 0.0, upstream) * scale
- else:
- return upstream
-
- return inputs, grad
-
- def call(self, inputs: tf.Tensor, training: tf.Tensor = False):
- if training:
- return self.xdropout(inputs)
- return inputs
-
-
-class TFDebertaLayerNorm(keras.layers.Layer):
- """LayerNorm module in the TF style (epsilon inside the square root)."""
-
- def __init__(self, size, eps=1e-12, **kwargs):
- super().__init__(**kwargs)
- self.size = size
- self.eps = eps
-
- def build(self, input_shape):
- self.gamma = self.add_weight(shape=[self.size], initializer=tf.ones_initializer(), name="weight")
- self.beta = self.add_weight(shape=[self.size], initializer=tf.zeros_initializer(), name="bias")
- return super().build(input_shape)
-
- def call(self, x: tf.Tensor) -> tf.Tensor:
- mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
- variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)
- std = tf.math.sqrt(variance + self.eps)
- return self.gamma * (x - mean) / std + self.beta
-
-
-class TFDebertaSelfOutput(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
- self.dense = keras.layers.Dense(config.hidden_size, name="dense")
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = TFDebertaStableDropout(config.hidden_dropout_prob, name="dropout")
- self.config = config
-
- def call(self, hidden_states, input_tensor, training: bool = False):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states, training=training)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
-
-
-class TFDebertaAttention(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
- self.self = TFDebertaDisentangledSelfAttention(config, name="self")
- self.dense_output = TFDebertaSelfOutput(config, name="output")
- self.config = config
-
- def call(
- self,
- input_tensor: tf.Tensor,
- attention_mask: tf.Tensor,
- query_states: tf.Tensor = None,
- relative_pos: tf.Tensor = None,
- rel_embeddings: tf.Tensor = None,
- output_attentions: bool = False,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- self_outputs = self.self(
- hidden_states=input_tensor,
- attention_mask=attention_mask,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- output_attentions=output_attentions,
- training=training,
- )
- if query_states is None:
- query_states = input_tensor
- attention_output = self.dense_output(
- hidden_states=self_outputs[0], input_tensor=query_states, training=training
- )
-
- output = (attention_output,) + self_outputs[1:]
-
- return output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "self", None) is not None:
- with tf.name_scope(self.self.name):
- self.self.build(None)
- if getattr(self, "dense_output", None) is not None:
- with tf.name_scope(self.dense_output.name):
- self.dense_output.build(None)
-
-
-class TFDebertaIntermediate(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
-
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.intermediate_act_fn = config.hidden_act
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFDebertaOutput(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = TFDebertaStableDropout(config.hidden_dropout_prob, name="dropout")
- self.config = config
-
- def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.dropout(hidden_states, training=training)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.intermediate_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
-
-
-class TFDebertaLayer(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.attention = TFDebertaAttention(config, name="attention")
- self.intermediate = TFDebertaIntermediate(config, name="intermediate")
- self.bert_output = TFDebertaOutput(config, name="output")
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- query_states: tf.Tensor = None,
- relative_pos: tf.Tensor = None,
- rel_embeddings: tf.Tensor = None,
- output_attentions: bool = False,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- attention_outputs = self.attention(
- input_tensor=hidden_states,
- attention_mask=attention_mask,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- output_attentions=output_attentions,
- training=training,
- )
- attention_output = attention_outputs[0]
- intermediate_output = self.intermediate(hidden_states=attention_output)
- layer_output = self.bert_output(
- hidden_states=intermediate_output, input_tensor=attention_output, training=training
- )
- outputs = (layer_output,) + attention_outputs[1:] # add attentions if we output them
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "intermediate", None) is not None:
- with tf.name_scope(self.intermediate.name):
- self.intermediate.build(None)
- if getattr(self, "bert_output", None) is not None:
- with tf.name_scope(self.bert_output.name):
- self.bert_output.build(None)
-
-
-class TFDebertaEncoder(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.layer = [TFDebertaLayer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]
- self.relative_attention = getattr(config, "relative_attention", False)
- self.config = config
- if self.relative_attention:
- self.max_relative_positions = getattr(config, "max_relative_positions", -1)
- if self.max_relative_positions < 1:
- self.max_relative_positions = config.max_position_embeddings
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if self.relative_attention:
- self.rel_embeddings = self.add_weight(
- name="rel_embeddings.weight",
- shape=[self.max_relative_positions * 2, self.config.hidden_size],
- initializer=get_initializer(self.config.initializer_range),
- )
- if getattr(self, "layer", None) is not None:
- for layer in self.layer:
- with tf.name_scope(layer.name):
- layer.build(None)
-
- def get_rel_embedding(self):
- rel_embeddings = self.rel_embeddings if self.relative_attention else None
- return rel_embeddings
-
- def get_attention_mask(self, attention_mask):
- if len(shape_list(attention_mask)) <= 2:
- extended_attention_mask = tf.expand_dims(tf.expand_dims(attention_mask, 1), 2)
- attention_mask = extended_attention_mask * tf.expand_dims(tf.squeeze(extended_attention_mask, -2), -1)
- attention_mask = tf.cast(attention_mask, tf.uint8)
- elif len(shape_list(attention_mask)) == 3:
- attention_mask = tf.expand_dims(attention_mask, 1)
-
- return attention_mask
-
- def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
- if self.relative_attention and relative_pos is None:
- q = shape_list(query_states)[-2] if query_states is not None else shape_list(hidden_states)[-2]
- relative_pos = build_relative_position(q, shape_list(hidden_states)[-2])
- return relative_pos
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- query_states: tf.Tensor = None,
- relative_pos: tf.Tensor = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- training: bool = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- attention_mask = self.get_attention_mask(attention_mask)
- relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
-
- if isinstance(hidden_states, Sequence):
- next_kv = hidden_states[0]
- else:
- next_kv = hidden_states
-
- rel_embeddings = self.get_rel_embedding()
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_outputs = layer_module(
- hidden_states=next_kv,
- attention_mask=attention_mask,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- output_attentions=output_attentions,
- training=training,
- )
- hidden_states = layer_outputs[0]
-
- if query_states is not None:
- query_states = hidden_states
- if isinstance(hidden_states, Sequence):
- next_kv = hidden_states[i + 1] if i + 1 < len(self.layer) else None
- else:
- next_kv = hidden_states
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- # Add last layer
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)
-
- return TFBaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
-
-def build_relative_position(query_size, key_size):
- """
- Build relative position according to the query and key
-
- We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
- \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
- P_k\\)
-
- Args:
- query_size (int): the length of query
- key_size (int): the length of key
-
- Return:
- `tf.Tensor`: A tensor with shape [1, query_size, key_size]
-
- """
- q_ids = tf.range(query_size, dtype=tf.int32)
- k_ids = tf.range(key_size, dtype=tf.int32)
- rel_pos_ids = q_ids[:, None] - tf.tile(tf.reshape(k_ids, [1, -1]), [query_size, 1])
- rel_pos_ids = rel_pos_ids[:query_size, :]
- rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)
- return tf.cast(rel_pos_ids, tf.int64)
-
-
-def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos):
- shapes = [
- shape_list(query_layer)[0],
- shape_list(query_layer)[1],
- shape_list(query_layer)[2],
- shape_list(relative_pos)[-1],
- ]
- return tf.broadcast_to(c2p_pos, shapes)
-
-
-def p2c_dynamic_expand(c2p_pos, query_layer, key_layer):
- shapes = [
- shape_list(query_layer)[0],
- shape_list(query_layer)[1],
- shape_list(key_layer)[-2],
- shape_list(key_layer)[-2],
- ]
- return tf.broadcast_to(c2p_pos, shapes)
-
-
-def pos_dynamic_expand(pos_index, p2c_att, key_layer):
- shapes = shape_list(p2c_att)[:2] + [shape_list(pos_index)[-2], shape_list(key_layer)[-2]]
- return tf.broadcast_to(pos_index, shapes)
-
-
-def torch_gather(x, indices, gather_axis):
- if gather_axis < 0:
- gather_axis = tf.rank(x) + gather_axis
-
- if gather_axis != tf.rank(x) - 1:
- pre_roll = tf.rank(x) - 1 - gather_axis
- permutation = tf.roll(tf.range(tf.rank(x)), pre_roll, axis=0)
- x = tf.transpose(x, perm=permutation)
- indices = tf.transpose(indices, perm=permutation)
- else:
- pre_roll = 0
-
- flat_x = tf.reshape(x, (-1, tf.shape(x)[-1]))
- flat_indices = tf.reshape(indices, (-1, tf.shape(indices)[-1]))
- gathered = tf.gather(flat_x, flat_indices, batch_dims=1)
- gathered = tf.reshape(gathered, tf.shape(indices))
-
- if pre_roll != 0:
- permutation = tf.roll(tf.range(tf.rank(x)), -pre_roll, axis=0)
- gathered = tf.transpose(gathered, perm=permutation)
-
- return gathered
-
-
-class TFDebertaDisentangledSelfAttention(keras.layers.Layer):
- """
- Disentangled self-attention module
-
- Parameters:
- config (`str`):
- A model config class instance with the configuration to build a new model. The schema is similar to
- *BertConfig*, for more details, please refer [`DebertaConfig`]
-
- """
-
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
- if config.hidden_size % config.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.in_proj = keras.layers.Dense(
- self.all_head_size * 3,
- kernel_initializer=get_initializer(config.initializer_range),
- name="in_proj",
- use_bias=False,
- )
- self.pos_att_type = config.pos_att_type if config.pos_att_type is not None else []
-
- self.relative_attention = getattr(config, "relative_attention", False)
- self.talking_head = getattr(config, "talking_head", False)
-
- if self.talking_head:
- self.head_logits_proj = keras.layers.Dense(
- self.num_attention_heads,
- kernel_initializer=get_initializer(config.initializer_range),
- name="head_logits_proj",
- use_bias=False,
- )
- self.head_weights_proj = keras.layers.Dense(
- self.num_attention_heads,
- kernel_initializer=get_initializer(config.initializer_range),
- name="head_weights_proj",
- use_bias=False,
- )
-
- self.softmax = TFDebertaXSoftmax(axis=-1)
-
- if self.relative_attention:
- self.max_relative_positions = getattr(config, "max_relative_positions", -1)
- if self.max_relative_positions < 1:
- self.max_relative_positions = config.max_position_embeddings
- self.pos_dropout = TFDebertaStableDropout(config.hidden_dropout_prob, name="pos_dropout")
- if "c2p" in self.pos_att_type:
- self.pos_proj = keras.layers.Dense(
- self.all_head_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="pos_proj",
- use_bias=False,
- )
- if "p2c" in self.pos_att_type:
- self.pos_q_proj = keras.layers.Dense(
- self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="pos_q_proj"
- )
-
- self.dropout = TFDebertaStableDropout(config.attention_probs_dropout_prob, name="dropout")
- self.config = config
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- self.q_bias = self.add_weight(
- name="q_bias", shape=(self.all_head_size), initializer=keras.initializers.Zeros()
- )
- self.v_bias = self.add_weight(
- name="v_bias", shape=(self.all_head_size), initializer=keras.initializers.Zeros()
- )
- if getattr(self, "in_proj", None) is not None:
- with tf.name_scope(self.in_proj.name):
- self.in_proj.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
- if getattr(self, "head_logits_proj", None) is not None:
- with tf.name_scope(self.head_logits_proj.name):
- self.head_logits_proj.build(None)
- if getattr(self, "head_weights_proj", None) is not None:
- with tf.name_scope(self.head_weights_proj.name):
- self.head_weights_proj.build(None)
- if getattr(self, "pos_dropout", None) is not None:
- with tf.name_scope(self.pos_dropout.name):
- self.pos_dropout.build(None)
- if getattr(self, "pos_proj", None) is not None:
- with tf.name_scope(self.pos_proj.name):
- self.pos_proj.build([self.config.hidden_size])
- if getattr(self, "pos_q_proj", None) is not None:
- with tf.name_scope(self.pos_q_proj.name):
- self.pos_q_proj.build([self.config.hidden_size])
-
- def transpose_for_scores(self, tensor: tf.Tensor) -> tf.Tensor:
- shape = shape_list(tensor)[:-1] + [self.num_attention_heads, -1]
- # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]
- tensor = tf.reshape(tensor=tensor, shape=shape)
-
- # Transpose the tensor from [batch_size, seq_length, num_attention_heads, attention_head_size] to [batch_size, num_attention_heads, seq_length, attention_head_size]
- return tf.transpose(tensor, perm=[0, 2, 1, 3])
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- query_states: tf.Tensor = None,
- relative_pos: tf.Tensor = None,
- rel_embeddings: tf.Tensor = None,
- output_attentions: bool = False,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- """
- Call the module
-
- Args:
- hidden_states (`tf.Tensor`):
- Input states to the module usually the output from previous layer, it will be the Q,K and V in
- *Attention(Q,K,V)*
-
- attention_mask (`tf.Tensor`):
- An attention mask matrix of shape [*B*, *N*, *N*] where *B* is the batch size, *N* is the maximum
- sequence length in which element [i,j] = *1* means the *i* th token in the input can attend to the *j*
- th token.
-
- return_att (`bool`, optional):
- Whether return the attention matrix.
-
- query_states (`tf.Tensor`, optional):
- The *Q* state in *Attention(Q,K,V)*.
-
- relative_pos (`tf.Tensor`):
- The relative position encoding between the tokens in the sequence. It's of shape [*B*, *N*, *N*] with
- values ranging in [*-max_relative_positions*, *max_relative_positions*].
-
- rel_embeddings (`tf.Tensor`):
- The embedding of relative distances. It's a tensor of shape [\\(2 \\times
- \\text{max_relative_positions}\\), *hidden_size*].
-
-
- """
- if query_states is None:
- qp = self.in_proj(hidden_states) # .split(self.all_head_size, dim=-1)
- query_layer, key_layer, value_layer = tf.split(
- self.transpose_for_scores(qp), num_or_size_splits=3, axis=-1
- )
- else:
-
- def linear(w, b, x):
- out = tf.matmul(x, w, transpose_b=True)
- if b is not None:
- out += tf.transpose(b)
- return out
-
- ws = tf.split(
- tf.transpose(self.in_proj.weight[0]), num_or_size_splits=self.num_attention_heads * 3, axis=0
- )
- qkvw = tf.TensorArray(dtype=tf.float32, size=3)
- for k in tf.range(3):
- qkvw_inside = tf.TensorArray(dtype=tf.float32, size=self.num_attention_heads)
- for i in tf.range(self.num_attention_heads):
- qkvw_inside = qkvw_inside.write(i, ws[i * 3 + k])
- qkvw = qkvw.write(k, qkvw_inside.concat())
- qkvb = [None] * 3
-
- q = linear(qkvw[0], qkvb[0], query_states)
- k = linear(qkvw[1], qkvb[1], hidden_states)
- v = linear(qkvw[2], qkvb[2], hidden_states)
- query_layer = self.transpose_for_scores(q)
- key_layer = self.transpose_for_scores(k)
- value_layer = self.transpose_for_scores(v)
-
- query_layer = query_layer + self.transpose_for_scores(self.q_bias[None, None, :])
- value_layer = value_layer + self.transpose_for_scores(self.v_bias[None, None, :])
-
- rel_att = None
- # Take the dot product between "query" and "key" to get the raw attention scores.
- scale_factor = 1 + len(self.pos_att_type)
- scale = math.sqrt(shape_list(query_layer)[-1] * scale_factor)
- query_layer = query_layer / scale
-
- attention_scores = tf.matmul(query_layer, tf.transpose(key_layer, [0, 1, 3, 2]))
- if self.relative_attention:
- rel_embeddings = self.pos_dropout(rel_embeddings, training=training)
- rel_att = self.disentangled_att_bias(query_layer, key_layer, relative_pos, rel_embeddings, scale_factor)
-
- if rel_att is not None:
- attention_scores = attention_scores + rel_att
-
- if self.talking_head:
- attention_scores = tf.transpose(
- self.head_logits_proj(tf.transpose(attention_scores, [0, 2, 3, 1])), [0, 3, 1, 2]
- )
-
- attention_probs = self.softmax(attention_scores, attention_mask)
- attention_probs = self.dropout(attention_probs, training=training)
- if self.talking_head:
- attention_probs = tf.transpose(
- self.head_weights_proj(tf.transpose(attention_probs, [0, 2, 3, 1])), [0, 3, 1, 2]
- )
-
- context_layer = tf.matmul(attention_probs, value_layer)
- context_layer = tf.transpose(context_layer, [0, 2, 1, 3])
- context_layer_shape = shape_list(context_layer)
- # Set the final dimension here explicitly.
- # Calling tf.reshape(context_layer, (*context_layer_shape[:-2], -1)) raises an error when executing
- # the model in graph mode as context_layer is reshaped to (None, 7, None) and Dense layer in TFDebertaV2SelfOutput
- # requires final input dimension to be defined
- new_context_layer_shape = context_layer_shape[:-2] + [context_layer_shape[-2] * context_layer_shape[-1]]
- context_layer = tf.reshape(context_layer, new_context_layer_shape)
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
- return outputs
-
- def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
- if relative_pos is None:
- q = shape_list(query_layer)[-2]
- relative_pos = build_relative_position(q, shape_list(key_layer)[-2])
- shape_list_pos = shape_list(relative_pos)
- if len(shape_list_pos) == 2:
- relative_pos = tf.expand_dims(tf.expand_dims(relative_pos, 0), 0)
- elif len(shape_list_pos) == 3:
- relative_pos = tf.expand_dims(relative_pos, 1)
- # bxhxqxk
- elif len(shape_list_pos) != 4:
- raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {len(shape_list_pos)}")
-
- att_span = tf.cast(
- tf.minimum(
- tf.maximum(shape_list(query_layer)[-2], shape_list(key_layer)[-2]), self.max_relative_positions
- ),
- tf.int64,
- )
- rel_embeddings = tf.expand_dims(
- rel_embeddings[self.max_relative_positions - att_span : self.max_relative_positions + att_span, :], 0
- )
-
- score = 0
-
- # content->position
- if "c2p" in self.pos_att_type:
- pos_key_layer = self.pos_proj(rel_embeddings)
- pos_key_layer = self.transpose_for_scores(pos_key_layer)
- c2p_att = tf.matmul(query_layer, tf.transpose(pos_key_layer, [0, 1, 3, 2]))
- c2p_pos = tf.clip_by_value(relative_pos + att_span, 0, att_span * 2 - 1)
- c2p_att = torch_gather(c2p_att, c2p_dynamic_expand(c2p_pos, query_layer, relative_pos), -1)
- score += c2p_att
-
- # position->content
- if "p2c" in self.pos_att_type:
- pos_query_layer = self.pos_q_proj(rel_embeddings)
- pos_query_layer = self.transpose_for_scores(pos_query_layer)
- pos_query_layer /= tf.math.sqrt(tf.cast(shape_list(pos_query_layer)[-1] * scale_factor, dtype=tf.float32))
- if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
- r_pos = build_relative_position(shape_list(key_layer)[-2], shape_list(key_layer)[-2])
- else:
- r_pos = relative_pos
- p2c_pos = tf.clip_by_value(-r_pos + att_span, 0, att_span * 2 - 1)
- p2c_att = tf.matmul(key_layer, tf.transpose(pos_query_layer, [0, 1, 3, 2]))
- p2c_att = tf.transpose(
- torch_gather(p2c_att, p2c_dynamic_expand(p2c_pos, query_layer, key_layer), -1), [0, 1, 3, 2]
- )
- if shape_list(query_layer)[-2] != shape_list(key_layer)[-2]:
- pos_index = tf.expand_dims(relative_pos[:, :, :, 0], -1)
- p2c_att = torch_gather(p2c_att, pos_dynamic_expand(pos_index, p2c_att, key_layer), -2)
- score += p2c_att
-
- return score
-
-
-class TFDebertaEmbeddings(keras.layers.Layer):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
- self.hidden_size = config.hidden_size
- self.max_position_embeddings = config.max_position_embeddings
- self.position_biased_input = getattr(config, "position_biased_input", True)
- self.initializer_range = config.initializer_range
- if self.embedding_size != config.hidden_size:
- self.embed_proj = keras.layers.Dense(
- config.hidden_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="embed_proj",
- use_bias=False,
- )
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = TFDebertaStableDropout(config.hidden_dropout_prob, name="dropout")
-
- def build(self, input_shape=None):
- with tf.name_scope("word_embeddings"):
- self.weight = self.add_weight(
- name="weight",
- shape=[self.config.vocab_size, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
-
- with tf.name_scope("token_type_embeddings"):
- if self.config.type_vocab_size > 0:
- self.token_type_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.config.type_vocab_size, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
- else:
- self.token_type_embeddings = None
-
- with tf.name_scope("position_embeddings"):
- if self.position_biased_input:
- self.position_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.max_position_embeddings, self.hidden_size],
- initializer=get_initializer(self.initializer_range),
- )
- else:
- self.position_embeddings = None
-
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
- if getattr(self, "embed_proj", None) is not None:
- with tf.name_scope(self.embed_proj.name):
- self.embed_proj.build([None, None, self.embedding_size])
-
- def call(
- self,
- input_ids: tf.Tensor = None,
- position_ids: tf.Tensor = None,
- token_type_ids: tf.Tensor = None,
- inputs_embeds: tf.Tensor = None,
- mask: tf.Tensor = None,
- training: bool = False,
- ) -> tf.Tensor:
- """
- Applies embedding based on inputs tensor.
-
- Returns:
- final_embeddings (`tf.Tensor`): output embedding tensor.
- """
- if input_ids is None and inputs_embeds is None:
- raise ValueError("Need to provide either `input_ids` or `input_embeds`.")
-
- if input_ids is not None:
- check_embeddings_within_bounds(input_ids, self.config.vocab_size)
- inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
-
- input_shape = shape_list(inputs_embeds)[:-1]
-
- if token_type_ids is None:
- token_type_ids = tf.fill(dims=input_shape, value=0)
-
- if position_ids is None:
- position_ids = tf.expand_dims(tf.range(start=0, limit=input_shape[-1]), axis=0)
-
- final_embeddings = inputs_embeds
- if self.position_biased_input:
- position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
- final_embeddings += position_embeds
- if self.config.type_vocab_size > 0:
- token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)
- final_embeddings += token_type_embeds
-
- if self.embedding_size != self.hidden_size:
- final_embeddings = self.embed_proj(final_embeddings)
-
- final_embeddings = self.LayerNorm(final_embeddings)
-
- if mask is not None:
- if len(shape_list(mask)) != len(shape_list(final_embeddings)):
- if len(shape_list(mask)) == 4:
- mask = tf.squeeze(tf.squeeze(mask, axis=1), axis=1)
- mask = tf.cast(tf.expand_dims(mask, axis=2), tf.float32)
-
- final_embeddings = final_embeddings * mask
-
- final_embeddings = self.dropout(final_embeddings, training=training)
-
- return final_embeddings
-
-
-class TFDebertaPredictionHeadTransform(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
-
- self.dense = keras.layers.Dense(
- units=self.embedding_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="dense",
- )
-
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.transform_act_fn = config.hidden_act
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.embedding_size])
-
-
-class TFDebertaLMPredictionHead(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, input_embeddings: keras.layers.Layer, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
-
- self.transform = TFDebertaPredictionHeadTransform(config, name="transform")
-
- # The output weights are the same as the input embeddings, but there is
- # an output-only bias for each token.
- self.input_embeddings = input_embeddings
-
- def build(self, input_shape=None):
- self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")
-
- if self.built:
- return
- self.built = True
- if getattr(self, "transform", None) is not None:
- with tf.name_scope(self.transform.name):
- self.transform.build(None)
-
- def get_output_embeddings(self) -> keras.layers.Layer:
- return self.input_embeddings
-
- def set_output_embeddings(self, value: tf.Variable):
- self.input_embeddings.weight = value
- self.input_embeddings.vocab_size = shape_list(value)[0]
-
- def get_bias(self) -> Dict[str, tf.Variable]:
- return {"bias": self.bias}
-
- def set_bias(self, value: tf.Variable):
- self.bias = value["bias"]
- self.config.vocab_size = shape_list(value["bias"])[0]
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.transform(hidden_states=hidden_states)
- seq_length = shape_list(hidden_states)[1]
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.embedding_size])
- hidden_states = tf.matmul(a=hidden_states, b=self.input_embeddings.weight, transpose_b=True)
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])
- hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)
-
- return hidden_states
-
-
-class TFDebertaOnlyMLMHead(keras.layers.Layer):
- def __init__(self, config: DebertaConfig, input_embeddings: keras.layers.Layer, **kwargs):
- super().__init__(**kwargs)
- self.predictions = TFDebertaLMPredictionHead(config, input_embeddings, name="predictions")
-
- def call(self, sequence_output: tf.Tensor) -> tf.Tensor:
- prediction_scores = self.predictions(hidden_states=sequence_output)
-
- return prediction_scores
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "predictions", None) is not None:
- with tf.name_scope(self.predictions.name):
- self.predictions.build(None)
-
-
-# @keras_serializable
-class TFDebertaMainLayer(keras.layers.Layer):
- config_class = DebertaConfig
-
- def __init__(self, config: DebertaConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
-
- self.embeddings = TFDebertaEmbeddings(config, name="embeddings")
- self.encoder = TFDebertaEncoder(config, name="encoder")
-
- def get_input_embeddings(self) -> keras.layers.Layer:
- return self.embeddings
-
- def set_input_embeddings(self, value: tf.Variable):
- self.embeddings.weight = value
- self.embeddings.vocab_size = shape_list(value)[0]
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- raise NotImplementedError
-
- @unpack_inputs
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = tf.fill(dims=input_shape, value=1)
-
- if token_type_ids is None:
- token_type_ids = tf.fill(dims=input_shape, value=0)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- mask=attention_mask,
- training=training,
- )
-
- encoder_outputs = self.encoder(
- hidden_states=embedding_output,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = encoder_outputs[0]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[1:]
-
- return TFBaseModelOutput(
- last_hidden_state=sequence_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
-
-
-class TFDebertaPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DebertaConfig
- base_model_prefix = "deberta"
-
-
-DEBERTA_START_DOCSTRING = r"""
- The DeBERTa model was proposed in [DeBERTa: Decoding-enhanced BERT with Disentangled
- Attention](https://arxiv.org/abs/2006.03654) by Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen. It's build
- on top of BERT/RoBERTa with two improvements, i.e. disentangled attention and enhanced mask decoder. With those two
- improvements, it out perform BERT/RoBERTa on a majority of tasks with 80GB pretraining data.
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`DebertaConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DEBERTA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`np.ndarray`, `tf.Tensor`, `List[tf.Tensor]` ``Dict[str, tf.Tensor]` or `Dict[str, np.ndarray]` and each example must have the shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- inputs_embeds (`np.ndarray` or `tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput``] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DeBERTa Model transformer outputting raw hidden-states without any specific head on top.",
- DEBERTA_START_DOCSTRING,
-)
-class TFDebertaModel(TFDebertaPreTrainedModel):
- def __init__(self, config: DebertaConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.deberta = TFDebertaMainLayer(config, name="deberta")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
-
-
-@add_start_docstrings("""DeBERTa Model with a `language modeling` head on top.""", DEBERTA_START_DOCSTRING)
-class TFDebertaForMaskedLM(TFDebertaPreTrainedModel, TFMaskedLanguageModelingLoss):
- def __init__(self, config: DebertaConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- if config.is_decoder:
- logger.warning(
- "If you want to use `TFDebertaForMaskedLM` make sure `config.is_decoder=False` for "
- "bi-directional self-attention."
- )
-
- self.deberta = TFDebertaMainLayer(config, name="deberta")
- self.mlm = TFDebertaOnlyMLMHead(config, input_embeddings=self.deberta.embeddings, name="cls")
-
- def get_lm_head(self) -> keras.layers.Layer:
- return self.mlm.predictions
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFMaskedLMOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- prediction_scores = self.mlm(sequence_output=sequence_output, training=training)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=prediction_scores)
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFMaskedLMOutput(
- loss=loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "mlm", None) is not None:
- with tf.name_scope(self.mlm.name):
- self.mlm.build(None)
-
-
-@add_start_docstrings(
- """
- DeBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-class TFDebertaForSequenceClassification(TFDebertaPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: DebertaConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
-
- self.deberta = TFDebertaMainLayer(config, name="deberta")
- self.pooler = TFDebertaContextPooler(config, name="pooler")
-
- drop_out = getattr(config, "cls_dropout", None)
- drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
- self.dropout = TFDebertaStableDropout(drop_out, name="cls_dropout")
- self.classifier = keras.layers.Dense(
- units=config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- name="classifier",
- )
- self.output_dim = self.pooler.output_dim
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- pooled_output = self.pooler(sequence_output, training=training)
- pooled_output = self.dropout(pooled_output, training=training)
- logits = self.classifier(pooled_output)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "pooler", None) is not None:
- with tf.name_scope(self.pooler.name):
- self.pooler.build(None)
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.output_dim])
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-class TFDebertaForTokenClassification(TFDebertaPreTrainedModel, TFTokenClassificationLoss):
- def __init__(self, config: DebertaConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
-
- self.deberta = TFDebertaMainLayer(config, name="deberta")
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.classifier = keras.layers.Dense(
- units=config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFTokenClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- sequence_output = self.dropout(sequence_output, training=training)
- logits = self.classifier(inputs=sequence_output)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFTokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- DEBERTA_START_DOCSTRING,
-)
-class TFDebertaForQuestionAnswering(TFDebertaPreTrainedModel, TFQuestionAnsweringLoss):
- def __init__(self, config: DebertaConfig, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
-
- self.deberta = TFDebertaMainLayer(config, name="deberta")
- self.qa_outputs = keras.layers.Dense(
- units=config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFQuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- start_positions: np.ndarray | tf.Tensor | None = None,
- end_positions: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFQuestionAnsweringModelOutput, Tuple[tf.Tensor]]:
- r"""
- start_positions (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- logits = self.qa_outputs(inputs=sequence_output)
- start_logits, end_logits = tf.split(value=logits, num_or_size_splits=2, axis=-1)
- start_logits = tf.squeeze(input=start_logits, axis=-1)
- end_logits = tf.squeeze(input=end_logits, axis=-1)
- loss = None
-
- if start_positions is not None and end_positions is not None:
- labels = {"start_position": start_positions}
- labels["end_position"] = end_positions
- loss = self.hf_compute_loss(labels=labels, logits=(start_logits, end_logits))
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFQuestionAnsweringModelOutput(
- loss=loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.config.hidden_size])
diff --git a/transformers/models/deberta/tokenization_deberta.py b/transformers/models/deberta/tokenization_deberta.py
deleted file mode 100644
index b846a7891562d6386d40f342c47211a5b53857e4..0000000000000000000000000000000000000000
--- a/transformers/models/deberta/tokenization_deberta.py
+++ /dev/null
@@ -1,393 +0,0 @@
-# coding=utf-8
-# Copyright 2020 Microsoft and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Tokenization class for model DeBERTa."""
-
-import json
-import os
-from typing import List, Optional, Tuple
-
-import regex as re
-
-from ...tokenization_utils import AddedToken, PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt"}
-
-
-# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode
-def bytes_to_unicode():
- """
- Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
- characters the bpe code barfs on.
-
- The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
- if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
- decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
- tables between utf-8 bytes and unicode strings.
- """
- bs = (
- list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
- )
- cs = bs[:]
- n = 0
- for b in range(2**8):
- if b not in bs:
- bs.append(b)
- cs.append(2**8 + n)
- n += 1
- cs = [chr(n) for n in cs]
- return dict(zip(bs, cs))
-
-
-# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs
-def get_pairs(word):
- """
- Return set of symbol pairs in a word.
-
- Word is represented as tuple of symbols (symbols being variable-length strings).
- """
- pairs = set()
- prev_char = word[0]
- for char in word[1:]:
- pairs.add((prev_char, char))
- prev_char = char
- return pairs
-
-
-class DebertaTokenizer(PreTrainedTokenizer):
- """
- Construct a DeBERTa tokenizer. Based on byte-level Byte-Pair-Encoding.
-
- This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
- be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
- ```python
- >>> from transformers import DebertaTokenizer
-
- >>> tokenizer = DebertaTokenizer.from_pretrained("microsoft/deberta-base")
- >>> tokenizer("Hello world")["input_ids"]
- [1, 31414, 232, 2]
-
- >>> tokenizer(" Hello world")["input_ids"]
- [1, 20920, 232, 2]
- ```
-
- You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer or when you
- call it on some text, but since the model was not pretrained this way, it might yield a decrease in performance.
-
-
-
- When used with `is_split_into_words=True`, this tokenizer will add a space before each word (even the first one).
-
-
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- merges_file (`str`):
- Path to the merges file.
- errors (`str`, *optional*, defaults to `"replace"`):
- Paradigm to follow when decoding bytes to UTF-8. See
- [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
- bos_token (`str`, *optional*, defaults to `"[CLS]"`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `"[SEP]"`):
- The end of sequence token.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial space to the input. This allows to treat the leading word just as any
- other word. (Deberta tokenizer detect beginning of words by the preceding space).
- add_bos_token (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial <|endoftext|> to the input. This allows to treat the leading word just as
- any other word.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
-
- def __init__(
- self,
- vocab_file,
- merges_file,
- errors="replace",
- bos_token="[CLS]",
- eos_token="[SEP]",
- sep_token="[SEP]",
- cls_token="[CLS]",
- unk_token="[UNK]",
- pad_token="[PAD]",
- mask_token="[MASK]",
- add_prefix_space=False,
- add_bos_token=False,
- **kwargs,
- ):
- bos_token = AddedToken(bos_token, special=True) if isinstance(bos_token, str) else bos_token
- eos_token = AddedToken(eos_token, special=True) if isinstance(eos_token, str) else eos_token
- sep_token = AddedToken(sep_token, special=True) if isinstance(sep_token, str) else sep_token
- cls_token = AddedToken(cls_token, special=True) if isinstance(cls_token, str) else cls_token
- unk_token = AddedToken(unk_token, special=True) if isinstance(unk_token, str) else unk_token
- pad_token = AddedToken(pad_token, special=True) if isinstance(pad_token, str) else pad_token
-
- # Mask token behave like a normal word, i.e. include the space before it
- mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token
- self.add_bos_token = add_bos_token
-
- with open(vocab_file, encoding="utf-8") as vocab_handle:
- self.encoder = json.load(vocab_handle)
- self.decoder = {v: k for k, v in self.encoder.items()}
- self.errors = errors # how to handle errors in decoding
- self.byte_encoder = bytes_to_unicode()
- self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
- with open(merges_file, encoding="utf-8") as merges_handle:
- bpe_merges = merges_handle.read().split("\n")[1:-1]
- bpe_merges = [tuple(merge.split()) for merge in bpe_merges]
- self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
- self.cache = {}
- self.add_prefix_space = add_prefix_space
-
- # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
- self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
-
- super().__init__(
- errors=errors,
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- sep_token=sep_token,
- cls_token=cls_token,
- pad_token=pad_token,
- mask_token=mask_token,
- add_prefix_space=add_prefix_space,
- add_bos_token=add_bos_token,
- **kwargs,
- )
-
- @property
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.vocab_size
- def vocab_size(self):
- return len(self.encoder)
-
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.get_vocab
- def get_vocab(self):
- return dict(self.encoder, **self.added_tokens_encoder)
-
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe
- def bpe(self, token):
- if token in self.cache:
- return self.cache[token]
- word = tuple(token)
- pairs = get_pairs(word)
-
- if not pairs:
- return token
-
- while True:
- bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
- if bigram not in self.bpe_ranks:
- break
- first, second = bigram
- new_word = []
- i = 0
- while i < len(word):
- try:
- j = word.index(first, i)
- except ValueError:
- new_word.extend(word[i:])
- break
- else:
- new_word.extend(word[i:j])
- i = j
-
- if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
- new_word.append(first + second)
- i += 2
- else:
- new_word.append(word[i])
- i += 1
- new_word = tuple(new_word)
- word = new_word
- if len(word) == 1:
- break
- else:
- pairs = get_pairs(word)
- word = " ".join(word)
- self.cache[token] = word
- return word
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A DeBERTa sequence has the following format:
-
- - single sequence: [CLS] X [SEP]
- - pair of sequences: [CLS] A [SEP] B [SEP]
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is None:
- return [1] + ([0] * len(token_ids_0)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A DeBERTa
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
-
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize
- def _tokenize(self, text):
- """Tokenize a string."""
- bpe_tokens = []
- for token in re.findall(self.pat, text):
- token = "".join(
- self.byte_encoder[b] for b in token.encode("utf-8")
- ) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
- bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
- return bpe_tokens
-
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index)
-
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- text = "".join(tokens)
- text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
- return text
-
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- merge_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
- )
-
- with open(vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
- index = 0
- with open(merge_file, "w", encoding="utf-8") as writer:
- writer.write("#version: 0.2\n")
- for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
- " Please check that the tokenizer is not corrupted!"
- )
- index = token_index
- writer.write(" ".join(bpe_tokens) + "\n")
- index += 1
-
- return vocab_file, merge_file
-
- def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
- add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space)
- if (is_split_into_words or add_prefix_space) and (len(text) > 0 and not text[0].isspace()):
- text = " " + text
- return (text, kwargs)
diff --git a/transformers/models/deberta/tokenization_deberta_fast.py b/transformers/models/deberta/tokenization_deberta_fast.py
deleted file mode 100644
index 07226443d30a9c0cbe3d9f970e32343dac04ca65..0000000000000000000000000000000000000000
--- a/transformers/models/deberta/tokenization_deberta_fast.py
+++ /dev/null
@@ -1,247 +0,0 @@
-# coding=utf-8
-# Copyright 2020 Microsoft and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Fast Tokenization class for model DeBERTa."""
-
-import json
-from typing import List, Optional, Tuple
-
-from tokenizers import pre_tokenizers
-
-from ...tokenization_utils_base import AddedToken, BatchEncoding
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import logging
-from .tokenization_deberta import DebertaTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}
-
-
-class DebertaTokenizerFast(PreTrainedTokenizerFast):
- """
- Construct a "fast" DeBERTa tokenizer (backed by HuggingFace's *tokenizers* library). Based on byte-level
- Byte-Pair-Encoding.
-
- This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
- be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
- ```python
- >>> from transformers import DebertaTokenizerFast
-
- >>> tokenizer = DebertaTokenizerFast.from_pretrained("microsoft/deberta-base")
- >>> tokenizer("Hello world")["input_ids"]
- [1, 31414, 232, 2]
-
- >>> tokenizer(" Hello world")["input_ids"]
- [1, 20920, 232, 2]
- ```
-
- You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer, but since
- the model was not pretrained this way, it might yield a decrease in performance.
-
-
-
- When used with `is_split_into_words=True`, this tokenizer needs to be instantiated with `add_prefix_space=True`.
-
-
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`, *optional*):
- Path to the vocabulary file.
- merges_file (`str`, *optional*):
- Path to the merges file.
- tokenizer_file (`str`, *optional*):
- The path to a tokenizer file to use instead of the vocab file.
- errors (`str`, *optional*, defaults to `"replace"`):
- Paradigm to follow when decoding bytes to UTF-8. See
- [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
- bos_token (`str`, *optional*, defaults to `"[CLS]"`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `"[SEP]"`):
- The end of sequence token.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial space to the input. This allows to treat the leading word just as any
- other word. (Deberta tokenizer detect beginning of words by the preceding space).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
- slow_tokenizer_class = DebertaTokenizer
-
- def __init__(
- self,
- vocab_file=None,
- merges_file=None,
- tokenizer_file=None,
- errors="replace",
- bos_token="[CLS]",
- eos_token="[SEP]",
- sep_token="[SEP]",
- cls_token="[CLS]",
- unk_token="[UNK]",
- pad_token="[PAD]",
- mask_token="[MASK]",
- add_prefix_space=False,
- **kwargs,
- ):
- super().__init__(
- vocab_file,
- merges_file,
- tokenizer_file=tokenizer_file,
- errors=errors,
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- sep_token=sep_token,
- cls_token=cls_token,
- pad_token=pad_token,
- mask_token=mask_token,
- add_prefix_space=add_prefix_space,
- **kwargs,
- )
- self.add_bos_token = kwargs.pop("add_bos_token", False)
-
- pre_tok_state = json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__())
- if pre_tok_state.get("add_prefix_space", add_prefix_space) != add_prefix_space:
- pre_tok_class = getattr(pre_tokenizers, pre_tok_state.pop("type"))
- pre_tok_state["add_prefix_space"] = add_prefix_space
- self.backend_tokenizer.pre_tokenizer = pre_tok_class(**pre_tok_state)
-
- self.add_prefix_space = add_prefix_space
-
- @property
- def mask_token(self) -> str:
- """
- `str`: Mask token, to use when training a model with masked-language modeling. Log an error if used while not
- having been set.
-
- Deberta tokenizer has a special mask token to be used in the fill-mask pipeline. The mask token will greedily
- comprise the space before the *[MASK]*.
- """
- if self._mask_token is None:
- if self.verbose:
- logger.error("Using mask_token, but it is not set yet.")
- return None
- return str(self._mask_token)
-
- @mask_token.setter
- def mask_token(self, value):
- """
- Overriding the default behavior of the mask token to have it eat the space before it.
- """
- # Mask token behave like a normal word, i.e. include the space before it
- # So we set lstrip to True
- value = AddedToken(value, lstrip=True, rstrip=False) if isinstance(value, str) else value
- self._mask_token = value
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A DeBERTa sequence has the following format:
-
- - single sequence: [CLS] X [SEP]
- - pair of sequences: [CLS] A [SEP] B [SEP]
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A DeBERTa
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
-
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.gpt2.tokenization_gpt2_fast.GPT2TokenizerFast._batch_encode_plus
- def _batch_encode_plus(self, *args, **kwargs) -> BatchEncoding:
- is_split_into_words = kwargs.get("is_split_into_words", False)
- assert self.add_prefix_space or not is_split_into_words, (
- f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
- "to use it with pretokenized inputs."
- )
-
- return super()._batch_encode_plus(*args, **kwargs)
-
- # Copied from transformers.models.gpt2.tokenization_gpt2_fast.GPT2TokenizerFast._encode_plus
- def _encode_plus(self, *args, **kwargs) -> BatchEncoding:
- is_split_into_words = kwargs.get("is_split_into_words", False)
-
- assert self.add_prefix_space or not is_split_into_words, (
- f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
- "to use it with pretokenized inputs."
- )
-
- return super()._encode_plus(*args, **kwargs)
-
- # Copied from transformers.models.gpt2.tokenization_gpt2_fast.GPT2TokenizerFast.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
diff --git a/transformers/models/deberta_v2/__init__.py b/transformers/models/deberta_v2/__init__.py
deleted file mode 100644
index fb1b20a331fe11dfa687c7550685de296ebafbe0..0000000000000000000000000000000000000000
--- a/transformers/models/deberta_v2/__init__.py
+++ /dev/null
@@ -1,127 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_tf_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_deberta_v2": ["DEBERTA_V2_PRETRAINED_CONFIG_ARCHIVE_MAP", "DebertaV2Config", "DebertaV2OnnxConfig"],
- "tokenization_deberta_v2": ["DebertaV2Tokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_deberta_v2_fast"] = ["DebertaV2TokenizerFast"]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_deberta_v2"] = [
- "TF_DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFDebertaV2ForMaskedLM",
- "TFDebertaV2ForQuestionAnswering",
- "TFDebertaV2ForMultipleChoice",
- "TFDebertaV2ForSequenceClassification",
- "TFDebertaV2ForTokenClassification",
- "TFDebertaV2Model",
- "TFDebertaV2PreTrainedModel",
- ]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_deberta_v2"] = [
- "DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DebertaV2ForMaskedLM",
- "DebertaV2ForMultipleChoice",
- "DebertaV2ForQuestionAnswering",
- "DebertaV2ForSequenceClassification",
- "DebertaV2ForTokenClassification",
- "DebertaV2Model",
- "DebertaV2PreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_deberta_v2 import (
- DEBERTA_V2_PRETRAINED_CONFIG_ARCHIVE_MAP,
- DebertaV2Config,
- DebertaV2OnnxConfig,
- )
- from .tokenization_deberta_v2 import DebertaV2Tokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_deberta_v2_fast import DebertaV2TokenizerFast
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_deberta_v2 import (
- TF_DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFDebertaV2ForMaskedLM,
- TFDebertaV2ForMultipleChoice,
- TFDebertaV2ForQuestionAnswering,
- TFDebertaV2ForSequenceClassification,
- TFDebertaV2ForTokenClassification,
- TFDebertaV2Model,
- TFDebertaV2PreTrainedModel,
- )
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_deberta_v2 import (
- DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST,
- DebertaV2ForMaskedLM,
- DebertaV2ForMultipleChoice,
- DebertaV2ForQuestionAnswering,
- DebertaV2ForSequenceClassification,
- DebertaV2ForTokenClassification,
- DebertaV2Model,
- DebertaV2PreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deberta_v2/__pycache__/__init__.cpython-310.pyc b/transformers/models/deberta_v2/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 1e95219c18e45241d742b1ef32acceae4a6266ec..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta_v2/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta_v2/__pycache__/configuration_deberta_v2.cpython-310.pyc b/transformers/models/deberta_v2/__pycache__/configuration_deberta_v2.cpython-310.pyc
deleted file mode 100644
index 1f86674666314a135fdbf87cf1488869fa2ecad8..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta_v2/__pycache__/configuration_deberta_v2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta_v2/__pycache__/modeling_deberta_v2.cpython-310.pyc b/transformers/models/deberta_v2/__pycache__/modeling_deberta_v2.cpython-310.pyc
deleted file mode 100644
index c8b19c1e9dc86c0d8adbef5014bbec6035e8b9fb..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta_v2/__pycache__/modeling_deberta_v2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta_v2/__pycache__/modeling_tf_deberta_v2.cpython-310.pyc b/transformers/models/deberta_v2/__pycache__/modeling_tf_deberta_v2.cpython-310.pyc
deleted file mode 100644
index bb3ab1262fdcf6354ab7345de9226cacf944cfff..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta_v2/__pycache__/modeling_tf_deberta_v2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta_v2/__pycache__/tokenization_deberta_v2.cpython-310.pyc b/transformers/models/deberta_v2/__pycache__/tokenization_deberta_v2.cpython-310.pyc
deleted file mode 100644
index a48e2accd3e8c481494f5cb2840f4609ddd743eb..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta_v2/__pycache__/tokenization_deberta_v2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta_v2/__pycache__/tokenization_deberta_v2_fast.cpython-310.pyc b/transformers/models/deberta_v2/__pycache__/tokenization_deberta_v2_fast.cpython-310.pyc
deleted file mode 100644
index 6e2c3d0939831490e7dfc65d7a5d6e5276b0b243..0000000000000000000000000000000000000000
Binary files a/transformers/models/deberta_v2/__pycache__/tokenization_deberta_v2_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deberta_v2/configuration_deberta_v2.py b/transformers/models/deberta_v2/configuration_deberta_v2.py
deleted file mode 100644
index 25348849e2f240e2e284d501282b9cf9266b8b00..0000000000000000000000000000000000000000
--- a/transformers/models/deberta_v2/configuration_deberta_v2.py
+++ /dev/null
@@ -1,192 +0,0 @@
-# coding=utf-8
-# Copyright 2020, Microsoft and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DeBERTa-v2 model configuration"""
-from collections import OrderedDict
-from typing import TYPE_CHECKING, Any, Mapping, Optional, Union
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-if TYPE_CHECKING:
- from ... import FeatureExtractionMixin, PreTrainedTokenizerBase, TensorType
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DEBERTA_V2_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DebertaV2Config(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DebertaV2Model`]. It is used to instantiate a
- DeBERTa-v2 model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the DeBERTa
- [microsoft/deberta-v2-xlarge](https://huggingface.co/microsoft/deberta-v2-xlarge) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Arguments:
- vocab_size (`int`, *optional*, defaults to 128100):
- Vocabulary size of the DeBERTa-v2 model. Defines the number of different tokens that can be represented by
- the `inputs_ids` passed when calling [`DebertaV2Model`].
- hidden_size (`int`, *optional*, defaults to 1536):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 24):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 24):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 6144):
- Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"`, `"gelu"`, `"tanh"`, `"gelu_fast"`, `"mish"`, `"linear"`, `"sigmoid"` and `"gelu_new"`
- are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 0):
- The vocabulary size of the `token_type_ids` passed when calling [`DebertaModel`] or [`TFDebertaModel`].
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-7):
- The epsilon used by the layer normalization layers.
- relative_attention (`bool`, *optional*, defaults to `True`):
- Whether use relative position encoding.
- max_relative_positions (`int`, *optional*, defaults to -1):
- The range of relative positions `[-max_position_embeddings, max_position_embeddings]`. Use the same value
- as `max_position_embeddings`.
- pad_token_id (`int`, *optional*, defaults to 0):
- The value used to pad input_ids.
- position_biased_input (`bool`, *optional*, defaults to `True`):
- Whether add absolute position embedding to content embedding.
- pos_att_type (`List[str]`, *optional*):
- The type of relative position attention, it can be a combination of `["p2c", "c2p"]`, e.g. `["p2c"]`,
- `["p2c", "c2p"]`, `["p2c", "c2p"]`.
- layer_norm_eps (`float`, optional, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
-
- Example:
-
- ```python
- >>> from transformers import DebertaV2Config, DebertaV2Model
-
- >>> # Initializing a DeBERTa-v2 microsoft/deberta-v2-xlarge style configuration
- >>> configuration = DebertaV2Config()
-
- >>> # Initializing a model (with random weights) from the microsoft/deberta-v2-xlarge style configuration
- >>> model = DebertaV2Model(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "deberta-v2"
-
- def __init__(
- self,
- vocab_size=128100,
- hidden_size=1536,
- num_hidden_layers=24,
- num_attention_heads=24,
- intermediate_size=6144,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=0,
- initializer_range=0.02,
- layer_norm_eps=1e-7,
- relative_attention=False,
- max_relative_positions=-1,
- pad_token_id=0,
- position_biased_input=True,
- pos_att_type=None,
- pooler_dropout=0,
- pooler_hidden_act="gelu",
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.initializer_range = initializer_range
- self.relative_attention = relative_attention
- self.max_relative_positions = max_relative_positions
- self.pad_token_id = pad_token_id
- self.position_biased_input = position_biased_input
-
- # Backwards compatibility
- if isinstance(pos_att_type, str):
- pos_att_type = [x.strip() for x in pos_att_type.lower().split("|")]
-
- self.pos_att_type = pos_att_type
- self.vocab_size = vocab_size
- self.layer_norm_eps = layer_norm_eps
-
- self.pooler_hidden_size = kwargs.get("pooler_hidden_size", hidden_size)
- self.pooler_dropout = pooler_dropout
- self.pooler_hidden_act = pooler_hidden_act
-
-
-class DebertaV2OnnxConfig(OnnxConfig):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- if self.task == "multiple-choice":
- dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
- else:
- dynamic_axis = {0: "batch", 1: "sequence"}
- if self._config.type_vocab_size > 0:
- return OrderedDict(
- [("input_ids", dynamic_axis), ("attention_mask", dynamic_axis), ("token_type_ids", dynamic_axis)]
- )
- else:
- return OrderedDict([("input_ids", dynamic_axis), ("attention_mask", dynamic_axis)])
-
- @property
- def default_onnx_opset(self) -> int:
- return 12
-
- def generate_dummy_inputs(
- self,
- preprocessor: Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"],
- batch_size: int = -1,
- seq_length: int = -1,
- num_choices: int = -1,
- is_pair: bool = False,
- framework: Optional["TensorType"] = None,
- num_channels: int = 3,
- image_width: int = 40,
- image_height: int = 40,
- tokenizer: "PreTrainedTokenizerBase" = None,
- ) -> Mapping[str, Any]:
- dummy_inputs = super().generate_dummy_inputs(preprocessor=preprocessor, framework=framework)
- if self._config.type_vocab_size == 0 and "token_type_ids" in dummy_inputs:
- del dummy_inputs["token_type_ids"]
- return dummy_inputs
diff --git a/transformers/models/deberta_v2/modeling_deberta_v2.py b/transformers/models/deberta_v2/modeling_deberta_v2.py
deleted file mode 100644
index dfe18b0d4964af94a2790c0d76b56ed3caacaeff..0000000000000000000000000000000000000000
--- a/transformers/models/deberta_v2/modeling_deberta_v2.py
+++ /dev/null
@@ -1,1629 +0,0 @@
-# coding=utf-8
-# Copyright 2020 Microsoft and the Hugging Face Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DeBERTa-v2 model."""
-
-from collections.abc import Sequence
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutput,
- MaskedLMOutput,
- MultipleChoiceModelOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import softmax_backward_data
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_deberta_v2 import DebertaV2Config
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "DebertaV2Config"
-_CHECKPOINT_FOR_DOC = "microsoft/deberta-v2-xlarge"
-_QA_TARGET_START_INDEX = 2
-_QA_TARGET_END_INDEX = 9
-
-
-from ..deprecated._archive_maps import DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.deberta.modeling_deberta.ContextPooler
-class ContextPooler(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.pooler_hidden_size, config.pooler_hidden_size)
- self.dropout = StableDropout(config.pooler_dropout)
- self.config = config
-
- def forward(self, hidden_states):
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
-
- context_token = hidden_states[:, 0]
- context_token = self.dropout(context_token)
- pooled_output = self.dense(context_token)
- pooled_output = ACT2FN[self.config.pooler_hidden_act](pooled_output)
- return pooled_output
-
- @property
- def output_dim(self):
- return self.config.hidden_size
-
-
-# Copied from transformers.models.deberta.modeling_deberta.XSoftmax with deberta->deberta_v2
-class XSoftmax(torch.autograd.Function):
- """
- Masked Softmax which is optimized for saving memory
-
- Args:
- input (`torch.tensor`): The input tensor that will apply softmax.
- mask (`torch.IntTensor`):
- The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
- dim (int): The dimension that will apply softmax
-
- Example:
-
- ```python
- >>> import torch
- >>> from transformers.models.deberta_v2.modeling_deberta_v2 import XSoftmax
-
- >>> # Make a tensor
- >>> x = torch.randn([4, 20, 100])
-
- >>> # Create a mask
- >>> mask = (x > 0).int()
-
- >>> # Specify the dimension to apply softmax
- >>> dim = -1
-
- >>> y = XSoftmax.apply(x, mask, dim)
- ```"""
-
- @staticmethod
- def forward(self, input, mask, dim):
- self.dim = dim
- rmask = ~(mask.to(torch.bool))
-
- output = input.masked_fill(rmask, torch.tensor(torch.finfo(input.dtype).min))
- output = torch.softmax(output, self.dim)
- output.masked_fill_(rmask, 0)
- self.save_for_backward(output)
- return output
-
- @staticmethod
- def backward(self, grad_output):
- (output,) = self.saved_tensors
- inputGrad = softmax_backward_data(self, grad_output, output, self.dim, output)
- return inputGrad, None, None
-
- @staticmethod
- def symbolic(g, self, mask, dim):
- import torch.onnx.symbolic_helper as sym_help
- from torch.onnx.symbolic_opset9 import masked_fill, softmax
-
- mask_cast_value = g.op("Cast", mask, to_i=sym_help.cast_pytorch_to_onnx["Long"])
- r_mask = g.op(
- "Cast",
- g.op("Sub", g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)), mask_cast_value),
- to_i=sym_help.cast_pytorch_to_onnx["Bool"],
- )
- output = masked_fill(
- g, self, r_mask, g.op("Constant", value_t=torch.tensor(torch.finfo(self.type().dtype()).min))
- )
- output = softmax(g, output, dim)
- return masked_fill(g, output, r_mask, g.op("Constant", value_t=torch.tensor(0, dtype=torch.bool)))
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DropoutContext
-class DropoutContext(object):
- def __init__(self):
- self.dropout = 0
- self.mask = None
- self.scale = 1
- self.reuse_mask = True
-
-
-# Copied from transformers.models.deberta.modeling_deberta.get_mask
-def get_mask(input, local_context):
- if not isinstance(local_context, DropoutContext):
- dropout = local_context
- mask = None
- else:
- dropout = local_context.dropout
- dropout *= local_context.scale
- mask = local_context.mask if local_context.reuse_mask else None
-
- if dropout > 0 and mask is None:
- mask = (1 - torch.empty_like(input).bernoulli_(1 - dropout)).to(torch.bool)
-
- if isinstance(local_context, DropoutContext):
- if local_context.mask is None:
- local_context.mask = mask
-
- return mask, dropout
-
-
-# Copied from transformers.models.deberta.modeling_deberta.XDropout
-class XDropout(torch.autograd.Function):
- """Optimized dropout function to save computation and memory by using mask operation instead of multiplication."""
-
- @staticmethod
- def forward(ctx, input, local_ctx):
- mask, dropout = get_mask(input, local_ctx)
- ctx.scale = 1.0 / (1 - dropout)
- if dropout > 0:
- ctx.save_for_backward(mask)
- return input.masked_fill(mask, 0) * ctx.scale
- else:
- return input
-
- @staticmethod
- def backward(ctx, grad_output):
- if ctx.scale > 1:
- (mask,) = ctx.saved_tensors
- return grad_output.masked_fill(mask, 0) * ctx.scale, None
- else:
- return grad_output, None
-
- @staticmethod
- def symbolic(g: torch._C.Graph, input: torch._C.Value, local_ctx: Union[float, DropoutContext]) -> torch._C.Value:
- from torch.onnx import symbolic_opset12
-
- dropout_p = local_ctx
- if isinstance(local_ctx, DropoutContext):
- dropout_p = local_ctx.dropout
- # StableDropout only calls this function when training.
- train = True
- # TODO: We should check if the opset_version being used to export
- # is > 12 here, but there's no good way to do that. As-is, if the
- # opset_version < 12, export will fail with a CheckerError.
- # Once https://github.com/pytorch/pytorch/issues/78391 is fixed, do something like:
- # if opset_version < 12:
- # return torch.onnx.symbolic_opset9.dropout(g, input, dropout_p, train)
- return symbolic_opset12.dropout(g, input, dropout_p, train)
-
-
-# Copied from transformers.models.deberta.modeling_deberta.StableDropout
-class StableDropout(nn.Module):
- """
- Optimized dropout module for stabilizing the training
-
- Args:
- drop_prob (float): the dropout probabilities
- """
-
- def __init__(self, drop_prob):
- super().__init__()
- self.drop_prob = drop_prob
- self.count = 0
- self.context_stack = None
-
- def forward(self, x):
- """
- Call the module
-
- Args:
- x (`torch.tensor`): The input tensor to apply dropout
- """
- if self.training and self.drop_prob > 0:
- return XDropout.apply(x, self.get_context())
- return x
-
- def clear_context(self):
- self.count = 0
- self.context_stack = None
-
- def init_context(self, reuse_mask=True, scale=1):
- if self.context_stack is None:
- self.context_stack = []
- self.count = 0
- for c in self.context_stack:
- c.reuse_mask = reuse_mask
- c.scale = scale
-
- def get_context(self):
- if self.context_stack is not None:
- if self.count >= len(self.context_stack):
- self.context_stack.append(DropoutContext())
- ctx = self.context_stack[self.count]
- ctx.dropout = self.drop_prob
- self.count += 1
- return ctx
- else:
- return self.drop_prob
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DebertaSelfOutput with DebertaLayerNorm->LayerNorm
-class DebertaV2SelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
- self.dropout = StableDropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DebertaAttention with Deberta->DebertaV2
-class DebertaV2Attention(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.self = DisentangledSelfAttention(config)
- self.output = DebertaV2SelfOutput(config)
- self.config = config
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- output_attentions=False,
- query_states=None,
- relative_pos=None,
- rel_embeddings=None,
- ):
- self_output = self.self(
- hidden_states,
- attention_mask,
- output_attentions,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- )
- if output_attentions:
- self_output, att_matrix = self_output
- if query_states is None:
- query_states = hidden_states
- attention_output = self.output(self_output, query_states)
-
- if output_attentions:
- return (attention_output, att_matrix)
- else:
- return attention_output
-
-
-# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->DebertaV2
-class DebertaV2Intermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DebertaOutput with DebertaLayerNorm->LayerNorm
-class DebertaV2Output(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
- self.dropout = StableDropout(config.hidden_dropout_prob)
- self.config = config
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DebertaLayer with Deberta->DebertaV2
-class DebertaV2Layer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.attention = DebertaV2Attention(config)
- self.intermediate = DebertaV2Intermediate(config)
- self.output = DebertaV2Output(config)
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- query_states=None,
- relative_pos=None,
- rel_embeddings=None,
- output_attentions=False,
- ):
- attention_output = self.attention(
- hidden_states,
- attention_mask,
- output_attentions=output_attentions,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- )
- if output_attentions:
- attention_output, att_matrix = attention_output
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.output(intermediate_output, attention_output)
- if output_attentions:
- return (layer_output, att_matrix)
- else:
- return layer_output
-
-
-class ConvLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- kernel_size = getattr(config, "conv_kernel_size", 3)
- groups = getattr(config, "conv_groups", 1)
- self.conv_act = getattr(config, "conv_act", "tanh")
- self.conv = nn.Conv1d(
- config.hidden_size, config.hidden_size, kernel_size, padding=(kernel_size - 1) // 2, groups=groups
- )
- self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
- self.dropout = StableDropout(config.hidden_dropout_prob)
- self.config = config
-
- def forward(self, hidden_states, residual_states, input_mask):
- out = self.conv(hidden_states.permute(0, 2, 1).contiguous()).permute(0, 2, 1).contiguous()
- rmask = (1 - input_mask).bool()
- out.masked_fill_(rmask.unsqueeze(-1).expand(out.size()), 0)
- out = ACT2FN[self.conv_act](self.dropout(out))
-
- layer_norm_input = residual_states + out
- output = self.LayerNorm(layer_norm_input).to(layer_norm_input)
-
- if input_mask is None:
- output_states = output
- else:
- if input_mask.dim() != layer_norm_input.dim():
- if input_mask.dim() == 4:
- input_mask = input_mask.squeeze(1).squeeze(1)
- input_mask = input_mask.unsqueeze(2)
-
- input_mask = input_mask.to(output.dtype)
- output_states = output * input_mask
-
- return output_states
-
-
-class DebertaV2Encoder(nn.Module):
- """Modified BertEncoder with relative position bias support"""
-
- def __init__(self, config):
- super().__init__()
-
- self.layer = nn.ModuleList([DebertaV2Layer(config) for _ in range(config.num_hidden_layers)])
- self.relative_attention = getattr(config, "relative_attention", False)
-
- if self.relative_attention:
- self.max_relative_positions = getattr(config, "max_relative_positions", -1)
- if self.max_relative_positions < 1:
- self.max_relative_positions = config.max_position_embeddings
-
- self.position_buckets = getattr(config, "position_buckets", -1)
- pos_ebd_size = self.max_relative_positions * 2
-
- if self.position_buckets > 0:
- pos_ebd_size = self.position_buckets * 2
-
- self.rel_embeddings = nn.Embedding(pos_ebd_size, config.hidden_size)
-
- self.norm_rel_ebd = [x.strip() for x in getattr(config, "norm_rel_ebd", "none").lower().split("|")]
-
- if "layer_norm" in self.norm_rel_ebd:
- self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=True)
-
- self.conv = ConvLayer(config) if getattr(config, "conv_kernel_size", 0) > 0 else None
- self.gradient_checkpointing = False
-
- def get_rel_embedding(self):
- rel_embeddings = self.rel_embeddings.weight if self.relative_attention else None
- if rel_embeddings is not None and ("layer_norm" in self.norm_rel_ebd):
- rel_embeddings = self.LayerNorm(rel_embeddings)
- return rel_embeddings
-
- def get_attention_mask(self, attention_mask):
- if attention_mask.dim() <= 2:
- extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
- attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1)
- elif attention_mask.dim() == 3:
- attention_mask = attention_mask.unsqueeze(1)
-
- return attention_mask
-
- def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
- if self.relative_attention and relative_pos is None:
- q = query_states.size(-2) if query_states is not None else hidden_states.size(-2)
- relative_pos = build_relative_position(
- q,
- hidden_states.size(-2),
- bucket_size=self.position_buckets,
- max_position=self.max_relative_positions,
- device=hidden_states.device,
- )
- return relative_pos
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- output_hidden_states=True,
- output_attentions=False,
- query_states=None,
- relative_pos=None,
- return_dict=True,
- ):
- if attention_mask.dim() <= 2:
- input_mask = attention_mask
- else:
- input_mask = attention_mask.sum(-2) > 0
- attention_mask = self.get_attention_mask(attention_mask)
- relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
-
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- if isinstance(hidden_states, Sequence):
- next_kv = hidden_states[0]
- else:
- next_kv = hidden_states
- rel_embeddings = self.get_rel_embedding()
- output_states = next_kv
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (output_states,)
-
- if self.gradient_checkpointing and self.training:
- output_states = self._gradient_checkpointing_func(
- layer_module.__call__,
- next_kv,
- attention_mask,
- query_states,
- relative_pos,
- rel_embeddings,
- output_attentions,
- )
- else:
- output_states = layer_module(
- next_kv,
- attention_mask,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- output_attentions=output_attentions,
- )
-
- if output_attentions:
- output_states, att_m = output_states
-
- if i == 0 and self.conv is not None:
- output_states = self.conv(hidden_states, output_states, input_mask)
-
- if query_states is not None:
- query_states = output_states
- if isinstance(hidden_states, Sequence):
- next_kv = hidden_states[i + 1] if i + 1 < len(self.layer) else None
- else:
- next_kv = output_states
-
- if output_attentions:
- all_attentions = all_attentions + (att_m,)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (output_states,)
-
- if not return_dict:
- return tuple(v for v in [output_states, all_hidden_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=output_states, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
-
-def make_log_bucket_position(relative_pos, bucket_size, max_position):
- sign = torch.sign(relative_pos)
- mid = bucket_size // 2
- abs_pos = torch.where(
- (relative_pos < mid) & (relative_pos > -mid),
- torch.tensor(mid - 1).type_as(relative_pos),
- torch.abs(relative_pos),
- )
- log_pos = (
- torch.ceil(torch.log(abs_pos / mid) / torch.log(torch.tensor((max_position - 1) / mid)) * (mid - 1)) + mid
- )
- bucket_pos = torch.where(abs_pos <= mid, relative_pos.type_as(log_pos), log_pos * sign)
- return bucket_pos
-
-
-def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1, device=None):
- """
- Build relative position according to the query and key
-
- We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
- \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
- P_k\\)
-
- Args:
- query_size (int): the length of query
- key_size (int): the length of key
- bucket_size (int): the size of position bucket
- max_position (int): the maximum allowed absolute position
- device (`torch.device`): the device on which tensors will be created.
-
- Return:
- `torch.LongTensor`: A tensor with shape [1, query_size, key_size]
- """
-
- q_ids = torch.arange(0, query_size, device=device)
- k_ids = torch.arange(0, key_size, device=device)
- rel_pos_ids = q_ids[:, None] - k_ids[None, :]
- if bucket_size > 0 and max_position > 0:
- rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
- rel_pos_ids = rel_pos_ids.to(torch.long)
- rel_pos_ids = rel_pos_ids[:query_size, :]
- rel_pos_ids = rel_pos_ids.unsqueeze(0)
- return rel_pos_ids
-
-
-@torch.jit.script
-# Copied from transformers.models.deberta.modeling_deberta.c2p_dynamic_expand
-def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos):
- return c2p_pos.expand([query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)])
-
-
-@torch.jit.script
-# Copied from transformers.models.deberta.modeling_deberta.p2c_dynamic_expand
-def p2c_dynamic_expand(c2p_pos, query_layer, key_layer):
- return c2p_pos.expand([query_layer.size(0), query_layer.size(1), key_layer.size(-2), key_layer.size(-2)])
-
-
-@torch.jit.script
-# Copied from transformers.models.deberta.modeling_deberta.pos_dynamic_expand
-def pos_dynamic_expand(pos_index, p2c_att, key_layer):
- return pos_index.expand(p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2)))
-
-
-class DisentangledSelfAttention(nn.Module):
- """
- Disentangled self-attention module
-
- Parameters:
- config (`DebertaV2Config`):
- A model config class instance with the configuration to build a new model. The schema is similar to
- *BertConfig*, for more details, please refer [`DebertaV2Config`]
-
- """
-
- def __init__(self, config):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
- self.num_attention_heads = config.num_attention_heads
- _attention_head_size = config.hidden_size // config.num_attention_heads
- self.attention_head_size = getattr(config, "attention_head_size", _attention_head_size)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.query_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
- self.key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
- self.value_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
-
- self.share_att_key = getattr(config, "share_att_key", False)
- self.pos_att_type = config.pos_att_type if config.pos_att_type is not None else []
- self.relative_attention = getattr(config, "relative_attention", False)
-
- if self.relative_attention:
- self.position_buckets = getattr(config, "position_buckets", -1)
- self.max_relative_positions = getattr(config, "max_relative_positions", -1)
- if self.max_relative_positions < 1:
- self.max_relative_positions = config.max_position_embeddings
- self.pos_ebd_size = self.max_relative_positions
- if self.position_buckets > 0:
- self.pos_ebd_size = self.position_buckets
-
- self.pos_dropout = StableDropout(config.hidden_dropout_prob)
-
- if not self.share_att_key:
- if "c2p" in self.pos_att_type:
- self.pos_key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
- if "p2c" in self.pos_att_type:
- self.pos_query_proj = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = StableDropout(config.attention_probs_dropout_prob)
-
- def transpose_for_scores(self, x, attention_heads):
- new_x_shape = x.size()[:-1] + (attention_heads, -1)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3).contiguous().view(-1, x.size(1), x.size(-1))
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- output_attentions=False,
- query_states=None,
- relative_pos=None,
- rel_embeddings=None,
- ):
- """
- Call the module
-
- Args:
- hidden_states (`torch.FloatTensor`):
- Input states to the module usually the output from previous layer, it will be the Q,K and V in
- *Attention(Q,K,V)*
-
- attention_mask (`torch.BoolTensor`):
- An attention mask matrix of shape [*B*, *N*, *N*] where *B* is the batch size, *N* is the maximum
- sequence length in which element [i,j] = *1* means the *i* th token in the input can attend to the *j*
- th token.
-
- output_attentions (`bool`, optional):
- Whether return the attention matrix.
-
- query_states (`torch.FloatTensor`, optional):
- The *Q* state in *Attention(Q,K,V)*.
-
- relative_pos (`torch.LongTensor`):
- The relative position encoding between the tokens in the sequence. It's of shape [*B*, *N*, *N*] with
- values ranging in [*-max_relative_positions*, *max_relative_positions*].
-
- rel_embeddings (`torch.FloatTensor`):
- The embedding of relative distances. It's a tensor of shape [\\(2 \\times
- \\text{max_relative_positions}\\), *hidden_size*].
-
-
- """
- if query_states is None:
- query_states = hidden_states
- query_layer = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads)
- key_layer = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads)
- value_layer = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads)
-
- rel_att = None
- # Take the dot product between "query" and "key" to get the raw attention scores.
- scale_factor = 1
- if "c2p" in self.pos_att_type:
- scale_factor += 1
- if "p2c" in self.pos_att_type:
- scale_factor += 1
- scale = torch.sqrt(torch.tensor(query_layer.size(-1), dtype=torch.float) * scale_factor)
- attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2) / scale.to(dtype=query_layer.dtype))
- if self.relative_attention:
- rel_embeddings = self.pos_dropout(rel_embeddings)
- rel_att = self.disentangled_attention_bias(
- query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
- )
-
- if rel_att is not None:
- attention_scores = attention_scores + rel_att
- attention_scores = attention_scores
- attention_scores = attention_scores.view(
- -1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1)
- )
-
- # bsz x height x length x dimension
- attention_probs = XSoftmax.apply(attention_scores, attention_mask, -1)
- attention_probs = self.dropout(attention_probs)
- context_layer = torch.bmm(
- attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1)), value_layer
- )
- context_layer = (
- context_layer.view(-1, self.num_attention_heads, context_layer.size(-2), context_layer.size(-1))
- .permute(0, 2, 1, 3)
- .contiguous()
- )
- new_context_layer_shape = context_layer.size()[:-2] + (-1,)
- context_layer = context_layer.view(new_context_layer_shape)
- if output_attentions:
- return (context_layer, attention_probs)
- else:
- return context_layer
-
- def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
- if relative_pos is None:
- q = query_layer.size(-2)
- relative_pos = build_relative_position(
- q,
- key_layer.size(-2),
- bucket_size=self.position_buckets,
- max_position=self.max_relative_positions,
- device=query_layer.device,
- )
- if relative_pos.dim() == 2:
- relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
- elif relative_pos.dim() == 3:
- relative_pos = relative_pos.unsqueeze(1)
- # bsz x height x query x key
- elif relative_pos.dim() != 4:
- raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {relative_pos.dim()}")
-
- att_span = self.pos_ebd_size
- relative_pos = relative_pos.long().to(query_layer.device)
-
- rel_embeddings = rel_embeddings[0 : att_span * 2, :].unsqueeze(0)
- if self.share_att_key:
- pos_query_layer = self.transpose_for_scores(
- self.query_proj(rel_embeddings), self.num_attention_heads
- ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1)
- pos_key_layer = self.transpose_for_scores(self.key_proj(rel_embeddings), self.num_attention_heads).repeat(
- query_layer.size(0) // self.num_attention_heads, 1, 1
- )
- else:
- if "c2p" in self.pos_att_type:
- pos_key_layer = self.transpose_for_scores(
- self.pos_key_proj(rel_embeddings), self.num_attention_heads
- ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1) # .split(self.all_head_size, dim=-1)
- if "p2c" in self.pos_att_type:
- pos_query_layer = self.transpose_for_scores(
- self.pos_query_proj(rel_embeddings), self.num_attention_heads
- ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1) # .split(self.all_head_size, dim=-1)
-
- score = 0
- # content->position
- if "c2p" in self.pos_att_type:
- scale = torch.sqrt(torch.tensor(pos_key_layer.size(-1), dtype=torch.float) * scale_factor)
- c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2))
- c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
- c2p_att = torch.gather(
- c2p_att,
- dim=-1,
- index=c2p_pos.squeeze(0).expand([query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]),
- )
- score += c2p_att / scale.to(dtype=c2p_att.dtype)
-
- # position->content
- if "p2c" in self.pos_att_type:
- scale = torch.sqrt(torch.tensor(pos_query_layer.size(-1), dtype=torch.float) * scale_factor)
- if key_layer.size(-2) != query_layer.size(-2):
- r_pos = build_relative_position(
- key_layer.size(-2),
- key_layer.size(-2),
- bucket_size=self.position_buckets,
- max_position=self.max_relative_positions,
- device=query_layer.device,
- )
- r_pos = r_pos.unsqueeze(0)
- else:
- r_pos = relative_pos
-
- p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
- p2c_att = torch.bmm(key_layer, pos_query_layer.transpose(-1, -2))
- p2c_att = torch.gather(
- p2c_att,
- dim=-1,
- index=p2c_pos.squeeze(0).expand([query_layer.size(0), key_layer.size(-2), key_layer.size(-2)]),
- ).transpose(-1, -2)
- score += p2c_att / scale.to(dtype=p2c_att.dtype)
-
- return score
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DebertaEmbeddings with DebertaLayerNorm->LayerNorm
-class DebertaV2Embeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config):
- super().__init__()
- pad_token_id = getattr(config, "pad_token_id", 0)
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
- self.word_embeddings = nn.Embedding(config.vocab_size, self.embedding_size, padding_idx=pad_token_id)
-
- self.position_biased_input = getattr(config, "position_biased_input", True)
- if not self.position_biased_input:
- self.position_embeddings = None
- else:
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.embedding_size)
-
- if config.type_vocab_size > 0:
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, self.embedding_size)
-
- if self.embedding_size != config.hidden_size:
- self.embed_proj = nn.Linear(self.embedding_size, config.hidden_size, bias=False)
- self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
- self.dropout = StableDropout(config.hidden_dropout_prob)
- self.config = config
-
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
-
- def forward(self, input_ids=None, token_type_ids=None, position_ids=None, mask=None, inputs_embeds=None):
- if input_ids is not None:
- input_shape = input_ids.size()
- else:
- input_shape = inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, :seq_length]
-
- if token_type_ids is None:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
-
- if self.position_embeddings is not None:
- position_embeddings = self.position_embeddings(position_ids.long())
- else:
- position_embeddings = torch.zeros_like(inputs_embeds)
-
- embeddings = inputs_embeds
- if self.position_biased_input:
- embeddings += position_embeddings
- if self.config.type_vocab_size > 0:
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
- embeddings += token_type_embeddings
-
- if self.embedding_size != self.config.hidden_size:
- embeddings = self.embed_proj(embeddings)
-
- embeddings = self.LayerNorm(embeddings)
-
- if mask is not None:
- if mask.dim() != embeddings.dim():
- if mask.dim() == 4:
- mask = mask.squeeze(1).squeeze(1)
- mask = mask.unsqueeze(2)
- mask = mask.to(embeddings.dtype)
-
- embeddings = embeddings * mask
-
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DebertaPreTrainedModel with Deberta->DebertaV2
-class DebertaV2PreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DebertaV2Config
- base_model_prefix = "deberta"
- _keys_to_ignore_on_load_unexpected = ["position_embeddings"]
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
-
-DEBERTA_START_DOCSTRING = r"""
- The DeBERTa model was proposed in [DeBERTa: Decoding-enhanced BERT with Disentangled
- Attention](https://arxiv.org/abs/2006.03654) by Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen. It's build
- on top of BERT/RoBERTa with two improvements, i.e. disentangled attention and enhanced mask decoder. With those two
- improvements, it out perform BERT/RoBERTa on a majority of tasks with 80GB pretraining data.
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
-
- Parameters:
- config ([`DebertaV2Config`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DEBERTA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DeBERTa Model transformer outputting raw hidden-states without any specific head on top.",
- DEBERTA_START_DOCSTRING,
-)
-# Copied from transformers.models.deberta.modeling_deberta.DebertaModel with Deberta->DebertaV2
-class DebertaV2Model(DebertaV2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.embeddings = DebertaV2Embeddings(config)
- self.encoder = DebertaV2Encoder(config)
- self.z_steps = 0
- self.config = config
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, new_embeddings):
- self.embeddings.word_embeddings = new_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- raise NotImplementedError("The prune function is not implemented in DeBERTa model.")
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device)
- if token_type_ids is None:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- mask=attention_mask,
- inputs_embeds=inputs_embeds,
- )
-
- encoder_outputs = self.encoder(
- embedding_output,
- attention_mask,
- output_hidden_states=True,
- output_attentions=output_attentions,
- return_dict=return_dict,
- )
- encoded_layers = encoder_outputs[1]
-
- if self.z_steps > 1:
- hidden_states = encoded_layers[-2]
- layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
- query_states = encoded_layers[-1]
- rel_embeddings = self.encoder.get_rel_embedding()
- attention_mask = self.encoder.get_attention_mask(attention_mask)
- rel_pos = self.encoder.get_rel_pos(embedding_output)
- for layer in layers[1:]:
- query_states = layer(
- hidden_states,
- attention_mask,
- output_attentions=False,
- query_states=query_states,
- relative_pos=rel_pos,
- rel_embeddings=rel_embeddings,
- )
- encoded_layers.append(query_states)
-
- sequence_output = encoded_layers[-1]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[(1 if output_hidden_states else 2) :]
-
- return BaseModelOutput(
- last_hidden_state=sequence_output,
- hidden_states=encoder_outputs.hidden_states if output_hidden_states else None,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings("""DeBERTa Model with a `language modeling` head on top.""", DEBERTA_START_DOCSTRING)
-class DebertaV2ForMaskedLM(DebertaV2PreTrainedModel):
- _tied_weights_keys = ["cls.predictions.decoder.weight", "cls.predictions.decoder.bias"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.deberta = DebertaV2Model(config)
- self.cls = DebertaV2OnlyMLMHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.cls.predictions.decoder
-
- def set_output_embeddings(self, new_embeddings):
- self.cls.predictions.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="[MASK]",
- )
- # Copied from transformers.models.deberta.modeling_deberta.DebertaForMaskedLM.forward with Deberta->DebertaV2
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deberta(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- prediction_scores = self.cls(sequence_output)
-
- masked_lm_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss() # -100 index = padding token
- masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[1:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return MaskedLMOutput(
- loss=masked_lm_loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DebertaPredictionHeadTransform with Deberta->DebertaV2
-class DebertaV2PredictionHeadTransform(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
-
- self.dense = nn.Linear(config.hidden_size, self.embedding_size)
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = ACT2FN[config.hidden_act]
- else:
- self.transform_act_fn = config.hidden_act
- self.LayerNorm = nn.LayerNorm(self.embedding_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.deberta.modeling_deberta.DebertaLMPredictionHead with Deberta->DebertaV2
-class DebertaV2LMPredictionHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.transform = DebertaV2PredictionHeadTransform(config)
-
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
- # The output weights are the same as the input embeddings, but there is
- # an output-only bias for each token.
- self.decoder = nn.Linear(self.embedding_size, config.vocab_size, bias=False)
-
- self.bias = nn.Parameter(torch.zeros(config.vocab_size))
-
- # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
- self.decoder.bias = self.bias
-
- def forward(self, hidden_states):
- hidden_states = self.transform(hidden_states)
- hidden_states = self.decoder(hidden_states)
- return hidden_states
-
-
-# copied from transformers.models.bert.BertOnlyMLMHead with bert -> deberta
-class DebertaV2OnlyMLMHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.predictions = DebertaV2LMPredictionHead(config)
-
- def forward(self, sequence_output):
- prediction_scores = self.predictions(sequence_output)
- return prediction_scores
-
-
-@add_start_docstrings(
- """
- DeBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-class DebertaV2ForSequenceClassification(DebertaV2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- num_labels = getattr(config, "num_labels", 2)
- self.num_labels = num_labels
-
- self.deberta = DebertaV2Model(config)
- self.pooler = ContextPooler(config)
- output_dim = self.pooler.output_dim
-
- self.classifier = nn.Linear(output_dim, num_labels)
- drop_out = getattr(config, "cls_dropout", None)
- drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
- self.dropout = StableDropout(drop_out)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.deberta.get_input_embeddings()
-
- def set_input_embeddings(self, new_embeddings):
- self.deberta.set_input_embeddings(new_embeddings)
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- # Copied from transformers.models.deberta.modeling_deberta.DebertaForSequenceClassification.forward with Deberta->DebertaV2
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deberta(
- input_ids,
- token_type_ids=token_type_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- encoder_layer = outputs[0]
- pooled_output = self.pooler(encoder_layer)
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- # regression task
- loss_fn = nn.MSELoss()
- logits = logits.view(-1).to(labels.dtype)
- loss = loss_fn(logits, labels.view(-1))
- elif labels.dim() == 1 or labels.size(-1) == 1:
- label_index = (labels >= 0).nonzero()
- labels = labels.long()
- if label_index.size(0) > 0:
- labeled_logits = torch.gather(
- logits, 0, label_index.expand(label_index.size(0), logits.size(1))
- )
- labels = torch.gather(labels, 0, label_index.view(-1))
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(labeled_logits.view(-1, self.num_labels).float(), labels.view(-1))
- else:
- loss = torch.tensor(0).to(logits)
- else:
- log_softmax = nn.LogSoftmax(-1)
- loss = -((log_softmax(logits) * labels).sum(-1)).mean()
- elif self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-# Copied from transformers.models.deberta.modeling_deberta.DebertaForTokenClassification with Deberta->DebertaV2
-class DebertaV2ForTokenClassification(DebertaV2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.deberta = DebertaV2Model(config)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deberta(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- DEBERTA_START_DOCSTRING,
-)
-class DebertaV2ForQuestionAnswering(DebertaV2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.deberta = DebertaV2Model(config)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- qa_target_start_index=_QA_TARGET_START_INDEX,
- qa_target_end_index=_QA_TARGET_END_INDEX,
- )
- # Copied from transformers.models.deberta.modeling_deberta.DebertaForQuestionAnswering.forward with Deberta->DebertaV2
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deberta(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[1:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-class DebertaV2ForMultipleChoice(DebertaV2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- num_labels = getattr(config, "num_labels", 2)
- self.num_labels = num_labels
-
- self.deberta = DebertaV2Model(config)
- self.pooler = ContextPooler(config)
- output_dim = self.pooler.output_dim
-
- self.classifier = nn.Linear(output_dim, 1)
- drop_out = getattr(config, "cls_dropout", None)
- drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
- self.dropout = StableDropout(drop_out)
-
- self.init_weights()
-
- def get_input_embeddings(self):
- return self.deberta.get_input_embeddings()
-
- def set_input_embeddings(self, new_embeddings):
- self.deberta.set_input_embeddings(new_embeddings)
-
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
- flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
- flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- flat_inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.deberta(
- flat_input_ids,
- position_ids=flat_position_ids,
- token_type_ids=flat_token_type_ids,
- attention_mask=flat_attention_mask,
- inputs_embeds=flat_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- encoder_layer = outputs[0]
- pooled_output = self.pooler(encoder_layer)
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/deberta_v2/modeling_tf_deberta_v2.py b/transformers/models/deberta_v2/modeling_tf_deberta_v2.py
deleted file mode 100644
index 546e7f1a8d003857ad31d70728b4fc0bc1d99744..0000000000000000000000000000000000000000
--- a/transformers/models/deberta_v2/modeling_tf_deberta_v2.py
+++ /dev/null
@@ -1,1874 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Microsoft and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 DeBERTa-v2 model."""
-
-from __future__ import annotations
-
-from typing import Dict, Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFMaskedLMOutput,
- TFMultipleChoiceModelOutput,
- TFQuestionAnsweringModelOutput,
- TFSequenceClassifierOutput,
- TFTokenClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFMaskedLanguageModelingLoss,
- TFModelInputType,
- TFMultipleChoiceLoss,
- TFPreTrainedModel,
- TFQuestionAnsweringLoss,
- TFSequenceClassificationLoss,
- TFTokenClassificationLoss,
- get_initializer,
- keras,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_deberta_v2 import DebertaV2Config
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "DebertaV2Config"
-_CHECKPOINT_FOR_DOC = "kamalkraj/deberta-v2-xlarge"
-
-
-from ..deprecated._archive_maps import TF_DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaContextPooler with Deberta->DebertaV2
-class TFDebertaV2ContextPooler(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
- self.dense = keras.layers.Dense(config.pooler_hidden_size, name="dense")
- self.dropout = TFDebertaV2StableDropout(config.pooler_dropout, name="dropout")
- self.config = config
-
- def call(self, hidden_states, training: bool = False):
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- context_token = hidden_states[:, 0]
- context_token = self.dropout(context_token, training=training)
- pooled_output = self.dense(context_token)
- pooled_output = get_tf_activation(self.config.pooler_hidden_act)(pooled_output)
- return pooled_output
-
- @property
- def output_dim(self) -> int:
- return self.config.hidden_size
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.pooler_hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaXSoftmax with Deberta->DebertaV2
-class TFDebertaV2XSoftmax(keras.layers.Layer):
- """
- Masked Softmax which is optimized for saving memory
-
- Args:
- input (`tf.Tensor`): The input tensor that will apply softmax.
- mask (`tf.Tensor`): The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
- dim (int): The dimension that will apply softmax
- """
-
- def __init__(self, axis=-1, **kwargs):
- super().__init__(**kwargs)
- self.axis = axis
-
- def call(self, inputs: tf.Tensor, mask: tf.Tensor):
- rmask = tf.logical_not(tf.cast(mask, tf.bool))
- output = tf.where(rmask, float("-inf"), inputs)
- output = stable_softmax(output, self.axis)
- output = tf.where(rmask, 0.0, output)
- return output
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaStableDropout with Deberta->DebertaV2
-class TFDebertaV2StableDropout(keras.layers.Layer):
- """
- Optimized dropout module for stabilizing the training
-
- Args:
- drop_prob (float): the dropout probabilities
- """
-
- def __init__(self, drop_prob, **kwargs):
- super().__init__(**kwargs)
- self.drop_prob = drop_prob
-
- @tf.custom_gradient
- def xdropout(self, inputs):
- """
- Applies dropout to the inputs, as vanilla dropout, but also scales the remaining elements up by 1/drop_prob.
- """
- mask = tf.cast(
- 1
- - tf.compat.v1.distributions.Bernoulli(probs=1.0 - self.drop_prob).sample(sample_shape=shape_list(inputs)),
- tf.bool,
- )
- scale = tf.convert_to_tensor(1.0 / (1 - self.drop_prob), dtype=tf.float32)
- if self.drop_prob > 0:
- inputs = tf.where(mask, 0.0, inputs) * scale
-
- def grad(upstream):
- if self.drop_prob > 0:
- return tf.where(mask, 0.0, upstream) * scale
- else:
- return upstream
-
- return inputs, grad
-
- def call(self, inputs: tf.Tensor, training: tf.Tensor = False):
- if training:
- return self.xdropout(inputs)
- return inputs
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaSelfOutput with Deberta->DebertaV2
-class TFDebertaV2SelfOutput(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
- self.dense = keras.layers.Dense(config.hidden_size, name="dense")
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = TFDebertaV2StableDropout(config.hidden_dropout_prob, name="dropout")
- self.config = config
-
- def call(self, hidden_states, input_tensor, training: bool = False):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states, training=training)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaAttention with Deberta->DebertaV2
-class TFDebertaV2Attention(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
- self.self = TFDebertaV2DisentangledSelfAttention(config, name="self")
- self.dense_output = TFDebertaV2SelfOutput(config, name="output")
- self.config = config
-
- def call(
- self,
- input_tensor: tf.Tensor,
- attention_mask: tf.Tensor,
- query_states: tf.Tensor = None,
- relative_pos: tf.Tensor = None,
- rel_embeddings: tf.Tensor = None,
- output_attentions: bool = False,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- self_outputs = self.self(
- hidden_states=input_tensor,
- attention_mask=attention_mask,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- output_attentions=output_attentions,
- training=training,
- )
- if query_states is None:
- query_states = input_tensor
- attention_output = self.dense_output(
- hidden_states=self_outputs[0], input_tensor=query_states, training=training
- )
-
- output = (attention_output,) + self_outputs[1:]
-
- return output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "self", None) is not None:
- with tf.name_scope(self.self.name):
- self.self.build(None)
- if getattr(self, "dense_output", None) is not None:
- with tf.name_scope(self.dense_output.name):
- self.dense_output.build(None)
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaIntermediate with Deberta->DebertaV2
-class TFDebertaV2Intermediate(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
-
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.intermediate_act_fn = config.hidden_act
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaOutput with Deberta->DebertaV2
-class TFDebertaV2Output(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = TFDebertaV2StableDropout(config.hidden_dropout_prob, name="dropout")
- self.config = config
-
- def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.dropout(hidden_states, training=training)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.intermediate_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaLayer with Deberta->DebertaV2
-class TFDebertaV2Layer(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
-
- self.attention = TFDebertaV2Attention(config, name="attention")
- self.intermediate = TFDebertaV2Intermediate(config, name="intermediate")
- self.bert_output = TFDebertaV2Output(config, name="output")
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- query_states: tf.Tensor = None,
- relative_pos: tf.Tensor = None,
- rel_embeddings: tf.Tensor = None,
- output_attentions: bool = False,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- attention_outputs = self.attention(
- input_tensor=hidden_states,
- attention_mask=attention_mask,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- output_attentions=output_attentions,
- training=training,
- )
- attention_output = attention_outputs[0]
- intermediate_output = self.intermediate(hidden_states=attention_output)
- layer_output = self.bert_output(
- hidden_states=intermediate_output, input_tensor=attention_output, training=training
- )
- outputs = (layer_output,) + attention_outputs[1:] # add attentions if we output them
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "intermediate", None) is not None:
- with tf.name_scope(self.intermediate.name):
- self.intermediate.build(None)
- if getattr(self, "bert_output", None) is not None:
- with tf.name_scope(self.bert_output.name):
- self.bert_output.build(None)
-
-
-class TFDebertaV2ConvLayer(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
-
- self.kernel_size = getattr(config, "conv_kernel_size", 3)
- # groups = getattr(config, "conv_groups", 1)
- self.conv_act = get_tf_activation(getattr(config, "conv_act", "tanh"))
- self.padding = (self.kernel_size - 1) // 2
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = TFDebertaV2StableDropout(config.hidden_dropout_prob, name="dropout")
- self.config = config
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- with tf.name_scope("conv"):
- self.conv_kernel = self.add_weight(
- name="kernel",
- shape=[self.kernel_size, self.config.hidden_size, self.config.hidden_size],
- initializer=get_initializer(self.config.initializer_range),
- )
- self.conv_bias = self.add_weight(
- name="bias", shape=[self.config.hidden_size], initializer=tf.zeros_initializer()
- )
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
-
- def call(
- self, hidden_states: tf.Tensor, residual_states: tf.Tensor, input_mask: tf.Tensor, training: bool = False
- ) -> tf.Tensor:
- out = tf.nn.conv2d(
- tf.expand_dims(hidden_states, 1),
- tf.expand_dims(self.conv_kernel, 0),
- strides=1,
- padding=[[0, 0], [0, 0], [self.padding, self.padding], [0, 0]],
- )
- out = tf.squeeze(tf.nn.bias_add(out, self.conv_bias), 1)
- rmask = tf.cast(1 - input_mask, tf.bool)
- out = tf.where(tf.broadcast_to(tf.expand_dims(rmask, -1), shape_list(out)), 0.0, out)
- out = self.dropout(out, training=training)
- out = self.conv_act(out)
-
- layer_norm_input = residual_states + out
- output = self.LayerNorm(layer_norm_input)
-
- if input_mask is None:
- output_states = output
- else:
- if len(shape_list(input_mask)) != len(shape_list(layer_norm_input)):
- if len(shape_list(input_mask)) == 4:
- input_mask = tf.squeeze(tf.squeeze(input_mask, axis=1), axis=1)
- input_mask = tf.cast(tf.expand_dims(input_mask, axis=2), tf.float32)
-
- output_states = output * input_mask
-
- return output_states
-
-
-class TFDebertaV2Encoder(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
-
- self.layer = [TFDebertaV2Layer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]
- self.relative_attention = getattr(config, "relative_attention", False)
- self.config = config
- if self.relative_attention:
- self.max_relative_positions = getattr(config, "max_relative_positions", -1)
- if self.max_relative_positions < 1:
- self.max_relative_positions = config.max_position_embeddings
-
- self.position_buckets = getattr(config, "position_buckets", -1)
- self.pos_ebd_size = self.max_relative_positions * 2
-
- if self.position_buckets > 0:
- self.pos_ebd_size = self.position_buckets * 2
-
- self.norm_rel_ebd = [x.strip() for x in getattr(config, "norm_rel_ebd", "none").lower().split("|")]
-
- if "layer_norm" in self.norm_rel_ebd:
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
-
- self.conv = TFDebertaV2ConvLayer(config, name="conv") if getattr(config, "conv_kernel_size", 0) > 0 else None
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if self.relative_attention:
- self.rel_embeddings = self.add_weight(
- name="rel_embeddings.weight",
- shape=[self.pos_ebd_size, self.config.hidden_size],
- initializer=get_initializer(self.config.initializer_range),
- )
- if getattr(self, "conv", None) is not None:
- with tf.name_scope(self.conv.name):
- self.conv.build(None)
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, self.config.hidden_size])
- if getattr(self, "layer", None) is not None:
- for layer in self.layer:
- with tf.name_scope(layer.name):
- layer.build(None)
-
- def get_rel_embedding(self):
- rel_embeddings = self.rel_embeddings if self.relative_attention else None
- if rel_embeddings is not None and ("layer_norm" in self.norm_rel_ebd):
- rel_embeddings = self.LayerNorm(rel_embeddings)
- return rel_embeddings
-
- def get_attention_mask(self, attention_mask):
- if len(shape_list(attention_mask)) <= 2:
- extended_attention_mask = tf.expand_dims(tf.expand_dims(attention_mask, 1), 2)
- attention_mask = extended_attention_mask * tf.expand_dims(tf.squeeze(extended_attention_mask, -2), -1)
- attention_mask = tf.cast(attention_mask, tf.uint8)
- elif len(shape_list(attention_mask)) == 3:
- attention_mask = tf.expand_dims(attention_mask, 1)
-
- return attention_mask
-
- def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
- if self.relative_attention and relative_pos is None:
- q = shape_list(query_states)[-2] if query_states is not None else shape_list(hidden_states)[-2]
- relative_pos = build_relative_position(
- q,
- shape_list(hidden_states)[-2],
- bucket_size=self.position_buckets,
- max_position=self.max_relative_positions,
- )
- return relative_pos
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- query_states: tf.Tensor = None,
- relative_pos: tf.Tensor = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- training: bool = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- if len(shape_list(attention_mask)) <= 2:
- input_mask = attention_mask
- else:
- input_mask = tf.cast(tf.math.reduce_sum(attention_mask, axis=-2) > 0, dtype=tf.uint8)
-
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- attention_mask = self.get_attention_mask(attention_mask)
- relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
-
- next_kv = hidden_states
-
- rel_embeddings = self.get_rel_embedding()
- output_states = next_kv
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (output_states,)
-
- layer_outputs = layer_module(
- hidden_states=next_kv,
- attention_mask=attention_mask,
- query_states=query_states,
- relative_pos=relative_pos,
- rel_embeddings=rel_embeddings,
- output_attentions=output_attentions,
- training=training,
- )
- output_states = layer_outputs[0]
-
- if i == 0 and self.conv is not None:
- output_states = self.conv(hidden_states, output_states, input_mask)
-
- next_kv = output_states
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- # Add last layer
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (output_states,)
-
- if not return_dict:
- return tuple(v for v in [output_states, all_hidden_states, all_attentions] if v is not None)
-
- return TFBaseModelOutput(
- last_hidden_state=output_states, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
-
-def make_log_bucket_position(relative_pos, bucket_size, max_position):
- sign = tf.math.sign(relative_pos)
- mid = bucket_size // 2
- abs_pos = tf.where((relative_pos < mid) & (relative_pos > -mid), mid - 1, tf.math.abs(relative_pos))
- log_pos = (
- tf.math.ceil(
- tf.cast(tf.math.log(abs_pos / mid), tf.float32) / tf.math.log((max_position - 1) / mid) * (mid - 1)
- )
- + mid
- )
- bucket_pos = tf.cast(
- tf.where(abs_pos <= mid, tf.cast(relative_pos, tf.float32), log_pos * tf.cast(sign, tf.float32)), tf.int32
- )
- return bucket_pos
-
-
-def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1):
- """
- Build relative position according to the query and key
-
- We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
- \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
- P_k\\)
-
- Args:
- query_size (int): the length of query
- key_size (int): the length of key
- bucket_size (int): the size of position bucket
- max_position (int): the maximum allowed absolute position
-
- Return:
- `tf.Tensor`: A tensor with shape [1, query_size, key_size]
-
- """
- q_ids = tf.range(query_size, dtype=tf.int32)
- k_ids = tf.range(key_size, dtype=tf.int32)
- rel_pos_ids = q_ids[:, None] - tf.tile(tf.expand_dims(k_ids, axis=0), [shape_list(q_ids)[0], 1])
- if bucket_size > 0 and max_position > 0:
- rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
- rel_pos_ids = rel_pos_ids[:query_size, :]
- rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)
- return tf.cast(rel_pos_ids, tf.int64)
-
-
-def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos):
- shapes = [
- shape_list(query_layer)[0],
- shape_list(query_layer)[1],
- shape_list(query_layer)[2],
- shape_list(relative_pos)[-1],
- ]
- return tf.broadcast_to(c2p_pos, shapes)
-
-
-def p2c_dynamic_expand(c2p_pos, query_layer, key_layer):
- shapes = [
- shape_list(query_layer)[0],
- shape_list(query_layer)[1],
- shape_list(key_layer)[-2],
- shape_list(key_layer)[-2],
- ]
- return tf.broadcast_to(c2p_pos, shapes)
-
-
-def pos_dynamic_expand(pos_index, p2c_att, key_layer):
- shapes = shape_list(p2c_att)[:2] + [shape_list(pos_index)[-2], shape_list(key_layer)[-2]]
- return tf.broadcast_to(pos_index, shapes)
-
-
-def take_along_axis(x, indices):
- # Only a valid port of np.take_along_axis when the gather axis is -1
-
- # TPU + gathers and reshapes don't go along well -- see https://github.com/huggingface/transformers/issues/18239
- if isinstance(tf.distribute.get_strategy(), tf.distribute.TPUStrategy):
- # [B, S, P] -> [B, S, P, D]
- one_hot_indices = tf.one_hot(indices, depth=x.shape[-1], dtype=x.dtype)
-
- # if we ignore the first two dims, this is equivalent to multiplying a matrix (one hot) by a vector (x)
- # grossly abusing notation: [B, S, P, D] . [B, S, D] = [B, S, P]
- gathered = tf.einsum("ijkl,ijl->ijk", one_hot_indices, x)
-
- # GPUs, on the other hand, prefer gathers instead of large one-hot+matmuls
- else:
- gathered = tf.gather(x, indices, batch_dims=2)
-
- return gathered
-
-
-class TFDebertaV2DisentangledSelfAttention(keras.layers.Layer):
- """
- Disentangled self-attention module
-
- Parameters:
- config (`DebertaV2Config`):
- A model config class instance with the configuration to build a new model. The schema is similar to
- *BertConfig*, for more details, please refer [`DebertaV2Config`]
-
- """
-
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
- if config.hidden_size % config.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
- self.num_attention_heads = config.num_attention_heads
- _attention_head_size = config.hidden_size // config.num_attention_heads
- self.attention_head_size = getattr(config, "attention_head_size", _attention_head_size)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.query_proj = keras.layers.Dense(
- self.all_head_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="query_proj",
- use_bias=True,
- )
- self.key_proj = keras.layers.Dense(
- self.all_head_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="key_proj",
- use_bias=True,
- )
- self.value_proj = keras.layers.Dense(
- self.all_head_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="value_proj",
- use_bias=True,
- )
-
- self.share_att_key = getattr(config, "share_att_key", False)
- self.pos_att_type = config.pos_att_type if config.pos_att_type is not None else []
- self.relative_attention = getattr(config, "relative_attention", False)
-
- if self.relative_attention:
- self.position_buckets = getattr(config, "position_buckets", -1)
- self.max_relative_positions = getattr(config, "max_relative_positions", -1)
- if self.max_relative_positions < 1:
- self.max_relative_positions = config.max_position_embeddings
- self.pos_ebd_size = self.max_relative_positions
- if self.position_buckets > 0:
- self.pos_ebd_size = self.position_buckets
-
- self.pos_dropout = TFDebertaV2StableDropout(config.hidden_dropout_prob, name="pos_dropout")
-
- if not self.share_att_key:
- if "c2p" in self.pos_att_type:
- self.pos_key_proj = keras.layers.Dense(
- self.all_head_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="pos_proj",
- use_bias=True,
- )
- if "p2c" in self.pos_att_type:
- self.pos_query_proj = keras.layers.Dense(
- self.all_head_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="pos_q_proj",
- )
- self.softmax = TFDebertaV2XSoftmax(axis=-1)
- self.dropout = TFDebertaV2StableDropout(config.attention_probs_dropout_prob, name="dropout")
- self.config = config
-
- def transpose_for_scores(self, tensor: tf.Tensor, attention_heads: int) -> tf.Tensor:
- tensor_shape = shape_list(tensor)
- # In graph mode mode, we can't reshape with -1 as the final dimension if the first dimension (batch size) is None
- shape = tensor_shape[:-1] + [attention_heads, tensor_shape[-1] // attention_heads]
- # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]
- tensor = tf.reshape(tensor=tensor, shape=shape)
- tensor = tf.transpose(tensor, perm=[0, 2, 1, 3])
- x_shape = shape_list(tensor)
- tensor = tf.reshape(tensor, shape=[-1, x_shape[-2], x_shape[-1]])
- return tensor
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- query_states: tf.Tensor = None,
- relative_pos: tf.Tensor = None,
- rel_embeddings: tf.Tensor = None,
- output_attentions: bool = False,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- """
- Call the module
-
- Args:
- hidden_states (`tf.Tensor`):
- Input states to the module usually the output from previous layer, it will be the Q,K and V in
- *Attention(Q,K,V)*
-
- attention_mask (`tf.Tensor`):
- An attention mask matrix of shape [*B*, *N*, *N*] where *B* is the batch size, *N* is the maximum
- sequence length in which element [i,j] = *1* means the *i* th token in the input can attend to the *j*
- th token.
-
- return_att (`bool`, optional):
- Whether return the attention matrix.
-
- query_states (`tf.Tensor`, optional):
- The *Q* state in *Attention(Q,K,V)*.
-
- relative_pos (`tf.Tensor`):
- The relative position encoding between the tokens in the sequence. It's of shape [*B*, *N*, *N*] with
- values ranging in [*-max_relative_positions*, *max_relative_positions*].
-
- rel_embeddings (`tf.Tensor`):
- The embedding of relative distances. It's a tensor of shape [\\(2 \\times
- \\text{max_relative_positions}\\), *hidden_size*].
-
-
- """
- if query_states is None:
- query_states = hidden_states
- query_layer = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads)
- key_layer = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads)
- value_layer = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads)
-
- rel_att = None
- # Take the dot product between "query" and "key" to get the raw attention scores.
- scale_factor = 1
- if "c2p" in self.pos_att_type:
- scale_factor += 1
- if "p2c" in self.pos_att_type:
- scale_factor += 1
- scale = tf.math.sqrt(tf.cast(shape_list(query_layer)[-1] * scale_factor, tf.float32))
- attention_scores = tf.matmul(query_layer, tf.transpose(key_layer, [0, 2, 1]) / scale)
- if self.relative_attention:
- rel_embeddings = self.pos_dropout(rel_embeddings)
- rel_att = self.disentangled_att_bias(query_layer, key_layer, relative_pos, rel_embeddings, scale_factor)
-
- if rel_att is not None:
- attention_scores = attention_scores + rel_att
- attention_scores = tf.reshape(
- attention_scores,
- (-1, self.num_attention_heads, shape_list(attention_scores)[-2], shape_list(attention_scores)[-1]),
- )
-
- # bsz x height x length x dimension
- attention_probs = self.softmax(attention_scores, attention_mask)
- attention_probs = self.dropout(attention_probs, training=training)
- context_layer = tf.matmul(
- tf.reshape(attention_probs, [-1, shape_list(attention_probs)[-2], shape_list(attention_probs)[-1]]),
- value_layer,
- )
- context_layer = tf.transpose(
- tf.reshape(
- context_layer,
- [-1, self.num_attention_heads, shape_list(context_layer)[-2], shape_list(context_layer)[-1]],
- ),
- [0, 2, 1, 3],
- )
- # Set the final dimension here explicitly.
- # Calling tf.reshape(context_layer, (*context_layer_shape[:-2], -1)) raises an error when executing
- # the model in graph mode as context_layer is reshaped to (None, 7, None) and Dense layer in TFDebertaV2SelfOutput
- # requires final input dimension to be defined
- context_layer_shape = shape_list(context_layer)
- new_context_layer_shape = context_layer_shape[:-2] + [context_layer_shape[-2] * context_layer_shape[-1]]
- context_layer = tf.reshape(context_layer, new_context_layer_shape)
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
- return outputs
-
- def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
- if relative_pos is None:
- q = shape_list(query_layer)[-2]
- relative_pos = build_relative_position(
- q,
- shape_list(key_layer)[-2],
- bucket_size=self.position_buckets,
- max_position=self.max_relative_positions,
- )
- shape_list_pos = shape_list(relative_pos)
- if len(shape_list_pos) == 2:
- relative_pos = tf.expand_dims(tf.expand_dims(relative_pos, 0), 0)
- elif len(shape_list_pos) == 3:
- relative_pos = tf.expand_dims(relative_pos, 1)
- # bsz x height x query x key
- elif len(shape_list_pos) != 4:
- raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {len(shape_list_pos)}")
-
- att_span = self.pos_ebd_size
- rel_embeddings = tf.expand_dims(
- rel_embeddings[self.pos_ebd_size - att_span : self.pos_ebd_size + att_span, :], 0
- )
- if self.share_att_key:
- pos_query_layer = tf.tile(
- self.transpose_for_scores(self.query_proj(rel_embeddings), self.num_attention_heads),
- [shape_list(query_layer)[0] // self.num_attention_heads, 1, 1],
- )
- pos_key_layer = tf.tile(
- self.transpose_for_scores(self.key_proj(rel_embeddings), self.num_attention_heads),
- [shape_list(query_layer)[0] // self.num_attention_heads, 1, 1],
- )
- else:
- if "c2p" in self.pos_att_type:
- pos_key_layer = tf.tile(
- self.transpose_for_scores(self.pos_key_proj(rel_embeddings), self.num_attention_heads),
- [shape_list(query_layer)[0] // self.num_attention_heads, 1, 1],
- ) # .split(self.all_head_size, dim=-1)
- if "p2c" in self.pos_att_type:
- pos_query_layer = tf.tile(
- self.transpose_for_scores(self.pos_query_proj(rel_embeddings), self.num_attention_heads),
- [shape_list(query_layer)[0] // self.num_attention_heads, 1, 1],
- ) # .split(self.all_head_size, dim=-1)
-
- score = 0
- # content->position
- if "c2p" in self.pos_att_type:
- scale = tf.math.sqrt(tf.cast(shape_list(pos_key_layer)[-1] * scale_factor, tf.float32))
- c2p_att = tf.matmul(query_layer, tf.transpose(pos_key_layer, [0, 2, 1]))
- c2p_pos = tf.clip_by_value(relative_pos + att_span, 0, att_span * 2 - 1)
- c2p_att = take_along_axis(
- c2p_att,
- tf.broadcast_to(
- tf.squeeze(c2p_pos, 0),
- [shape_list(query_layer)[0], shape_list(query_layer)[1], shape_list(relative_pos)[-1]],
- ),
- )
- score += c2p_att / scale
-
- # position->content
- if "p2c" in self.pos_att_type:
- scale = tf.math.sqrt(tf.cast(shape_list(pos_query_layer)[-1] * scale_factor, tf.float32))
- if shape_list(key_layer)[-2] != shape_list(query_layer)[-2]:
- r_pos = build_relative_position(
- shape_list(key_layer)[-2],
- shape_list(key_layer)[-2],
- bucket_size=self.position_buckets,
- max_position=self.max_relative_positions,
- )
- r_pos = tf.expand_dims(r_pos, 0)
- else:
- r_pos = relative_pos
-
- p2c_pos = tf.clip_by_value(-r_pos + att_span, 0, att_span * 2 - 1)
-
- p2c_att = tf.matmul(key_layer, tf.transpose(pos_query_layer, [0, 2, 1]))
- p2c_att = tf.transpose(
- take_along_axis(
- p2c_att,
- tf.broadcast_to(
- tf.squeeze(p2c_pos, 0),
- [shape_list(query_layer)[0], shape_list(key_layer)[-2], shape_list(key_layer)[-2]],
- ),
- ),
- [0, 2, 1],
- )
- score += p2c_att / scale
-
- return score
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "query_proj", None) is not None:
- with tf.name_scope(self.query_proj.name):
- self.query_proj.build([None, None, self.config.hidden_size])
- if getattr(self, "key_proj", None) is not None:
- with tf.name_scope(self.key_proj.name):
- self.key_proj.build([None, None, self.config.hidden_size])
- if getattr(self, "value_proj", None) is not None:
- with tf.name_scope(self.value_proj.name):
- self.value_proj.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
- if getattr(self, "pos_dropout", None) is not None:
- with tf.name_scope(self.pos_dropout.name):
- self.pos_dropout.build(None)
- if getattr(self, "pos_key_proj", None) is not None:
- with tf.name_scope(self.pos_key_proj.name):
- self.pos_key_proj.build([None, None, self.config.hidden_size])
- if getattr(self, "pos_query_proj", None) is not None:
- with tf.name_scope(self.pos_query_proj.name):
- self.pos_query_proj.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaEmbeddings Deberta->DebertaV2
-class TFDebertaV2Embeddings(keras.layers.Layer):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
- self.hidden_size = config.hidden_size
- self.max_position_embeddings = config.max_position_embeddings
- self.position_biased_input = getattr(config, "position_biased_input", True)
- self.initializer_range = config.initializer_range
- if self.embedding_size != config.hidden_size:
- self.embed_proj = keras.layers.Dense(
- config.hidden_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="embed_proj",
- use_bias=False,
- )
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = TFDebertaV2StableDropout(config.hidden_dropout_prob, name="dropout")
-
- def build(self, input_shape=None):
- with tf.name_scope("word_embeddings"):
- self.weight = self.add_weight(
- name="weight",
- shape=[self.config.vocab_size, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
-
- with tf.name_scope("token_type_embeddings"):
- if self.config.type_vocab_size > 0:
- self.token_type_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.config.type_vocab_size, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
- else:
- self.token_type_embeddings = None
-
- with tf.name_scope("position_embeddings"):
- if self.position_biased_input:
- self.position_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.max_position_embeddings, self.hidden_size],
- initializer=get_initializer(self.initializer_range),
- )
- else:
- self.position_embeddings = None
-
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
- if getattr(self, "embed_proj", None) is not None:
- with tf.name_scope(self.embed_proj.name):
- self.embed_proj.build([None, None, self.embedding_size])
-
- def call(
- self,
- input_ids: tf.Tensor = None,
- position_ids: tf.Tensor = None,
- token_type_ids: tf.Tensor = None,
- inputs_embeds: tf.Tensor = None,
- mask: tf.Tensor = None,
- training: bool = False,
- ) -> tf.Tensor:
- """
- Applies embedding based on inputs tensor.
-
- Returns:
- final_embeddings (`tf.Tensor`): output embedding tensor.
- """
- if input_ids is None and inputs_embeds is None:
- raise ValueError("Need to provide either `input_ids` or `input_embeds`.")
-
- if input_ids is not None:
- check_embeddings_within_bounds(input_ids, self.config.vocab_size)
- inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
-
- input_shape = shape_list(inputs_embeds)[:-1]
-
- if token_type_ids is None:
- token_type_ids = tf.fill(dims=input_shape, value=0)
-
- if position_ids is None:
- position_ids = tf.expand_dims(tf.range(start=0, limit=input_shape[-1]), axis=0)
-
- final_embeddings = inputs_embeds
- if self.position_biased_input:
- position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
- final_embeddings += position_embeds
- if self.config.type_vocab_size > 0:
- token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)
- final_embeddings += token_type_embeds
-
- if self.embedding_size != self.hidden_size:
- final_embeddings = self.embed_proj(final_embeddings)
-
- final_embeddings = self.LayerNorm(final_embeddings)
-
- if mask is not None:
- if len(shape_list(mask)) != len(shape_list(final_embeddings)):
- if len(shape_list(mask)) == 4:
- mask = tf.squeeze(tf.squeeze(mask, axis=1), axis=1)
- mask = tf.cast(tf.expand_dims(mask, axis=2), tf.float32)
-
- final_embeddings = final_embeddings * mask
-
- final_embeddings = self.dropout(final_embeddings, training=training)
-
- return final_embeddings
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaPredictionHeadTransform with Deberta->DebertaV2
-class TFDebertaV2PredictionHeadTransform(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
-
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
-
- self.dense = keras.layers.Dense(
- units=self.embedding_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="dense",
- )
-
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.transform_act_fn = config.hidden_act
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.embedding_size])
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaLMPredictionHead with Deberta->DebertaV2
-class TFDebertaV2LMPredictionHead(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, input_embeddings: keras.layers.Layer, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
-
- self.transform = TFDebertaV2PredictionHeadTransform(config, name="transform")
-
- # The output weights are the same as the input embeddings, but there is
- # an output-only bias for each token.
- self.input_embeddings = input_embeddings
-
- def build(self, input_shape=None):
- self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")
-
- if self.built:
- return
- self.built = True
- if getattr(self, "transform", None) is not None:
- with tf.name_scope(self.transform.name):
- self.transform.build(None)
-
- def get_output_embeddings(self) -> keras.layers.Layer:
- return self.input_embeddings
-
- def set_output_embeddings(self, value: tf.Variable):
- self.input_embeddings.weight = value
- self.input_embeddings.vocab_size = shape_list(value)[0]
-
- def get_bias(self) -> Dict[str, tf.Variable]:
- return {"bias": self.bias}
-
- def set_bias(self, value: tf.Variable):
- self.bias = value["bias"]
- self.config.vocab_size = shape_list(value["bias"])[0]
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.transform(hidden_states=hidden_states)
- seq_length = shape_list(hidden_states)[1]
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.embedding_size])
- hidden_states = tf.matmul(a=hidden_states, b=self.input_embeddings.weight, transpose_b=True)
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])
- hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)
-
- return hidden_states
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaOnlyMLMHead with Deberta->DebertaV2
-class TFDebertaV2OnlyMLMHead(keras.layers.Layer):
- def __init__(self, config: DebertaV2Config, input_embeddings: keras.layers.Layer, **kwargs):
- super().__init__(**kwargs)
- self.predictions = TFDebertaV2LMPredictionHead(config, input_embeddings, name="predictions")
-
- def call(self, sequence_output: tf.Tensor) -> tf.Tensor:
- prediction_scores = self.predictions(hidden_states=sequence_output)
-
- return prediction_scores
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "predictions", None) is not None:
- with tf.name_scope(self.predictions.name):
- self.predictions.build(None)
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaMainLayer with Deberta->DebertaV2
-class TFDebertaV2MainLayer(keras.layers.Layer):
- config_class = DebertaV2Config
-
- def __init__(self, config: DebertaV2Config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
-
- self.embeddings = TFDebertaV2Embeddings(config, name="embeddings")
- self.encoder = TFDebertaV2Encoder(config, name="encoder")
-
- def get_input_embeddings(self) -> keras.layers.Layer:
- return self.embeddings
-
- def set_input_embeddings(self, value: tf.Variable):
- self.embeddings.weight = value
- self.embeddings.vocab_size = shape_list(value)[0]
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- raise NotImplementedError
-
- @unpack_inputs
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = tf.fill(dims=input_shape, value=1)
-
- if token_type_ids is None:
- token_type_ids = tf.fill(dims=input_shape, value=0)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- mask=attention_mask,
- training=training,
- )
-
- encoder_outputs = self.encoder(
- hidden_states=embedding_output,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = encoder_outputs[0]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[1:]
-
- return TFBaseModelOutput(
- last_hidden_state=sequence_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
-
-
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaPreTrainedModel with Deberta->DebertaV2
-class TFDebertaV2PreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DebertaV2Config
- base_model_prefix = "deberta"
-
-
-DEBERTA_START_DOCSTRING = r"""
- The DeBERTa model was proposed in [DeBERTa: Decoding-enhanced BERT with Disentangled
- Attention](https://arxiv.org/abs/2006.03654) by Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen. It's build
- on top of BERT/RoBERTa with two improvements, i.e. disentangled attention and enhanced mask decoder. With those two
- improvements, it out perform BERT/RoBERTa on a majority of tasks with 80GB pretraining data.
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`DebertaV2Config`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DEBERTA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`np.ndarray`, `tf.Tensor`, `List[tf.Tensor]` ``Dict[str, tf.Tensor]` or `Dict[str, np.ndarray]` and each example must have the shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- inputs_embeds (`np.ndarray` or `tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput``] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DeBERTa Model transformer outputting raw hidden-states without any specific head on top.",
- DEBERTA_START_DOCSTRING,
-)
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaModel with Deberta->DebertaV2
-class TFDebertaV2Model(TFDebertaV2PreTrainedModel):
- def __init__(self, config: DebertaV2Config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.deberta = TFDebertaV2MainLayer(config, name="deberta")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
-
-
-@add_start_docstrings("""DeBERTa Model with a `language modeling` head on top.""", DEBERTA_START_DOCSTRING)
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaForMaskedLM with Deberta->DebertaV2
-class TFDebertaV2ForMaskedLM(TFDebertaV2PreTrainedModel, TFMaskedLanguageModelingLoss):
- def __init__(self, config: DebertaV2Config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- if config.is_decoder:
- logger.warning(
- "If you want to use `TFDebertaV2ForMaskedLM` make sure `config.is_decoder=False` for "
- "bi-directional self-attention."
- )
-
- self.deberta = TFDebertaV2MainLayer(config, name="deberta")
- self.mlm = TFDebertaV2OnlyMLMHead(config, input_embeddings=self.deberta.embeddings, name="cls")
-
- def get_lm_head(self) -> keras.layers.Layer:
- return self.mlm.predictions
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFMaskedLMOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- prediction_scores = self.mlm(sequence_output=sequence_output, training=training)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=prediction_scores)
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFMaskedLMOutput(
- loss=loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "mlm", None) is not None:
- with tf.name_scope(self.mlm.name):
- self.mlm.build(None)
-
-
-@add_start_docstrings(
- """
- DeBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaForSequenceClassification with Deberta->DebertaV2
-class TFDebertaV2ForSequenceClassification(TFDebertaV2PreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: DebertaV2Config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
-
- self.deberta = TFDebertaV2MainLayer(config, name="deberta")
- self.pooler = TFDebertaV2ContextPooler(config, name="pooler")
-
- drop_out = getattr(config, "cls_dropout", None)
- drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
- self.dropout = TFDebertaV2StableDropout(drop_out, name="cls_dropout")
- self.classifier = keras.layers.Dense(
- units=config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- name="classifier",
- )
- self.output_dim = self.pooler.output_dim
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- pooled_output = self.pooler(sequence_output, training=training)
- pooled_output = self.dropout(pooled_output, training=training)
- logits = self.classifier(pooled_output)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "pooler", None) is not None:
- with tf.name_scope(self.pooler.name):
- self.pooler.build(None)
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.output_dim])
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaForTokenClassification with Deberta->DebertaV2
-class TFDebertaV2ForTokenClassification(TFDebertaV2PreTrainedModel, TFTokenClassificationLoss):
- def __init__(self, config: DebertaV2Config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
-
- self.deberta = TFDebertaV2MainLayer(config, name="deberta")
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.classifier = keras.layers.Dense(
- units=config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFTokenClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- sequence_output = self.dropout(sequence_output, training=training)
- logits = self.classifier(inputs=sequence_output)
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFTokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- DEBERTA_START_DOCSTRING,
-)
-# Copied from transformers.models.deberta.modeling_tf_deberta.TFDebertaForQuestionAnswering with Deberta->DebertaV2
-class TFDebertaV2ForQuestionAnswering(TFDebertaV2PreTrainedModel, TFQuestionAnsweringLoss):
- def __init__(self, config: DebertaV2Config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
-
- self.deberta = TFDebertaV2MainLayer(config, name="deberta")
- self.qa_outputs = keras.layers.Dense(
- units=config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFQuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- start_positions: np.ndarray | tf.Tensor | None = None,
- end_positions: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFQuestionAnsweringModelOutput, Tuple[tf.Tensor]]:
- r"""
- start_positions (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- outputs = self.deberta(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- logits = self.qa_outputs(inputs=sequence_output)
- start_logits, end_logits = tf.split(value=logits, num_or_size_splits=2, axis=-1)
- start_logits = tf.squeeze(input=start_logits, axis=-1)
- end_logits = tf.squeeze(input=end_logits, axis=-1)
- loss = None
-
- if start_positions is not None and end_positions is not None:
- labels = {"start_position": start_positions}
- labels["end_position"] = end_positions
- loss = self.hf_compute_loss(labels=labels, logits=(start_logits, end_logits))
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFQuestionAnsweringModelOutput(
- loss=loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- DeBERTa Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- DEBERTA_START_DOCSTRING,
-)
-class TFDebertaV2ForMultipleChoice(TFDebertaV2PreTrainedModel, TFMultipleChoiceLoss):
- # names with a '.' represents the authorized unexpected/missing layers when a TF model is loaded from a PT model
- # _keys_to_ignore_on_load_unexpected = [r"mlm___cls", r"nsp___cls", r"cls.predictions", r"cls.seq_relationship"]
- # _keys_to_ignore_on_load_missing = [r"dropout"]
-
- def __init__(self, config: DebertaV2Config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.deberta = TFDebertaV2MainLayer(config, name="deberta")
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.pooler = TFDebertaV2ContextPooler(config, name="pooler")
- self.classifier = keras.layers.Dense(
- units=1, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.output_dim = self.pooler.output_dim
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFMultipleChoiceModelOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]`
- where `num_choices` is the size of the second dimension of the input tensors. (See `input_ids` above)
- """
- if input_ids is not None:
- num_choices = shape_list(input_ids)[1]
- seq_length = shape_list(input_ids)[2]
- else:
- num_choices = shape_list(inputs_embeds)[1]
- seq_length = shape_list(inputs_embeds)[2]
-
- flat_input_ids = tf.reshape(tensor=input_ids, shape=(-1, seq_length)) if input_ids is not None else None
- flat_attention_mask = (
- tf.reshape(tensor=attention_mask, shape=(-1, seq_length)) if attention_mask is not None else None
- )
- flat_token_type_ids = (
- tf.reshape(tensor=token_type_ids, shape=(-1, seq_length)) if token_type_ids is not None else None
- )
- flat_position_ids = (
- tf.reshape(tensor=position_ids, shape=(-1, seq_length)) if position_ids is not None else None
- )
- flat_inputs_embeds = (
- tf.reshape(tensor=inputs_embeds, shape=(-1, seq_length, shape_list(inputs_embeds)[3]))
- if inputs_embeds is not None
- else None
- )
- outputs = self.deberta(
- input_ids=flat_input_ids,
- attention_mask=flat_attention_mask,
- token_type_ids=flat_token_type_ids,
- position_ids=flat_position_ids,
- inputs_embeds=flat_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- pooled_output = self.pooler(sequence_output, training=training)
- pooled_output = self.dropout(pooled_output, training=training)
- logits = self.classifier(pooled_output)
- reshaped_logits = tf.reshape(tensor=logits, shape=(-1, num_choices))
- loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=reshaped_logits)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFMultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deberta", None) is not None:
- with tf.name_scope(self.deberta.name):
- self.deberta.build(None)
- if getattr(self, "pooler", None) is not None:
- with tf.name_scope(self.pooler.name):
- self.pooler.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.output_dim])
diff --git a/transformers/models/deberta_v2/tokenization_deberta_v2.py b/transformers/models/deberta_v2/tokenization_deberta_v2.py
deleted file mode 100644
index a92103945416d78c24d7f77b541eac9c4d6b62e6..0000000000000000000000000000000000000000
--- a/transformers/models/deberta_v2/tokenization_deberta_v2.py
+++ /dev/null
@@ -1,521 +0,0 @@
-# coding=utf-8
-# Copyright 2020 Microsoft and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Tokenization class for model DeBERTa."""
-
-import os
-import unicodedata
-from typing import Any, Dict, List, Optional, Tuple
-
-import sentencepiece as sp
-
-from ...tokenization_utils import AddedToken, PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-VOCAB_FILES_NAMES = {"vocab_file": "spm.model"}
-
-
-class DebertaV2Tokenizer(PreTrainedTokenizer):
- r"""
- Constructs a DeBERTa-v2 tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
-
- Args:
- vocab_file (`str`):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- do_lower_case (`bool`, *optional*, defaults to `False`):
- Whether or not to lowercase the input when tokenizing.
- bos_token (`string`, *optional*, defaults to `"[CLS]"`):
- The beginning of sequence token that was used during pre-training. Can be used a sequence classifier token.
- When building a sequence using special tokens, this is not the token that is used for the beginning of
- sequence. The token used is the `cls_token`.
- eos_token (`string`, *optional*, defaults to `"[SEP]"`):
- The end of sequence token. When building a sequence using special tokens, this is not the token that is
- used for the end of sequence. The token used is the `sep_token`.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- sp_model_kwargs (`dict`, *optional*):
- Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
- SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
- to set:
-
- - `enable_sampling`: Enable subword regularization.
- - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
-
- - `nbest_size = {0,1}`: No sampling is performed.
- - `nbest_size > 1`: samples from the nbest_size results.
- - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
- using forward-filtering-and-backward-sampling algorithm.
-
- - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
- BPE-dropout.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
-
- def __init__(
- self,
- vocab_file,
- do_lower_case=False,
- split_by_punct=False,
- bos_token="[CLS]",
- eos_token="[SEP]",
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- sp_model_kwargs: Optional[Dict[str, Any]] = None,
- **kwargs,
- ) -> None:
- self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
-
- if not os.path.isfile(vocab_file):
- raise ValueError(
- f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
- " model use `tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- self.do_lower_case = do_lower_case
- self.split_by_punct = split_by_punct
- self.vocab_file = vocab_file
- self._tokenizer = SPMTokenizer(
- vocab_file, None, split_by_punct=split_by_punct, sp_model_kwargs=self.sp_model_kwargs
- )
- unk_token = AddedToken(unk_token, normalized=True, special=True) if isinstance(unk_token, str) else unk_token
- super().__init__(
- do_lower_case=do_lower_case,
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- split_by_punct=split_by_punct,
- sp_model_kwargs=self.sp_model_kwargs,
- **kwargs,
- )
- self._tokenizer.special_tokens = self.all_special_tokens
-
- @property
- def vocab_size(self):
- return len(self.vocab)
-
- @property
- def vocab(self):
- return self._tokenizer.vocab
-
- def get_vocab(self):
- vocab = self.vocab.copy()
- vocab.update(self.get_added_vocab())
- return vocab
-
- def _tokenize(self, text: str) -> List[str]:
- """Take as input a string and return a list of strings (tokens) for words/sub-words"""
- if self.do_lower_case:
- text = text.lower()
- return self._tokenizer.tokenize(text)
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self._tokenizer.spm.PieceToId(token)
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self._tokenizer.spm.IdToPiece(index) if index < self.vocab_size else self.unk_token
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- return self._tokenizer.decode(tokens)
-
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A DeBERTa sequence has the following format:
-
- - single sequence: [CLS] X [SEP]
- - pair of sequences: [CLS] A [SEP] B [SEP]
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
-
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
- """
- Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A DeBERTa
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
- add_prefix_space = kwargs.pop("add_prefix_space", False)
- if is_split_into_words or add_prefix_space:
- text = " " + text
- return (text, kwargs)
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- return self._tokenizer.save_pretrained(save_directory, filename_prefix=filename_prefix)
-
-
-class SPMTokenizer:
- r"""
- Constructs a tokenizer based on [SentencePiece](https://github.com/google/sentencepiece).
-
- Args:
- vocab_file (`str`):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- sp_model_kwargs (`dict`, *optional*):
- Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
- SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
- to set:
-
- - `enable_sampling`: Enable subword regularization.
- - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
-
- - `nbest_size = {0,1}`: No sampling is performed.
- - `nbest_size > 1`: samples from the nbest_size results.
- - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
- using forward-filtering-and-backward-sampling algorithm.
-
- - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
- BPE-dropout.
- """
-
- def __init__(
- self, vocab_file, special_tokens, split_by_punct=False, sp_model_kwargs: Optional[Dict[str, Any]] = None
- ):
- self.split_by_punct = split_by_punct
- self.vocab_file = vocab_file
- self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
- spm = sp.SentencePieceProcessor(**self.sp_model_kwargs)
- if not os.path.exists(vocab_file):
- raise FileNotFoundError(f"{vocab_file} does not exist!")
- spm.load(vocab_file)
- bpe_vocab_size = spm.GetPieceSize()
- # Token map
- # 0+1
- # 1+1
- # 2+1
- self.vocab = {spm.IdToPiece(i): i for i in range(bpe_vocab_size)}
- self.ids_to_tokens = [spm.IdToPiece(i) for i in range(bpe_vocab_size)]
- # self.vocab['[PAD]'] = 0
- # self.vocab['[CLS]'] = 1
- # self.vocab['[SEP]'] = 2
- # self.vocab['[UNK]'] = 3
-
- self.spm = spm
- self.special_tokens = special_tokens
-
- def __getstate__(self):
- state = self.__dict__.copy()
- state["spm"] = None
- return state
-
- def __setstate__(self, d):
- self.__dict__ = d
-
- # for backward compatibility
- if not hasattr(self, "sp_model_kwargs"):
- self.sp_model_kwargs = {}
-
- self.spm = sp.SentencePieceProcessor(**self.sp_model_kwargs)
- self.spm.Load(self.vocab_file)
-
- def tokenize(self, text):
- return self._encode_as_pieces(text)
-
- def convert_ids_to_tokens(self, ids):
- tokens = []
- for i in ids:
- tokens.append(self.ids_to_tokens[i])
- return tokens
-
- def decode(self, tokens, start=-1, end=-1, raw_text=None):
- if raw_text is None:
- current_sub_tokens = []
- out_string = ""
- prev_is_special = False
- for token in tokens:
- # make sure that special tokens are not decoded using sentencepiece model
- if token in self.special_tokens:
- if not prev_is_special:
- out_string += " "
- out_string += self.spm.decode_pieces(current_sub_tokens) + token
- prev_is_special = True
- current_sub_tokens = []
- else:
- current_sub_tokens.append(token)
- prev_is_special = False
- out_string += self.spm.decode_pieces(current_sub_tokens)
- return out_string.strip()
- else:
- words = self.split_to_words(raw_text)
- word_tokens = [self.tokenize(w) for w in words]
- token2words = [0] * len(tokens)
- tid = 0
- for i, w in enumerate(word_tokens):
- for k, t in enumerate(w):
- token2words[tid] = i
- tid += 1
- word_start = token2words[start]
- word_end = token2words[end] if end < len(tokens) else len(words)
- text = "".join(words[word_start:word_end])
- return text
-
- # TODO add a deprecation cycle as this can have different behaviour from our API
- def add_special_token(self, token):
- if token not in self.special_tokens:
- self.special_tokens.append(token)
- if token not in self.vocab:
- self.vocab[token] = len(self.vocab) - 1
- self.ids_to_tokens.append(token)
- return self.id(token)
-
- def part_of_whole_word(self, token, is_bos=False):
- logger.warning_once(
- "The `DebertaTokenizer.part_of_whole_word` method is deprecated and will be removed in `transformers==4.35`"
- )
- if is_bos:
- return True
- if (
- len(token) == 1
- and (_is_whitespace(list(token)[0]) or _is_control(list(token)[0]) or _is_punctuation(list(token)[0]))
- ) or token in self.special_tokens:
- return False
-
- word_start = b"\xe2\x96\x81".decode("utf-8")
- return not token.startswith(word_start)
-
- def pad(self):
- return "[PAD]"
-
- def bos(self):
- return "[CLS]"
-
- def eos(self):
- return "[SEP]"
-
- def unk(self):
- return "[UNK]"
-
- def mask(self):
- return "[MASK]"
-
- def sym(self, id):
- return self.ids_to_tokens[id]
-
- def id(self, sym):
- logger.warning_once(
- "The `DebertaTokenizer.id` method is deprecated and will be removed in `transformers==4.35`"
- )
- return self.vocab[sym] if sym in self.vocab else 1
-
- def _encode_as_pieces(self, text):
- text = convert_to_unicode(text)
- if self.split_by_punct:
- words = self._run_split_on_punc(text)
- pieces = [self.spm.encode(w, out_type=str) for w in words]
- return [p for w in pieces for p in w]
- else:
- return self.spm.encode(text, out_type=str)
-
- def split_to_words(self, text):
- pieces = self._encode_as_pieces(text)
- word_start = b"\xe2\x96\x81".decode("utf-8")
- words = []
- offset = 0
- prev_end = 0
- for i, p in enumerate(pieces):
- if p.startswith(word_start):
- if offset > prev_end:
- words.append(text[prev_end:offset])
- prev_end = offset
- w = p.replace(word_start, "")
- else:
- w = p
- try:
- s = text.index(w, offset)
- pn = ""
- k = i + 1
- while k < len(pieces):
- pn = pieces[k].replace(word_start, "")
- if len(pn) > 0:
- break
- k += 1
-
- if len(pn) > 0 and pn in text[offset:s]:
- offset = offset + 1
- else:
- offset = s + len(w)
- except Exception:
- offset = offset + 1
-
- if prev_end < offset:
- words.append(text[prev_end:offset])
-
- return words
-
- def _run_split_on_punc(self, text):
- """Splits punctuation on a piece of text."""
- chars = list(text)
- i = 0
- start_new_word = True
- output = []
- while i < len(chars):
- char = chars[i]
- if _is_punctuation(char):
- output.append([char])
- start_new_word = True
- else:
- if start_new_word:
- output.append([])
- start_new_word = False
- output[-1].append(char)
- i += 1
-
- return ["".join(x) for x in output]
-
- def save_pretrained(self, path: str, filename_prefix: str = None):
- filename = VOCAB_FILES_NAMES[list(VOCAB_FILES_NAMES.keys())[0]]
- if filename_prefix is not None:
- filename = filename_prefix + "-" + filename
- full_path = os.path.join(path, filename)
- with open(full_path, "wb") as fs:
- fs.write(self.spm.serialized_model_proto())
- return (full_path,)
-
-
-def _is_whitespace(char):
- """Checks whether `chars` is a whitespace character."""
- # \t, \n, and \r are technically control characters but we treat them
- # as whitespace since they are generally considered as such.
- if char == " " or char == "\t" or char == "\n" or char == "\r":
- return True
- cat = unicodedata.category(char)
- if cat == "Zs":
- return True
- return False
-
-
-def _is_control(char):
- """Checks whether `chars` is a control character."""
- # These are technically control characters but we count them as whitespace
- # characters.
- if char == "\t" or char == "\n" or char == "\r":
- return False
- cat = unicodedata.category(char)
- if cat.startswith("C"):
- return True
- return False
-
-
-def _is_punctuation(char):
- """Checks whether `chars` is a punctuation character."""
- cp = ord(char)
- # We treat all non-letter/number ASCII as punctuation.
- # Characters such as "^", "$", and "`" are not in the Unicode
- # Punctuation class but we treat them as punctuation anyways, for
- # consistency.
- if (cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126):
- return True
- cat = unicodedata.category(char)
- if cat.startswith("P"):
- return True
- return False
-
-
-def convert_to_unicode(text):
- """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
- if isinstance(text, str):
- return text
- elif isinstance(text, bytes):
- return text.decode("utf-8", "ignore")
- else:
- raise ValueError(f"Unsupported string type: {type(text)}")
diff --git a/transformers/models/deberta_v2/tokenization_deberta_v2_fast.py b/transformers/models/deberta_v2/tokenization_deberta_v2_fast.py
deleted file mode 100644
index cb92a61edf1afbde7c21f4be7130bb649ef3a8ab..0000000000000000000000000000000000000000
--- a/transformers/models/deberta_v2/tokenization_deberta_v2_fast.py
+++ /dev/null
@@ -1,220 +0,0 @@
-# coding=utf-8
-# Copyright 2020 Microsoft and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Fast Tokenization class for model DeBERTa."""
-
-import os
-from shutil import copyfile
-from typing import Optional, Tuple
-
-from ...file_utils import is_sentencepiece_available
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import logging
-
-
-if is_sentencepiece_available():
- from .tokenization_deberta_v2 import DebertaV2Tokenizer
-else:
- DebertaV2Tokenizer = None
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "spm.model", "tokenizer_file": "tokenizer.json"}
-
-
-class DebertaV2TokenizerFast(PreTrainedTokenizerFast):
- r"""
- Constructs a DeBERTa-v2 fast tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
-
- Args:
- vocab_file (`str`):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- do_lower_case (`bool`, *optional*, defaults to `False`):
- Whether or not to lowercase the input when tokenizing.
- bos_token (`string`, *optional*, defaults to `"[CLS]"`):
- The beginning of sequence token that was used during pre-training. Can be used a sequence classifier token.
- When building a sequence using special tokens, this is not the token that is used for the beginning of
- sequence. The token used is the `cls_token`.
- eos_token (`string`, *optional*, defaults to `"[SEP]"`):
- The end of sequence token. When building a sequence using special tokens, this is not the token that is
- used for the end of sequence. The token used is the `sep_token`.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- sp_model_kwargs (`dict`, *optional*):
- Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
- SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
- to set:
-
- - `enable_sampling`: Enable subword regularization.
- - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
-
- - `nbest_size = {0,1}`: No sampling is performed.
- - `nbest_size > 1`: samples from the nbest_size results.
- - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
- using forward-filtering-and-backward-sampling algorithm.
-
- - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
- BPE-dropout.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = DebertaV2Tokenizer
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- do_lower_case=False,
- split_by_punct=False,
- bos_token="[CLS]",
- eos_token="[SEP]",
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- **kwargs,
- ) -> None:
- super().__init__(
- vocab_file,
- tokenizer_file=tokenizer_file,
- do_lower_case=do_lower_case,
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- split_by_punct=split_by_punct,
- **kwargs,
- )
-
- self.do_lower_case = do_lower_case
- self.split_by_punct = split_by_punct
- self.vocab_file = vocab_file
-
- @property
- def can_save_slow_tokenizer(self) -> bool:
- return os.path.isfile(self.vocab_file) if self.vocab_file else False
-
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A DeBERTa sequence has the following format:
-
- - single sequence: [CLS] X [SEP]
- - pair of sequences: [CLS] A [SEP] B [SEP]
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
-
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
- """
- Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A DeBERTa
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not self.can_save_slow_tokenizer:
- raise ValueError(
- "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
- "tokenizer."
- )
-
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
-
- return (out_vocab_file,)
diff --git a/transformers/models/decision_transformer/__init__.py b/transformers/models/decision_transformer/__init__.py
deleted file mode 100644
index 44070229aaa8591cb967a4ca7ff4867873072f8a..0000000000000000000000000000000000000000
--- a/transformers/models/decision_transformer/__init__.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {
- "configuration_decision_transformer": [
- "DECISION_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "DecisionTransformerConfig",
- ],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_decision_transformer"] = [
- "DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DecisionTransformerGPT2Model",
- "DecisionTransformerGPT2PreTrainedModel",
- "DecisionTransformerModel",
- "DecisionTransformerPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_decision_transformer import (
- DECISION_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
- DecisionTransformerConfig,
- )
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_decision_transformer import (
- DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
- DecisionTransformerGPT2Model,
- DecisionTransformerGPT2PreTrainedModel,
- DecisionTransformerModel,
- DecisionTransformerPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/decision_transformer/__pycache__/__init__.cpython-310.pyc b/transformers/models/decision_transformer/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 8207a229fe9131d4c04ead0e321ddd3d71ac24d5..0000000000000000000000000000000000000000
Binary files a/transformers/models/decision_transformer/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/decision_transformer/__pycache__/configuration_decision_transformer.cpython-310.pyc b/transformers/models/decision_transformer/__pycache__/configuration_decision_transformer.cpython-310.pyc
deleted file mode 100644
index 21a28437ffc6fa414c2da43480c86e8117798cba..0000000000000000000000000000000000000000
Binary files a/transformers/models/decision_transformer/__pycache__/configuration_decision_transformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/decision_transformer/__pycache__/modeling_decision_transformer.cpython-310.pyc b/transformers/models/decision_transformer/__pycache__/modeling_decision_transformer.cpython-310.pyc
deleted file mode 100644
index 8f420cb6a86eb01431df4d88705c0302a73e5c7a..0000000000000000000000000000000000000000
Binary files a/transformers/models/decision_transformer/__pycache__/modeling_decision_transformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/decision_transformer/configuration_decision_transformer.py b/transformers/models/decision_transformer/configuration_decision_transformer.py
deleted file mode 100644
index d2c1914bee06eec71dc15fae565b45c0d673dbe7..0000000000000000000000000000000000000000
--- a/transformers/models/decision_transformer/configuration_decision_transformer.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Team and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Decision Transformer model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DECISION_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DecisionTransformerConfig(PretrainedConfig):
- """
- This is the configuration class to store the configuration of a [`DecisionTransformerModel`]. It is used to
- instantiate a Decision Transformer model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the standard
- DecisionTransformer architecture. Many of the config options are used to instatiate the GPT2 model that is used as
- part of the architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- state_dim (`int`, *optional*, defaults to 17):
- The state size for the RL environment
- act_dim (`int`, *optional*, defaults to 4):
- The size of the output action space
- hidden_size (`int`, *optional*, defaults to 128):
- The size of the hidden layers
- max_ep_len (`int`, *optional*, defaults to 4096):
- The maximum length of an episode in the environment
- action_tanh (`bool`, *optional*, defaults to True):
- Whether to use a tanh activation on action prediction
- vocab_size (`int`, *optional*, defaults to 50257):
- Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`DecisionTransformerModel`].
- n_positions (`int`, *optional*, defaults to 1024):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- n_layer (`int`, *optional*, defaults to 3):
- Number of hidden layers in the Transformer encoder.
- n_head (`int`, *optional*, defaults to 1):
- Number of attention heads for each attention layer in the Transformer encoder.
- n_inner (`int`, *optional*):
- Dimensionality of the inner feed-forward layers. If unset, will default to 4 times `n_embd`.
- activation_function (`str`, *optional*, defaults to `"gelu"`):
- Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new"]`.
- resid_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- embd_pdrop (`int`, *optional*, defaults to 0.1):
- The dropout ratio for the embeddings.
- attn_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention.
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
- The epsilon to use in the layer normalization layers.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- scale_attn_weights (`bool`, *optional*, defaults to `True`):
- Scale attention weights by dividing by sqrt(hidden_size)..
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
- scale_attn_by_inverse_layer_idx (`bool`, *optional*, defaults to `False`):
- Whether to additionally scale attention weights by `1 / layer_idx + 1`.
- reorder_and_upcast_attn (`bool`, *optional*, defaults to `False`):
- Whether to scale keys (K) prior to computing attention (dot-product) and upcast attention
- dot-product/softmax to float() when training with mixed precision.
-
- Example:
-
- ```python
- >>> from transformers import DecisionTransformerConfig, DecisionTransformerModel
-
- >>> # Initializing a DecisionTransformer configuration
- >>> configuration = DecisionTransformerConfig()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = DecisionTransformerModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "decision_transformer"
- keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {
- "max_position_embeddings": "n_positions",
- "num_attention_heads": "n_head",
- "num_hidden_layers": "n_layer",
- }
-
- def __init__(
- self,
- state_dim=17,
- act_dim=4,
- hidden_size=128,
- max_ep_len=4096,
- action_tanh=True,
- vocab_size=1,
- n_positions=1024,
- n_layer=3,
- n_head=1,
- n_inner=None,
- activation_function="relu",
- resid_pdrop=0.1,
- embd_pdrop=0.1,
- attn_pdrop=0.1,
- layer_norm_epsilon=1e-5,
- initializer_range=0.02,
- scale_attn_weights=True,
- use_cache=True,
- bos_token_id=50256,
- eos_token_id=50256,
- scale_attn_by_inverse_layer_idx=False,
- reorder_and_upcast_attn=False,
- **kwargs,
- ):
- self.state_dim = state_dim
- self.act_dim = act_dim
- self.hidden_size = hidden_size
- self.max_ep_len = max_ep_len
- self.action_tanh = action_tanh
- self.vocab_size = vocab_size
- self.n_positions = n_positions
- self.n_layer = n_layer
- self.n_head = n_head
- self.n_inner = n_inner
- self.activation_function = activation_function
- self.resid_pdrop = resid_pdrop
- self.embd_pdrop = embd_pdrop
- self.attn_pdrop = attn_pdrop
- self.layer_norm_epsilon = layer_norm_epsilon
- self.initializer_range = initializer_range
- self.scale_attn_weights = scale_attn_weights
- self.use_cache = use_cache
- self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx
- self.reorder_and_upcast_attn = reorder_and_upcast_attn
-
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
-
- super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
diff --git a/transformers/models/decision_transformer/modeling_decision_transformer.py b/transformers/models/decision_transformer/modeling_decision_transformer.py
deleted file mode 100644
index 6f939460aab86f16fee393c992fa37b98f61de7a..0000000000000000000000000000000000000000
--- a/transformers/models/decision_transformer/modeling_decision_transformer.py
+++ /dev/null
@@ -1,937 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Team The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DecisionTransformer model."""
-
-import math
-import os
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.cuda.amp import autocast
-
-from ...activations import ACT2FN
-from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import Conv1D, find_pruneable_heads_and_indices, prune_conv1d_layer
-from ...utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_decision_transformer import DecisionTransformerConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "edbeeching/decision-transformer-gym-hopper-medium"
-_CONFIG_FOR_DOC = "DecisionTransformerConfig"
-
-
-from ..deprecated._archive_maps import DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.gpt2.modeling_gpt2.load_tf_weights_in_gpt2
-def load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path):
- """Load tf checkpoints in a pytorch model"""
- try:
- import re
-
- import tensorflow as tf
- except ImportError:
- logger.error(
- "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
- "https://www.tensorflow.org/install/ for installation instructions."
- )
- raise
- tf_path = os.path.abspath(gpt2_checkpoint_path)
- logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
- # Load weights from TF model
- init_vars = tf.train.list_variables(tf_path)
- names = []
- arrays = []
- for name, shape in init_vars:
- logger.info(f"Loading TF weight {name} with shape {shape}")
- array = tf.train.load_variable(tf_path, name)
- names.append(name)
- arrays.append(array.squeeze())
-
- for name, array in zip(names, arrays):
- name = name[6:] # skip "model/"
- name = name.split("/")
- pointer = model
- for m_name in name:
- if re.fullmatch(r"[A-Za-z]+\d+", m_name):
- scope_names = re.split(r"(\d+)", m_name)
- else:
- scope_names = [m_name]
- if scope_names[0] == "w" or scope_names[0] == "g":
- pointer = getattr(pointer, "weight")
- elif scope_names[0] == "b":
- pointer = getattr(pointer, "bias")
- elif scope_names[0] == "wpe" or scope_names[0] == "wte":
- pointer = getattr(pointer, scope_names[0])
- pointer = getattr(pointer, "weight")
- else:
- pointer = getattr(pointer, scope_names[0])
- if len(scope_names) >= 2:
- num = int(scope_names[1])
- pointer = pointer[num]
- try:
- if pointer.shape != array.shape:
- raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
- except ValueError as e:
- e.args += (pointer.shape, array.shape)
- raise
- logger.info(f"Initialize PyTorch weight {name}")
- pointer.data = torch.from_numpy(array)
- return model
-
-
-# Copied from transformers.models.gpt2.modeling_gpt2.GPT2Attention with GPT2->DecisionTransformerGPT2
-class DecisionTransformerGPT2Attention(nn.Module):
- def __init__(self, config, is_cross_attention=False, layer_idx=None):
- super().__init__()
- self.config = config
- max_positions = config.max_position_embeddings
- self.register_buffer(
- "bias",
- torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
- 1, 1, max_positions, max_positions
- ),
- persistent=False,
- )
- self.register_buffer("masked_bias", torch.tensor(-1e4), persistent=False)
-
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_heads
- self.split_size = self.embed_dim
- if self.head_dim * self.num_heads != self.embed_dim:
- raise ValueError(
- f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {self.num_heads})."
- )
-
- self.scale_attn_weights = config.scale_attn_weights
- self.is_cross_attention = is_cross_attention
-
- # Layer-wise attention scaling, reordering, and upcasting
- self.scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx
- self.layer_idx = layer_idx
- self.reorder_and_upcast_attn = config.reorder_and_upcast_attn
-
- if self.is_cross_attention:
- self.c_attn = Conv1D(2 * self.embed_dim, self.embed_dim)
- self.q_attn = Conv1D(self.embed_dim, self.embed_dim)
- else:
- self.c_attn = Conv1D(3 * self.embed_dim, self.embed_dim)
- self.c_proj = Conv1D(self.embed_dim, self.embed_dim)
-
- self.attn_dropout = nn.Dropout(config.attn_pdrop)
- self.resid_dropout = nn.Dropout(config.resid_pdrop)
- self.is_causal = True
-
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(heads, self.num_heads, self.head_dim, self.pruned_heads)
- index_attn = torch.cat([index, index + self.split_size, index + (2 * self.split_size)])
-
- # Prune conv1d layers
- self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
- self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0)
-
- # Update hyper params
- self.split_size = (self.split_size // self.num_heads) * (self.num_heads - len(heads))
- self.num_heads = self.num_heads - len(heads)
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def _attn(self, query, key, value, attention_mask=None, head_mask=None):
- attn_weights = torch.matmul(query, key.transpose(-1, -2))
-
- if self.scale_attn_weights:
- attn_weights = attn_weights / torch.full(
- [], value.size(-1) ** 0.5, dtype=attn_weights.dtype, device=attn_weights.device
- )
-
- # Layer-wise attention scaling
- if self.scale_attn_by_inverse_layer_idx:
- attn_weights = attn_weights / float(self.layer_idx + 1)
-
- if not self.is_cross_attention:
- # if only "normal" attention layer implements causal mask
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
- mask_value = torch.finfo(attn_weights.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.full([], mask_value, dtype=attn_weights.dtype, device=attn_weights.device)
- attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value)
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise
- attn_weights = attn_weights.type(value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, head_mask=None):
- # Use `torch.baddbmm` (a bit more efficient w/ alpha param for scaling -- from Megatron-LM)
- bsz, num_heads, q_seq_len, dk = query.size()
- _, _, k_seq_len, _ = key.size()
-
- # Preallocate attn_weights for `baddbmm`
- attn_weights = torch.empty(bsz * num_heads, q_seq_len, k_seq_len, dtype=torch.float32, device=query.device)
-
- # Compute Scale Factor
- scale_factor = 1.0
- if self.scale_attn_weights:
- scale_factor /= float(value.size(-1)) ** 0.5
-
- if self.scale_attn_by_inverse_layer_idx:
- scale_factor /= float(self.layer_idx + 1)
-
- # Upcast (turn off autocast) and reorder (Scale K by 1 / root(dk))
- with autocast(enabled=False):
- q, k = query.reshape(-1, q_seq_len, dk), key.transpose(-1, -2).reshape(-1, dk, k_seq_len)
- attn_weights = torch.baddbmm(attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor)
- attn_weights = attn_weights.reshape(bsz, num_heads, q_seq_len, k_seq_len)
-
- if not self.is_cross_attention:
- # if only "normal" attention layer implements causal mask
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
- mask_value = torch.finfo(attn_weights.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
- attn_weights = torch.where(causal_mask, attn_weights, mask_value)
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op if otherwise
- if attn_weights.dtype != torch.float32:
- raise RuntimeError("Error with upcasting, attn_weights does not have dtype torch.float32")
- attn_weights = attn_weights.type(value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def _split_heads(self, tensor, num_heads, attn_head_size):
- """
- Splits hidden_size dim into attn_head_size and num_heads
- """
- new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
- tensor = tensor.view(new_shape)
- return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)
-
- def _merge_heads(self, tensor, num_heads, attn_head_size):
- """
- Merges attn_head_size dim and num_attn_heads dim into hidden_size
- """
- tensor = tensor.permute(0, 2, 1, 3).contiguous()
- new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
- return tensor.view(new_shape)
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
- if encoder_hidden_states is not None:
- if not hasattr(self, "q_attn"):
- raise ValueError(
- "If class is used as cross attention, the weights `q_attn` have to be defined. "
- "Please make sure to instantiate class with `DecisionTransformerGPT2Attention(..., is_cross_attention=True)`."
- )
-
- query = self.q_attn(hidden_states)
- key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
- attention_mask = encoder_attention_mask
- else:
- query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
-
- query = self._split_heads(query, self.num_heads, self.head_dim)
- key = self._split_heads(key, self.num_heads, self.head_dim)
- value = self._split_heads(value, self.num_heads, self.head_dim)
-
- if layer_past is not None:
- past_key, past_value = layer_past
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- if use_cache is True:
- present = (key, value)
- else:
- present = None
-
- if self.reorder_and_upcast_attn:
- attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
- else:
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
- attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
- attn_output = self.c_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs # a, present, (attentions)
-
-
-# Copied from transformers.models.gpt2.modeling_gpt2.GPT2MLP with GPT2->DecisionTransformerGPT2
-class DecisionTransformerGPT2MLP(nn.Module):
- def __init__(self, intermediate_size, config):
- super().__init__()
- embed_dim = config.hidden_size
- self.c_fc = Conv1D(intermediate_size, embed_dim)
- self.c_proj = Conv1D(embed_dim, intermediate_size)
- self.act = ACT2FN[config.activation_function]
- self.dropout = nn.Dropout(config.resid_pdrop)
-
- def forward(self, hidden_states: Optional[Tuple[torch.FloatTensor]]) -> torch.FloatTensor:
- hidden_states = self.c_fc(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.c_proj(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.gpt2.modeling_gpt2.GPT2Block with GPT2->DecisionTransformerGPT2
-class DecisionTransformerGPT2Block(nn.Module):
- # Ignore copy
- def __init__(self, config, layer_idx=None):
- super().__init__()
- hidden_size = config.hidden_size
- inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
-
- self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
- self.attn = DecisionTransformerGPT2Attention(config, layer_idx=layer_idx)
- self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
- if config.add_cross_attention:
- self.crossattention = DecisionTransformerGPT2Attention(
- config, is_cross_attention=True, layer_idx=layer_idx
- )
- self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
- self.mlp = DecisionTransformerGPT2MLP(inner_dim, config)
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
- outputs = attn_outputs[1:]
- # residual connection
- hidden_states = attn_output + residual
-
- if encoder_hidden_states is not None:
- # add one self-attention block for cross-attention
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
- "cross-attention layers by setting `config.add_cross_attention=True`"
- )
- residual = hidden_states
- hidden_states = self.ln_cross_attn(hidden_states)
- cross_attn_outputs = self.crossattention(
- hidden_states,
- attention_mask=attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- )
- attn_output = cross_attn_outputs[0]
- # residual connection
- hidden_states = residual + attn_output
- outputs = outputs + cross_attn_outputs[2:] # add cross attentions if we output attention weights
-
- residual = hidden_states
- hidden_states = self.ln_2(hidden_states)
- feed_forward_hidden_states = self.mlp(hidden_states)
- # residual connection
- hidden_states = residual + feed_forward_hidden_states
-
- if use_cache:
- outputs = (hidden_states,) + outputs
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- return outputs # hidden_states, present, (attentions, cross_attentions)
-
-
-class DecisionTransformerGPT2PreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DecisionTransformerConfig
- load_tf_weights = load_tf_weights_in_gpt2
- base_model_prefix = "transformer"
- is_parallelizable = True
- supports_gradient_checkpointing = True
-
- def __init__(self, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, (nn.Linear, Conv1D)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
- # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
- # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
- # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
- # > -- GPT-2 :: https://openai.com/blog/better-language-models/
- #
- # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
- for name, p in module.named_parameters():
- if "c_proj" in name and "weight" in name:
- # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
- p.data.normal_(mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.n_layer)))
-
-
-class DecisionTransformerGPT2Model(DecisionTransformerGPT2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.embed_dim = config.hidden_size
-
- self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
- self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
-
- self.drop = nn.Dropout(config.embd_pdrop)
- self.h = nn.ModuleList(
- [DecisionTransformerGPT2Block(config, layer_idx=i) for i in range(config.num_hidden_layers)]
- )
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.wte
-
- def set_input_embeddings(self, new_embeddings):
- self.wte = new_embeddings
-
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- input_ids = input_ids.view(-1, input_shape[-1])
- batch_size = input_ids.shape[0]
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- batch_size = inputs_embeds.shape[0]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if token_type_ids is not None:
- token_type_ids = token_type_ids.view(-1, input_shape[-1])
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * len(self.h))
- else:
- past_length = past_key_values[0][0].size(-2)
- if position_ids is None:
- position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0)
-
- # Attention mask.
- if attention_mask is not None:
- if batch_size <= 0:
- raise ValueError("batch_size has to be defined and > 0")
- attention_mask = attention_mask.view(batch_size, -1)
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask = attention_mask[:, None, None, :]
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and the dtype's smallest value for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
- attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
- # If a 2D or 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- if self.config.add_cross_attention and encoder_hidden_states is not None:
- encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
- encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
- if encoder_attention_mask is None:
- encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
- encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
- else:
- encoder_attention_mask = None
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # head_mask has shape n_layer x batch x n_heads x N x N
- head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-
- if inputs_embeds is None:
- inputs_embeds = self.wte(input_ids)
- position_embeds = self.wpe(position_ids)
- hidden_states = inputs_embeds + position_embeds
-
- if token_type_ids is not None:
- token_type_embeds = self.wte(token_type_ids)
- hidden_states = hidden_states + token_type_embeds
-
- hidden_states = self.drop(hidden_states)
-
- output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- presents = () if use_cache else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
- all_hidden_states = () if output_hidden_states else None
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- # Model parallel
- if self.model_parallel:
- torch.cuda.set_device(hidden_states.device)
- # Ensure layer_past is on same device as hidden_states (might not be correct)
- if layer_past is not None:
- layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past)
- # Ensure that attention_mask is always on the same device as hidden_states
- if attention_mask is not None:
- attention_mask = attention_mask.to(hidden_states.device)
- if isinstance(head_mask, torch.Tensor):
- head_mask = head_mask.to(hidden_states.device)
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- None,
- attention_mask,
- head_mask[i],
- encoder_hidden_states,
- encoder_attention_mask,
- use_cache,
- output_attentions,
- )
- else:
- outputs = block(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask[i],
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
-
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],)
-
- # Model Parallel: If it's the last layer for that device, put things on the next device
- if self.model_parallel:
- for k, v in self.device_map.items():
- if i == v[-1] and "cuda:" + str(k) != self.last_device:
- hidden_states = hidden_states.to("cuda:" + str(k + 1))
-
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = hidden_states.view(output_shape)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions]
- if v is not None
- )
-
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-@dataclass
-class DecisionTransformerOutput(ModelOutput):
- """
- Base class for model's outputs that also contains a pooling of the last hidden states.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- state_preds (`torch.FloatTensor` of shape `(batch_size, sequence_length, state_dim)`):
- Environment state predictions
- action_preds (`torch.FloatTensor` of shape `(batch_size, sequence_length, action_dim)`):
- Model action predictions
- return_preds (`torch.FloatTensor` of shape `(batch_size, sequence_length, 1)`):
- Predicted returns for each state
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- state_preds: torch.FloatTensor = None
- action_preds: torch.FloatTensor = None
- return_preds: torch.FloatTensor = None
- hidden_states: torch.FloatTensor = None
- attentions: torch.FloatTensor = None
- last_hidden_state: torch.FloatTensor = None
-
-
-class DecisionTransformerPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DecisionTransformerConfig
- base_model_prefix = "decision_transformer"
- main_input_name = "states"
- supports_gradient_checkpointing = False
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-DECISION_TRANSFORMER_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`~DecisionTransformerConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DECISION_TRANSFORMER_INPUTS_DOCSTRING = r"""
- Args:
- states (`torch.FloatTensor` of shape `(batch_size, episode_length, state_dim)`):
- The states for each step in the trajectory
- actions (`torch.FloatTensor` of shape `(batch_size, episode_length, act_dim)`):
- The actions taken by the "expert" policy for the current state, these are masked for auto regressive
- prediction
- rewards (`torch.FloatTensor` of shape `(batch_size, episode_length, 1)`):
- The rewards for each state, action
- returns_to_go (`torch.FloatTensor` of shape `(batch_size, episode_length, 1)`):
- The returns for each state in the trajectory
- timesteps (`torch.LongTensor` of shape `(batch_size, episode_length)`):
- The timestep for each step in the trajectory
- attention_mask (`torch.FloatTensor` of shape `(batch_size, episode_length)`):
- Masking, used to mask the actions when performing autoregressive prediction
-"""
-
-
-@add_start_docstrings("The Decision Transformer Model", DECISION_TRANSFORMER_START_DOCSTRING)
-class DecisionTransformerModel(DecisionTransformerPreTrainedModel):
- """
-
- The model builds upon the GPT2 architecture to perform autoregressive prediction of actions in an offline RL
- setting. Refer to the paper for more details: https://arxiv.org/abs/2106.01345
-
- """
-
- def __init__(self, config):
- super().__init__(config)
- self.config = config
- self.hidden_size = config.hidden_size
- # note: the only difference between this GPT2Model and the default Huggingface version
- # is that the positional embeddings are removed (since we'll add those ourselves)
- self.encoder = DecisionTransformerGPT2Model(config)
-
- self.embed_timestep = nn.Embedding(config.max_ep_len, config.hidden_size)
- self.embed_return = torch.nn.Linear(1, config.hidden_size)
- self.embed_state = torch.nn.Linear(config.state_dim, config.hidden_size)
- self.embed_action = torch.nn.Linear(config.act_dim, config.hidden_size)
-
- self.embed_ln = nn.LayerNorm(config.hidden_size)
-
- # note: we don't predict states or returns for the paper
- self.predict_state = torch.nn.Linear(config.hidden_size, config.state_dim)
- self.predict_action = nn.Sequential(
- *([nn.Linear(config.hidden_size, config.act_dim)] + ([nn.Tanh()] if config.action_tanh else []))
- )
- self.predict_return = torch.nn.Linear(config.hidden_size, 1)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DECISION_TRANSFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=DecisionTransformerOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- states: Optional[torch.FloatTensor] = None,
- actions: Optional[torch.FloatTensor] = None,
- rewards: Optional[torch.FloatTensor] = None,
- returns_to_go: Optional[torch.FloatTensor] = None,
- timesteps: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- output_hidden_states: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], DecisionTransformerOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import DecisionTransformerModel
- >>> import torch
-
- >>> model = DecisionTransformerModel.from_pretrained("edbeeching/decision-transformer-gym-hopper-medium")
- >>> # evaluation
- >>> model = model.to(device)
- >>> model.eval()
-
- >>> env = gym.make("Hopper-v3")
- >>> state_dim = env.observation_space.shape[0]
- >>> act_dim = env.action_space.shape[0]
-
- >>> state = env.reset()
- >>> states = torch.from_numpy(state).reshape(1, 1, state_dim).to(device=device, dtype=torch.float32)
- >>> actions = torch.zeros((1, 1, act_dim), device=device, dtype=torch.float32)
- >>> rewards = torch.zeros(1, 1, device=device, dtype=torch.float32)
- >>> target_return = torch.tensor(TARGET_RETURN, dtype=torch.float32).reshape(1, 1)
- >>> timesteps = torch.tensor(0, device=device, dtype=torch.long).reshape(1, 1)
- >>> attention_mask = torch.zeros(1, 1, device=device, dtype=torch.float32)
-
- >>> # forward pass
- >>> with torch.no_grad():
- ... state_preds, action_preds, return_preds = model(
- ... states=states,
- ... actions=actions,
- ... rewards=rewards,
- ... returns_to_go=target_return,
- ... timesteps=timesteps,
- ... attention_mask=attention_mask,
- ... return_dict=False,
- ... )
- ```"""
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- batch_size, seq_length = states.shape[0], states.shape[1]
-
- if attention_mask is None:
- # attention mask for GPT: 1 if can be attended to, 0 if not
- attention_mask = torch.ones((batch_size, seq_length), dtype=torch.long)
-
- # embed each modality with a different head
- state_embeddings = self.embed_state(states)
- action_embeddings = self.embed_action(actions)
- returns_embeddings = self.embed_return(returns_to_go)
- time_embeddings = self.embed_timestep(timesteps)
-
- # time embeddings are treated similar to positional embeddings
- state_embeddings = state_embeddings + time_embeddings
- action_embeddings = action_embeddings + time_embeddings
- returns_embeddings = returns_embeddings + time_embeddings
-
- # this makes the sequence look like (R_1, s_1, a_1, R_2, s_2, a_2, ...)
- # which works nice in an autoregressive sense since states predict actions
- stacked_inputs = (
- torch.stack((returns_embeddings, state_embeddings, action_embeddings), dim=1)
- .permute(0, 2, 1, 3)
- .reshape(batch_size, 3 * seq_length, self.hidden_size)
- )
- stacked_inputs = self.embed_ln(stacked_inputs)
-
- # to make the attention mask fit the stacked inputs, have to stack it as well
- stacked_attention_mask = (
- torch.stack((attention_mask, attention_mask, attention_mask), dim=1)
- .permute(0, 2, 1)
- .reshape(batch_size, 3 * seq_length)
- )
- device = stacked_inputs.device
- # we feed in the input embeddings (not word indices as in NLP) to the model
- encoder_outputs = self.encoder(
- inputs_embeds=stacked_inputs,
- attention_mask=stacked_attention_mask,
- position_ids=torch.zeros(stacked_attention_mask.shape, device=device, dtype=torch.long),
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- x = encoder_outputs[0]
-
- # reshape x so that the second dimension corresponds to the original
- # returns (0), states (1), or actions (2); i.e. x[:,1,t] is the token for s_t
- x = x.reshape(batch_size, seq_length, 3, self.hidden_size).permute(0, 2, 1, 3)
-
- # get predictions
- return_preds = self.predict_return(x[:, 2]) # predict next return given state and action
- state_preds = self.predict_state(x[:, 2]) # predict next state given state and action
- action_preds = self.predict_action(x[:, 1]) # predict next action given state
- if not return_dict:
- return (state_preds, action_preds, return_preds)
-
- return DecisionTransformerOutput(
- last_hidden_state=encoder_outputs.last_hidden_state,
- state_preds=state_preds,
- action_preds=action_preds,
- return_preds=return_preds,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
diff --git a/transformers/models/deformable_detr/__init__.py b/transformers/models/deformable_detr/__init__.py
deleted file mode 100644
index a560265f4bfcb8d43f88d2b3cd55f751409016ec..0000000000000000000000000000000000000000
--- a/transformers/models/deformable_detr/__init__.py
+++ /dev/null
@@ -1,75 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {
- "configuration_deformable_detr": ["DEFORMABLE_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP", "DeformableDetrConfig"],
-}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_deformable_detr"] = ["DeformableDetrFeatureExtractor"]
- _import_structure["image_processing_deformable_detr"] = ["DeformableDetrImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_deformable_detr"] = [
- "DEFORMABLE_DETR_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DeformableDetrForObjectDetection",
- "DeformableDetrModel",
- "DeformableDetrPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_deformable_detr import DEFORMABLE_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP, DeformableDetrConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_deformable_detr import DeformableDetrFeatureExtractor
- from .image_processing_deformable_detr import DeformableDetrImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_deformable_detr import (
- DEFORMABLE_DETR_PRETRAINED_MODEL_ARCHIVE_LIST,
- DeformableDetrForObjectDetection,
- DeformableDetrModel,
- DeformableDetrPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deformable_detr/__pycache__/__init__.cpython-310.pyc b/transformers/models/deformable_detr/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 3a77aa0189d34486352365078f2daf7965220d73..0000000000000000000000000000000000000000
Binary files a/transformers/models/deformable_detr/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deformable_detr/__pycache__/configuration_deformable_detr.cpython-310.pyc b/transformers/models/deformable_detr/__pycache__/configuration_deformable_detr.cpython-310.pyc
deleted file mode 100644
index 081701a56606ba3e11536c373c7f2d0455da2d99..0000000000000000000000000000000000000000
Binary files a/transformers/models/deformable_detr/__pycache__/configuration_deformable_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deformable_detr/__pycache__/convert_deformable_detr_to_pytorch.cpython-310.pyc b/transformers/models/deformable_detr/__pycache__/convert_deformable_detr_to_pytorch.cpython-310.pyc
deleted file mode 100644
index f9f1991c5b1d455ee3afcb0b87be5970626485f9..0000000000000000000000000000000000000000
Binary files a/transformers/models/deformable_detr/__pycache__/convert_deformable_detr_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deformable_detr/__pycache__/feature_extraction_deformable_detr.cpython-310.pyc b/transformers/models/deformable_detr/__pycache__/feature_extraction_deformable_detr.cpython-310.pyc
deleted file mode 100644
index 94dc84372611f0a9174adacb88c63c3b43247c21..0000000000000000000000000000000000000000
Binary files a/transformers/models/deformable_detr/__pycache__/feature_extraction_deformable_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deformable_detr/__pycache__/image_processing_deformable_detr.cpython-310.pyc b/transformers/models/deformable_detr/__pycache__/image_processing_deformable_detr.cpython-310.pyc
deleted file mode 100644
index c470fe58dde175b557159a7133641b2801c367b1..0000000000000000000000000000000000000000
Binary files a/transformers/models/deformable_detr/__pycache__/image_processing_deformable_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deformable_detr/__pycache__/load_custom.cpython-310.pyc b/transformers/models/deformable_detr/__pycache__/load_custom.cpython-310.pyc
deleted file mode 100644
index db5572d4d840807cd74fae37624037e27d069c37..0000000000000000000000000000000000000000
Binary files a/transformers/models/deformable_detr/__pycache__/load_custom.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deformable_detr/__pycache__/modeling_deformable_detr.cpython-310.pyc b/transformers/models/deformable_detr/__pycache__/modeling_deformable_detr.cpython-310.pyc
deleted file mode 100644
index baa4d8a78e419603ae15fff55c7111a00b3fde94..0000000000000000000000000000000000000000
Binary files a/transformers/models/deformable_detr/__pycache__/modeling_deformable_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deformable_detr/configuration_deformable_detr.py b/transformers/models/deformable_detr/configuration_deformable_detr.py
deleted file mode 100644
index 6d32f6220df5868a7d7aca620444306f5ad33ba5..0000000000000000000000000000000000000000
--- a/transformers/models/deformable_detr/configuration_deformable_detr.py
+++ /dev/null
@@ -1,277 +0,0 @@
-# coding=utf-8
-# Copyright 2022 SenseTime and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Deformable DETR model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-from ..auto import CONFIG_MAPPING
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DEFORMABLE_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DeformableDetrConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DeformableDetrModel`]. It is used to instantiate
- a Deformable DETR model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the Deformable DETR
- [SenseTime/deformable-detr](https://huggingface.co/SenseTime/deformable-detr) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- use_timm_backbone (`bool`, *optional*, defaults to `True`):
- Whether or not to use the `timm` library for the backbone. If set to `False`, will use the [`AutoBackbone`]
- API.
- backbone_config (`PretrainedConfig` or `dict`, *optional*):
- The configuration of the backbone model. Only used in case `use_timm_backbone` is set to `False` in which
- case it will default to `ResNetConfig()`.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- num_queries (`int`, *optional*, defaults to 300):
- Number of object queries, i.e. detection slots. This is the maximal number of objects
- [`DeformableDetrModel`] can detect in a single image. In case `two_stage` is set to `True`, we use
- `two_stage_num_proposals` instead.
- d_model (`int`, *optional*, defaults to 256):
- Dimension of the layers.
- encoder_layers (`int`, *optional*, defaults to 6):
- Number of encoder layers.
- decoder_layers (`int`, *optional*, defaults to 6):
- Number of decoder layers.
- encoder_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the Transformer encoder.
- decoder_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the Transformer decoder.
- decoder_ffn_dim (`int`, *optional*, defaults to 1024):
- Dimension of the "intermediate" (often named feed-forward) layer in decoder.
- encoder_ffn_dim (`int`, *optional*, defaults to 1024):
- Dimension of the "intermediate" (often named feed-forward) layer in decoder.
- activation_function (`str` or `function`, *optional*, defaults to `"relu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- activation_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for activations inside the fully connected layer.
- init_std (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- init_xavier_std (`float`, *optional*, defaults to 1):
- The scaling factor used for the Xavier initialization gain in the HM Attention map module.
- encoder_layerdrop (`float`, *optional*, defaults to 0.0):
- The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
- for more details.
- auxiliary_loss (`bool`, *optional*, defaults to `False`):
- Whether auxiliary decoding losses (loss at each decoder layer) are to be used.
- position_embedding_type (`str`, *optional*, defaults to `"sine"`):
- Type of position embeddings to be used on top of the image features. One of `"sine"` or `"learned"`.
- backbone (`str`, *optional*, defaults to `"resnet50"`):
- Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
- will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
- is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
- use_pretrained_backbone (`bool`, *optional*, defaults to `True`):
- Whether to use pretrained weights for the backbone.
- backbone_kwargs (`dict`, *optional*):
- Keyword arguments to be passed to AutoBackbone when loading from a checkpoint
- e.g. `{'out_indices': (0, 1, 2, 3)}`. Cannot be specified if `backbone_config` is set.
- dilation (`bool`, *optional*, defaults to `False`):
- Whether to replace stride with dilation in the last convolutional block (DC5). Only supported when
- `use_timm_backbone` = `True`.
- class_cost (`float`, *optional*, defaults to 1):
- Relative weight of the classification error in the Hungarian matching cost.
- bbox_cost (`float`, *optional*, defaults to 5):
- Relative weight of the L1 error of the bounding box coordinates in the Hungarian matching cost.
- giou_cost (`float`, *optional*, defaults to 2):
- Relative weight of the generalized IoU loss of the bounding box in the Hungarian matching cost.
- mask_loss_coefficient (`float`, *optional*, defaults to 1):
- Relative weight of the Focal loss in the panoptic segmentation loss.
- dice_loss_coefficient (`float`, *optional*, defaults to 1):
- Relative weight of the DICE/F-1 loss in the panoptic segmentation loss.
- bbox_loss_coefficient (`float`, *optional*, defaults to 5):
- Relative weight of the L1 bounding box loss in the object detection loss.
- giou_loss_coefficient (`float`, *optional*, defaults to 2):
- Relative weight of the generalized IoU loss in the object detection loss.
- eos_coefficient (`float`, *optional*, defaults to 0.1):
- Relative classification weight of the 'no-object' class in the object detection loss.
- num_feature_levels (`int`, *optional*, defaults to 4):
- The number of input feature levels.
- encoder_n_points (`int`, *optional*, defaults to 4):
- The number of sampled keys in each feature level for each attention head in the encoder.
- decoder_n_points (`int`, *optional*, defaults to 4):
- The number of sampled keys in each feature level for each attention head in the decoder.
- two_stage (`bool`, *optional*, defaults to `False`):
- Whether to apply a two-stage deformable DETR, where the region proposals are also generated by a variant of
- Deformable DETR, which are further fed into the decoder for iterative bounding box refinement.
- two_stage_num_proposals (`int`, *optional*, defaults to 300):
- The number of region proposals to be generated, in case `two_stage` is set to `True`.
- with_box_refine (`bool`, *optional*, defaults to `False`):
- Whether to apply iterative bounding box refinement, where each decoder layer refines the bounding boxes
- based on the predictions from the previous layer.
- focal_alpha (`float`, *optional*, defaults to 0.25):
- Alpha parameter in the focal loss.
- disable_custom_kernels (`bool`, *optional*, defaults to `False`):
- Disable the use of custom CUDA and CPU kernels. This option is necessary for the ONNX export, as custom
- kernels are not supported by PyTorch ONNX export.
-
- Examples:
-
- ```python
- >>> from transformers import DeformableDetrConfig, DeformableDetrModel
-
- >>> # Initializing a Deformable DETR SenseTime/deformable-detr style configuration
- >>> configuration = DeformableDetrConfig()
-
- >>> # Initializing a model (with random weights) from the SenseTime/deformable-detr style configuration
- >>> model = DeformableDetrModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "deformable_detr"
- attribute_map = {
- "hidden_size": "d_model",
- "num_attention_heads": "encoder_attention_heads",
- }
-
- def __init__(
- self,
- use_timm_backbone=True,
- backbone_config=None,
- num_channels=3,
- num_queries=300,
- max_position_embeddings=1024,
- encoder_layers=6,
- encoder_ffn_dim=1024,
- encoder_attention_heads=8,
- decoder_layers=6,
- decoder_ffn_dim=1024,
- decoder_attention_heads=8,
- encoder_layerdrop=0.0,
- is_encoder_decoder=True,
- activation_function="relu",
- d_model=256,
- dropout=0.1,
- attention_dropout=0.0,
- activation_dropout=0.0,
- init_std=0.02,
- init_xavier_std=1.0,
- return_intermediate=True,
- auxiliary_loss=False,
- position_embedding_type="sine",
- backbone="resnet50",
- use_pretrained_backbone=True,
- backbone_kwargs=None,
- dilation=False,
- num_feature_levels=4,
- encoder_n_points=4,
- decoder_n_points=4,
- two_stage=False,
- two_stage_num_proposals=300,
- with_box_refine=False,
- class_cost=1,
- bbox_cost=5,
- giou_cost=2,
- mask_loss_coefficient=1,
- dice_loss_coefficient=1,
- bbox_loss_coefficient=5,
- giou_loss_coefficient=2,
- eos_coefficient=0.1,
- focal_alpha=0.25,
- disable_custom_kernels=False,
- **kwargs,
- ):
- if not use_timm_backbone and use_pretrained_backbone:
- raise ValueError(
- "Loading pretrained backbone weights from the transformers library is not supported yet. `use_timm_backbone` must be set to `True` when `use_pretrained_backbone=True`"
- )
-
- if backbone_config is not None and backbone is not None:
- raise ValueError("You can't specify both `backbone` and `backbone_config`.")
-
- if backbone_config is not None and use_timm_backbone:
- raise ValueError("You can't specify both `backbone_config` and `use_timm_backbone`.")
-
- if backbone_kwargs is not None and backbone_kwargs and backbone_config is not None:
- raise ValueError("You can't specify both `backbone_kwargs` and `backbone_config`.")
-
- if not use_timm_backbone:
- if backbone_config is None:
- logger.info("`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.")
- backbone_config = CONFIG_MAPPING["resnet"](out_features=["stage4"])
- elif isinstance(backbone_config, dict):
- backbone_model_type = backbone_config.get("model_type")
- config_class = CONFIG_MAPPING[backbone_model_type]
- backbone_config = config_class.from_dict(backbone_config)
- self.use_timm_backbone = use_timm_backbone
- self.backbone_config = backbone_config
- self.num_channels = num_channels
- self.num_queries = num_queries
- self.max_position_embeddings = max_position_embeddings
- self.d_model = d_model
- self.encoder_ffn_dim = encoder_ffn_dim
- self.encoder_layers = encoder_layers
- self.encoder_attention_heads = encoder_attention_heads
- self.decoder_ffn_dim = decoder_ffn_dim
- self.decoder_layers = decoder_layers
- self.decoder_attention_heads = decoder_attention_heads
- self.dropout = dropout
- self.attention_dropout = attention_dropout
- self.activation_dropout = activation_dropout
- self.activation_function = activation_function
- self.init_std = init_std
- self.init_xavier_std = init_xavier_std
- self.encoder_layerdrop = encoder_layerdrop
- self.auxiliary_loss = auxiliary_loss
- self.position_embedding_type = position_embedding_type
- self.backbone = backbone
- self.use_pretrained_backbone = use_pretrained_backbone
- self.backbone_kwargs = backbone_kwargs
- self.dilation = dilation
- # deformable attributes
- self.num_feature_levels = num_feature_levels
- self.encoder_n_points = encoder_n_points
- self.decoder_n_points = decoder_n_points
- self.two_stage = two_stage
- self.two_stage_num_proposals = two_stage_num_proposals
- self.with_box_refine = with_box_refine
- if two_stage is True and with_box_refine is False:
- raise ValueError("If two_stage is True, with_box_refine must be True.")
- # Hungarian matcher
- self.class_cost = class_cost
- self.bbox_cost = bbox_cost
- self.giou_cost = giou_cost
- # Loss coefficients
- self.mask_loss_coefficient = mask_loss_coefficient
- self.dice_loss_coefficient = dice_loss_coefficient
- self.bbox_loss_coefficient = bbox_loss_coefficient
- self.giou_loss_coefficient = giou_loss_coefficient
- self.eos_coefficient = eos_coefficient
- self.focal_alpha = focal_alpha
- self.disable_custom_kernels = disable_custom_kernels
- super().__init__(is_encoder_decoder=is_encoder_decoder, **kwargs)
-
- @property
- def num_attention_heads(self) -> int:
- return self.encoder_attention_heads
-
- @property
- def hidden_size(self) -> int:
- return self.d_model
diff --git a/transformers/models/deformable_detr/convert_deformable_detr_to_pytorch.py b/transformers/models/deformable_detr/convert_deformable_detr_to_pytorch.py
deleted file mode 100644
index 928fa368ed34c2d3f59baa8038aded596df2d58e..0000000000000000000000000000000000000000
--- a/transformers/models/deformable_detr/convert_deformable_detr_to_pytorch.py
+++ /dev/null
@@ -1,237 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Deformable DETR checkpoints."""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import cached_download, hf_hub_url
-from PIL import Image
-
-from transformers import DeformableDetrConfig, DeformableDetrForObjectDetection, DeformableDetrImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def rename_key(orig_key):
- if "backbone.0.body" in orig_key:
- orig_key = orig_key.replace("backbone.0.body", "backbone.conv_encoder.model")
- if "transformer" in orig_key:
- orig_key = orig_key.replace("transformer.", "")
- if "norm1" in orig_key:
- if "encoder" in orig_key:
- orig_key = orig_key.replace("norm1", "self_attn_layer_norm")
- else:
- orig_key = orig_key.replace("norm1", "encoder_attn_layer_norm")
- if "norm2" in orig_key:
- if "encoder" in orig_key:
- orig_key = orig_key.replace("norm2", "final_layer_norm")
- else:
- orig_key = orig_key.replace("norm2", "self_attn_layer_norm")
- if "norm3" in orig_key:
- orig_key = orig_key.replace("norm3", "final_layer_norm")
- if "linear1" in orig_key:
- orig_key = orig_key.replace("linear1", "fc1")
- if "linear2" in orig_key:
- orig_key = orig_key.replace("linear2", "fc2")
- if "query_embed" in orig_key:
- orig_key = orig_key.replace("query_embed", "query_position_embeddings")
- if "cross_attn" in orig_key:
- orig_key = orig_key.replace("cross_attn", "encoder_attn")
-
- return orig_key
-
-
-def read_in_q_k_v(state_dict):
- # transformer decoder self-attention layers
- for i in range(6):
- # read in weights + bias of input projection layer of self-attention
- in_proj_weight = state_dict.pop(f"decoder.layers.{i}.self_attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"decoder.layers.{i}.self_attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"decoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
- state_dict[f"decoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
- state_dict[f"decoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
- state_dict[f"decoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
- state_dict[f"decoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
- state_dict[f"decoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
-
- return im
-
-
-@torch.no_grad()
-def convert_deformable_detr_checkpoint(
- checkpoint_path,
- single_scale,
- dilation,
- with_box_refine,
- two_stage,
- pytorch_dump_folder_path,
- push_to_hub,
-):
- """
- Copy/paste/tweak model's weights to our Deformable DETR structure.
- """
-
- # load default config
- config = DeformableDetrConfig()
- # set config attributes
- if single_scale:
- config.num_feature_levels = 1
- config.dilation = dilation
- config.with_box_refine = with_box_refine
- config.two_stage = two_stage
- # set labels
- config.num_labels = 91
- repo_id = "huggingface/label-files"
- filename = "coco-detection-id2label.json"
- id2label = json.load(open(cached_download(hf_hub_url(repo_id, filename, repo_type="dataset")), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
-
- # load image processor
- image_processor = DeformableDetrImageProcessor(format="coco_detection")
-
- # prepare image
- img = prepare_img()
- encoding = image_processor(images=img, return_tensors="pt")
- pixel_values = encoding["pixel_values"]
-
- logger.info("Converting model...")
-
- # load original state dict
- state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
- # rename keys
- for key in state_dict.copy().keys():
- val = state_dict.pop(key)
- state_dict[rename_key(key)] = val
- # query, key and value matrices need special treatment
- read_in_q_k_v(state_dict)
- # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
- prefix = "model."
- for key in state_dict.copy().keys():
- if not key.startswith("class_embed") and not key.startswith("bbox_embed"):
- val = state_dict.pop(key)
- state_dict[prefix + key] = val
- # finally, create HuggingFace model and load state dict
- model = DeformableDetrForObjectDetection(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- device = "cuda" if torch.cuda.is_available() else "cpu"
- model.to(device)
- # verify our conversion
- outputs = model(pixel_values.to(device))
-
- expected_logits = torch.tensor(
- [[-9.6645, -4.3449, -5.8705], [-9.7035, -3.8504, -5.0724], [-10.5634, -5.3379, -7.5116]]
- )
- expected_boxes = torch.tensor([[0.8693, 0.2289, 0.2492], [0.3150, 0.5489, 0.5845], [0.5563, 0.7580, 0.8518]])
-
- if single_scale:
- expected_logits = torch.tensor(
- [[-9.9051, -4.2541, -6.4852], [-9.6947, -4.0854, -6.8033], [-10.0665, -5.8470, -7.7003]]
- )
- expected_boxes = torch.tensor([[0.7292, 0.4991, 0.5532], [0.7959, 0.2426, 0.4236], [0.7582, 0.3518, 0.4451]])
-
- if single_scale and dilation:
- expected_logits = torch.tensor(
- [[-8.9652, -4.1074, -5.6635], [-9.0596, -4.9447, -6.6075], [-10.1178, -4.5275, -6.2671]]
- )
- expected_boxes = torch.tensor([[0.7665, 0.4130, 0.4769], [0.8364, 0.1841, 0.3391], [0.6261, 0.3895, 0.7978]])
-
- if with_box_refine:
- expected_logits = torch.tensor(
- [[-8.8895, -5.4187, -6.8153], [-8.4706, -6.1668, -7.6184], [-9.0042, -5.5359, -6.9141]]
- )
- expected_boxes = torch.tensor([[0.7828, 0.2208, 0.4323], [0.0892, 0.5996, 0.1319], [0.5524, 0.6389, 0.8914]])
-
- if with_box_refine and two_stage:
- expected_logits = torch.tensor(
- [[-6.7108, -4.3213, -6.3777], [-8.9014, -6.1799, -6.7240], [-6.9315, -4.4735, -6.2298]]
- )
- expected_boxes = torch.tensor([[0.2583, 0.5499, 0.4683], [0.7652, 0.9068, 0.4882], [0.5490, 0.2763, 0.0564]])
-
- print("Logits:", outputs.logits[0, :3, :3])
-
- assert torch.allclose(outputs.logits[0, :3, :3], expected_logits.to(device), atol=1e-4)
- assert torch.allclose(outputs.pred_boxes[0, :3, :3], expected_boxes.to(device), atol=1e-4)
-
- print("Everything ok!")
-
- # Save model and image processor
- logger.info(f"Saving PyTorch model and image processor to {pytorch_dump_folder_path}...")
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- model.save_pretrained(pytorch_dump_folder_path)
- image_processor.save_pretrained(pytorch_dump_folder_path)
-
- # Push to hub
- if push_to_hub:
- model_name = "deformable-detr"
- model_name += "-single-scale" if single_scale else ""
- model_name += "-dc5" if dilation else ""
- model_name += "-with-box-refine" if with_box_refine else ""
- model_name += "-two-stage" if two_stage else ""
- print("Pushing model to hub...")
- model.push_to_hub(repo_path_or_name=model_name, organization="nielsr", commit_message="Add model")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- "--checkpoint_path",
- type=str,
- default="/home/niels/checkpoints/deformable_detr/r50_deformable_detr-checkpoint.pth",
- help="Path to Pytorch checkpoint (.pth file) you'd like to convert.",
- )
- parser.add_argument("--single_scale", action="store_true", help="Whether to set config.num_features_levels = 1.")
- parser.add_argument("--dilation", action="store_true", help="Whether to set config.dilation=True.")
- parser.add_argument("--with_box_refine", action="store_true", help="Whether to set config.with_box_refine=True.")
- parser.add_argument("--two_stage", action="store_true", help="Whether to set config.two_stage=True.")
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- required=True,
- help="Path to the folder to output PyTorch model.",
- )
- parser.add_argument(
- "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
- )
- args = parser.parse_args()
- convert_deformable_detr_checkpoint(
- args.checkpoint_path,
- args.single_scale,
- args.dilation,
- args.with_box_refine,
- args.two_stage,
- args.pytorch_dump_folder_path,
- args.push_to_hub,
- )
diff --git a/transformers/models/deformable_detr/feature_extraction_deformable_detr.py b/transformers/models/deformable_detr/feature_extraction_deformable_detr.py
deleted file mode 100644
index f04743e91ceefe5fbad2485e9767f0a97dd6db49..0000000000000000000000000000000000000000
--- a/transformers/models/deformable_detr/feature_extraction_deformable_detr.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for Deformable DETR."""
-
-import warnings
-
-from ...image_transforms import rgb_to_id as _rgb_to_id
-from ...utils import logging
-from .image_processing_deformable_detr import DeformableDetrImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-def rgb_to_id(x):
- warnings.warn(
- "rgb_to_id has moved and will not be importable from this module from v5. "
- "Please import from transformers.image_transforms instead.",
- FutureWarning,
- )
- return _rgb_to_id(x)
-
-
-class DeformableDetrFeatureExtractor(DeformableDetrImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class DeformableDetrFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
- " Please use DeformableDetrImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/deformable_detr/image_processing_deformable_detr.py b/transformers/models/deformable_detr/image_processing_deformable_detr.py
deleted file mode 100644
index 5525eeeb8c58d53f20af014e8c83071af1189ccd..0000000000000000000000000000000000000000
--- a/transformers/models/deformable_detr/image_processing_deformable_detr.py
+++ /dev/null
@@ -1,1553 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for Deformable DETR."""
-
-import io
-import pathlib
-from collections import defaultdict
-from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
-
-import numpy as np
-
-from ...feature_extraction_utils import BatchFeature
-from ...image_processing_utils import BaseImageProcessor, get_size_dict
-from ...image_transforms import (
- PaddingMode,
- center_to_corners_format,
- corners_to_center_format,
- id_to_rgb,
- pad,
- rescale,
- resize,
- rgb_to_id,
- to_channel_dimension_format,
-)
-from ...image_utils import (
- IMAGENET_DEFAULT_MEAN,
- IMAGENET_DEFAULT_STD,
- AnnotationFormat,
- AnnotationType,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- get_image_size,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_annotations,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import (
- TensorType,
- is_flax_available,
- is_jax_tensor,
- is_scipy_available,
- is_tf_available,
- is_tf_tensor,
- is_torch_available,
- is_torch_tensor,
- is_vision_available,
- logging,
-)
-
-
-if is_torch_available():
- import torch
- from torch import nn
-
-
-if is_vision_available():
- import PIL
-
-if is_scipy_available():
- import scipy.special
- import scipy.stats
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-SUPPORTED_ANNOTATION_FORMATS = (AnnotationFormat.COCO_DETECTION, AnnotationFormat.COCO_PANOPTIC)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_size_with_aspect_ratio
-def get_size_with_aspect_ratio(image_size, size, max_size=None) -> Tuple[int, int]:
- """
- Computes the output image size given the input image size and the desired output size.
-
- Args:
- image_size (`Tuple[int, int]`):
- The input image size.
- size (`int`):
- The desired output size.
- max_size (`int`, *optional*):
- The maximum allowed output size.
- """
- height, width = image_size
- if max_size is not None:
- min_original_size = float(min((height, width)))
- max_original_size = float(max((height, width)))
- if max_original_size / min_original_size * size > max_size:
- size = int(round(max_size * min_original_size / max_original_size))
-
- if (height <= width and height == size) or (width <= height and width == size):
- return height, width
-
- if width < height:
- ow = size
- oh = int(size * height / width)
- else:
- oh = size
- ow = int(size * width / height)
- return (oh, ow)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_resize_output_image_size
-def get_resize_output_image_size(
- input_image: np.ndarray,
- size: Union[int, Tuple[int, int], List[int]],
- max_size: Optional[int] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
-) -> Tuple[int, int]:
- """
- Computes the output image size given the input image size and the desired output size. If the desired output size
- is a tuple or list, the output image size is returned as is. If the desired output size is an integer, the output
- image size is computed by keeping the aspect ratio of the input image size.
-
- Args:
- input_image (`np.ndarray`):
- The image to resize.
- size (`int` or `Tuple[int, int]` or `List[int]`):
- The desired output size.
- max_size (`int`, *optional*):
- The maximum allowed output size.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred from the input image.
- """
- image_size = get_image_size(input_image, input_data_format)
- if isinstance(size, (list, tuple)):
- return size
-
- return get_size_with_aspect_ratio(image_size, size, max_size)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_numpy_to_framework_fn
-def get_numpy_to_framework_fn(arr) -> Callable:
- """
- Returns a function that converts a numpy array to the framework of the input array.
-
- Args:
- arr (`np.ndarray`): The array to convert.
- """
- if isinstance(arr, np.ndarray):
- return np.array
- if is_tf_available() and is_tf_tensor(arr):
- import tensorflow as tf
-
- return tf.convert_to_tensor
- if is_torch_available() and is_torch_tensor(arr):
- import torch
-
- return torch.tensor
- if is_flax_available() and is_jax_tensor(arr):
- import jax.numpy as jnp
-
- return jnp.array
- raise ValueError(f"Cannot convert arrays of type {type(arr)}")
-
-
-# Copied from transformers.models.detr.image_processing_detr.safe_squeeze
-def safe_squeeze(arr: np.ndarray, axis: Optional[int] = None) -> np.ndarray:
- """
- Squeezes an array, but only if the axis specified has dim 1.
- """
- if axis is None:
- return arr.squeeze()
-
- try:
- return arr.squeeze(axis=axis)
- except ValueError:
- return arr
-
-
-# Copied from transformers.models.detr.image_processing_detr.normalize_annotation
-def normalize_annotation(annotation: Dict, image_size: Tuple[int, int]) -> Dict:
- image_height, image_width = image_size
- norm_annotation = {}
- for key, value in annotation.items():
- if key == "boxes":
- boxes = value
- boxes = corners_to_center_format(boxes)
- boxes /= np.asarray([image_width, image_height, image_width, image_height], dtype=np.float32)
- norm_annotation[key] = boxes
- else:
- norm_annotation[key] = value
- return norm_annotation
-
-
-# Copied from transformers.models.detr.image_processing_detr.max_across_indices
-def max_across_indices(values: Iterable[Any]) -> List[Any]:
- """
- Return the maximum value across all indices of an iterable of values.
- """
- return [max(values_i) for values_i in zip(*values)]
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_max_height_width
-def get_max_height_width(
- images: List[np.ndarray], input_data_format: Optional[Union[str, ChannelDimension]] = None
-) -> List[int]:
- """
- Get the maximum height and width across all images in a batch.
- """
- if input_data_format is None:
- input_data_format = infer_channel_dimension_format(images[0])
-
- if input_data_format == ChannelDimension.FIRST:
- _, max_height, max_width = max_across_indices([img.shape for img in images])
- elif input_data_format == ChannelDimension.LAST:
- max_height, max_width, _ = max_across_indices([img.shape for img in images])
- else:
- raise ValueError(f"Invalid channel dimension format: {input_data_format}")
- return (max_height, max_width)
-
-
-# Copied from transformers.models.detr.image_processing_detr.make_pixel_mask
-def make_pixel_mask(
- image: np.ndarray, output_size: Tuple[int, int], input_data_format: Optional[Union[str, ChannelDimension]] = None
-) -> np.ndarray:
- """
- Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.
-
- Args:
- image (`np.ndarray`):
- Image to make the pixel mask for.
- output_size (`Tuple[int, int]`):
- Output size of the mask.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- mask = np.zeros(output_size, dtype=np.int64)
- mask[:input_height, :input_width] = 1
- return mask
-
-
-# Copied from transformers.models.detr.image_processing_detr.convert_coco_poly_to_mask
-def convert_coco_poly_to_mask(segmentations, height: int, width: int) -> np.ndarray:
- """
- Convert a COCO polygon annotation to a mask.
-
- Args:
- segmentations (`List[List[float]]`):
- List of polygons, each polygon represented by a list of x-y coordinates.
- height (`int`):
- Height of the mask.
- width (`int`):
- Width of the mask.
- """
- try:
- from pycocotools import mask as coco_mask
- except ImportError:
- raise ImportError("Pycocotools is not installed in your environment.")
-
- masks = []
- for polygons in segmentations:
- rles = coco_mask.frPyObjects(polygons, height, width)
- mask = coco_mask.decode(rles)
- if len(mask.shape) < 3:
- mask = mask[..., None]
- mask = np.asarray(mask, dtype=np.uint8)
- mask = np.any(mask, axis=2)
- masks.append(mask)
- if masks:
- masks = np.stack(masks, axis=0)
- else:
- masks = np.zeros((0, height, width), dtype=np.uint8)
-
- return masks
-
-
-# Copied from transformers.models.detr.image_processing_detr.prepare_coco_detection_annotation with DETR->DeformableDetr
-def prepare_coco_detection_annotation(
- image,
- target,
- return_segmentation_masks: bool = False,
- input_data_format: Optional[Union[ChannelDimension, str]] = None,
-):
- """
- Convert the target in COCO format into the format expected by DeformableDetr.
- """
- image_height, image_width = get_image_size(image, channel_dim=input_data_format)
-
- image_id = target["image_id"]
- image_id = np.asarray([image_id], dtype=np.int64)
-
- # Get all COCO annotations for the given image.
- annotations = target["annotations"]
- annotations = [obj for obj in annotations if "iscrowd" not in obj or obj["iscrowd"] == 0]
-
- classes = [obj["category_id"] for obj in annotations]
- classes = np.asarray(classes, dtype=np.int64)
-
- # for conversion to coco api
- area = np.asarray([obj["area"] for obj in annotations], dtype=np.float32)
- iscrowd = np.asarray([obj["iscrowd"] if "iscrowd" in obj else 0 for obj in annotations], dtype=np.int64)
-
- boxes = [obj["bbox"] for obj in annotations]
- # guard against no boxes via resizing
- boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4)
- boxes[:, 2:] += boxes[:, :2]
- boxes[:, 0::2] = boxes[:, 0::2].clip(min=0, max=image_width)
- boxes[:, 1::2] = boxes[:, 1::2].clip(min=0, max=image_height)
-
- keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
-
- new_target = {}
- new_target["image_id"] = image_id
- new_target["class_labels"] = classes[keep]
- new_target["boxes"] = boxes[keep]
- new_target["area"] = area[keep]
- new_target["iscrowd"] = iscrowd[keep]
- new_target["orig_size"] = np.asarray([int(image_height), int(image_width)], dtype=np.int64)
-
- if annotations and "keypoints" in annotations[0]:
- keypoints = [obj["keypoints"] for obj in annotations]
- # Converting the filtered keypoints list to a numpy array
- keypoints = np.asarray(keypoints, dtype=np.float32)
- # Apply the keep mask here to filter the relevant annotations
- keypoints = keypoints[keep]
- num_keypoints = keypoints.shape[0]
- keypoints = keypoints.reshape((-1, 3)) if num_keypoints else keypoints
- new_target["keypoints"] = keypoints
-
- if return_segmentation_masks:
- segmentation_masks = [obj["segmentation"] for obj in annotations]
- masks = convert_coco_poly_to_mask(segmentation_masks, image_height, image_width)
- new_target["masks"] = masks[keep]
-
- return new_target
-
-
-# Copied from transformers.models.detr.image_processing_detr.masks_to_boxes
-def masks_to_boxes(masks: np.ndarray) -> np.ndarray:
- """
- Compute the bounding boxes around the provided panoptic segmentation masks.
-
- Args:
- masks: masks in format `[number_masks, height, width]` where N is the number of masks
-
- Returns:
- boxes: bounding boxes in format `[number_masks, 4]` in xyxy format
- """
- if masks.size == 0:
- return np.zeros((0, 4))
-
- h, w = masks.shape[-2:]
- y = np.arange(0, h, dtype=np.float32)
- x = np.arange(0, w, dtype=np.float32)
- # see https://github.com/pytorch/pytorch/issues/50276
- y, x = np.meshgrid(y, x, indexing="ij")
-
- x_mask = masks * np.expand_dims(x, axis=0)
- x_max = x_mask.reshape(x_mask.shape[0], -1).max(-1)
- x = np.ma.array(x_mask, mask=~(np.array(masks, dtype=bool)))
- x_min = x.filled(fill_value=1e8)
- x_min = x_min.reshape(x_min.shape[0], -1).min(-1)
-
- y_mask = masks * np.expand_dims(y, axis=0)
- y_max = y_mask.reshape(x_mask.shape[0], -1).max(-1)
- y = np.ma.array(y_mask, mask=~(np.array(masks, dtype=bool)))
- y_min = y.filled(fill_value=1e8)
- y_min = y_min.reshape(y_min.shape[0], -1).min(-1)
-
- return np.stack([x_min, y_min, x_max, y_max], 1)
-
-
-# Copied from transformers.models.detr.image_processing_detr.prepare_coco_panoptic_annotation with DETR->DeformableDetr
-def prepare_coco_panoptic_annotation(
- image: np.ndarray,
- target: Dict,
- masks_path: Union[str, pathlib.Path],
- return_masks: bool = True,
- input_data_format: Union[ChannelDimension, str] = None,
-) -> Dict:
- """
- Prepare a coco panoptic annotation for DeformableDetr.
- """
- image_height, image_width = get_image_size(image, channel_dim=input_data_format)
- annotation_path = pathlib.Path(masks_path) / target["file_name"]
-
- new_target = {}
- new_target["image_id"] = np.asarray([target["image_id"] if "image_id" in target else target["id"]], dtype=np.int64)
- new_target["size"] = np.asarray([image_height, image_width], dtype=np.int64)
- new_target["orig_size"] = np.asarray([image_height, image_width], dtype=np.int64)
-
- if "segments_info" in target:
- masks = np.asarray(PIL.Image.open(annotation_path), dtype=np.uint32)
- masks = rgb_to_id(masks)
-
- ids = np.array([segment_info["id"] for segment_info in target["segments_info"]])
- masks = masks == ids[:, None, None]
- masks = masks.astype(np.uint8)
- if return_masks:
- new_target["masks"] = masks
- new_target["boxes"] = masks_to_boxes(masks)
- new_target["class_labels"] = np.array(
- [segment_info["category_id"] for segment_info in target["segments_info"]], dtype=np.int64
- )
- new_target["iscrowd"] = np.asarray(
- [segment_info["iscrowd"] for segment_info in target["segments_info"]], dtype=np.int64
- )
- new_target["area"] = np.asarray(
- [segment_info["area"] for segment_info in target["segments_info"]], dtype=np.float32
- )
-
- return new_target
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_segmentation_image
-def get_segmentation_image(
- masks: np.ndarray, input_size: Tuple, target_size: Tuple, stuff_equiv_classes, deduplicate=False
-):
- h, w = input_size
- final_h, final_w = target_size
-
- m_id = scipy.special.softmax(masks.transpose(0, 1), -1)
-
- if m_id.shape[-1] == 0:
- # We didn't detect any mask :(
- m_id = np.zeros((h, w), dtype=np.int64)
- else:
- m_id = m_id.argmax(-1).reshape(h, w)
-
- if deduplicate:
- # Merge the masks corresponding to the same stuff class
- for equiv in stuff_equiv_classes.values():
- for eq_id in equiv:
- m_id[m_id == eq_id] = equiv[0]
-
- seg_img = id_to_rgb(m_id)
- seg_img = resize(seg_img, (final_w, final_h), resample=PILImageResampling.NEAREST)
- return seg_img
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_mask_area
-def get_mask_area(seg_img: np.ndarray, target_size: Tuple[int, int], n_classes: int) -> np.ndarray:
- final_h, final_w = target_size
- np_seg_img = seg_img.astype(np.uint8)
- np_seg_img = np_seg_img.reshape(final_h, final_w, 3)
- m_id = rgb_to_id(np_seg_img)
- area = [(m_id == i).sum() for i in range(n_classes)]
- return area
-
-
-# Copied from transformers.models.detr.image_processing_detr.score_labels_from_class_probabilities
-def score_labels_from_class_probabilities(logits: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
- probs = scipy.special.softmax(logits, axis=-1)
- labels = probs.argmax(-1, keepdims=True)
- scores = np.take_along_axis(probs, labels, axis=-1)
- scores, labels = scores.squeeze(-1), labels.squeeze(-1)
- return scores, labels
-
-
-# Copied from transformers.models.detr.image_processing_detr.post_process_panoptic_sample
-def post_process_panoptic_sample(
- out_logits: np.ndarray,
- masks: np.ndarray,
- boxes: np.ndarray,
- processed_size: Tuple[int, int],
- target_size: Tuple[int, int],
- is_thing_map: Dict,
- threshold=0.85,
-) -> Dict:
- """
- Converts the output of [`DetrForSegmentation`] into panoptic segmentation predictions for a single sample.
-
- Args:
- out_logits (`torch.Tensor`):
- The logits for this sample.
- masks (`torch.Tensor`):
- The predicted segmentation masks for this sample.
- boxes (`torch.Tensor`):
- The prediced bounding boxes for this sample. The boxes are in the normalized format `(center_x, center_y,
- width, height)` and values between `[0, 1]`, relative to the size the image (disregarding padding).
- processed_size (`Tuple[int, int]`):
- The processed size of the image `(height, width)`, as returned by the preprocessing step i.e. the size
- after data augmentation but before batching.
- target_size (`Tuple[int, int]`):
- The target size of the image, `(height, width)` corresponding to the requested final size of the
- prediction.
- is_thing_map (`Dict`):
- A dictionary mapping class indices to a boolean value indicating whether the class is a thing or not.
- threshold (`float`, *optional*, defaults to 0.85):
- The threshold used to binarize the segmentation masks.
- """
- # we filter empty queries and detection below threshold
- scores, labels = score_labels_from_class_probabilities(out_logits)
- keep = (labels != out_logits.shape[-1] - 1) & (scores > threshold)
-
- cur_scores = scores[keep]
- cur_classes = labels[keep]
- cur_boxes = center_to_corners_format(boxes[keep])
-
- if len(cur_boxes) != len(cur_classes):
- raise ValueError("Not as many boxes as there are classes")
-
- cur_masks = masks[keep]
- cur_masks = resize(cur_masks[:, None], processed_size, resample=PILImageResampling.BILINEAR)
- cur_masks = safe_squeeze(cur_masks, 1)
- b, h, w = cur_masks.shape
-
- # It may be that we have several predicted masks for the same stuff class.
- # In the following, we track the list of masks ids for each stuff class (they are merged later on)
- cur_masks = cur_masks.reshape(b, -1)
- stuff_equiv_classes = defaultdict(list)
- for k, label in enumerate(cur_classes):
- if not is_thing_map[label]:
- stuff_equiv_classes[label].append(k)
-
- seg_img = get_segmentation_image(cur_masks, processed_size, target_size, stuff_equiv_classes, deduplicate=True)
- area = get_mask_area(cur_masks, processed_size, n_classes=len(cur_scores))
-
- # We filter out any mask that is too small
- if cur_classes.size() > 0:
- # We know filter empty masks as long as we find some
- filtered_small = np.array([a <= 4 for a in area], dtype=bool)
- while filtered_small.any():
- cur_masks = cur_masks[~filtered_small]
- cur_scores = cur_scores[~filtered_small]
- cur_classes = cur_classes[~filtered_small]
- seg_img = get_segmentation_image(cur_masks, (h, w), target_size, stuff_equiv_classes, deduplicate=True)
- area = get_mask_area(seg_img, target_size, n_classes=len(cur_scores))
- filtered_small = np.array([a <= 4 for a in area], dtype=bool)
- else:
- cur_classes = np.ones((1, 1), dtype=np.int64)
-
- segments_info = [
- {"id": i, "isthing": is_thing_map[cat], "category_id": int(cat), "area": a}
- for i, (cat, a) in enumerate(zip(cur_classes, area))
- ]
- del cur_classes
-
- with io.BytesIO() as out:
- PIL.Image.fromarray(seg_img).save(out, format="PNG")
- predictions = {"png_string": out.getvalue(), "segments_info": segments_info}
-
- return predictions
-
-
-# Copied from transformers.models.detr.image_processing_detr.resize_annotation
-def resize_annotation(
- annotation: Dict[str, Any],
- orig_size: Tuple[int, int],
- target_size: Tuple[int, int],
- threshold: float = 0.5,
- resample: PILImageResampling = PILImageResampling.NEAREST,
-):
- """
- Resizes an annotation to a target size.
-
- Args:
- annotation (`Dict[str, Any]`):
- The annotation dictionary.
- orig_size (`Tuple[int, int]`):
- The original size of the input image.
- target_size (`Tuple[int, int]`):
- The target size of the image, as returned by the preprocessing `resize` step.
- threshold (`float`, *optional*, defaults to 0.5):
- The threshold used to binarize the segmentation masks.
- resample (`PILImageResampling`, defaults to `PILImageResampling.NEAREST`):
- The resampling filter to use when resizing the masks.
- """
- ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(target_size, orig_size))
- ratio_height, ratio_width = ratios
-
- new_annotation = {}
- new_annotation["size"] = target_size
-
- for key, value in annotation.items():
- if key == "boxes":
- boxes = value
- scaled_boxes = boxes * np.asarray([ratio_width, ratio_height, ratio_width, ratio_height], dtype=np.float32)
- new_annotation["boxes"] = scaled_boxes
- elif key == "area":
- area = value
- scaled_area = area * (ratio_width * ratio_height)
- new_annotation["area"] = scaled_area
- elif key == "masks":
- masks = value[:, None]
- masks = np.array([resize(mask, target_size, resample=resample) for mask in masks])
- masks = masks.astype(np.float32)
- masks = masks[:, 0] > threshold
- new_annotation["masks"] = masks
- elif key == "size":
- new_annotation["size"] = target_size
- else:
- new_annotation[key] = value
-
- return new_annotation
-
-
-# Copied from transformers.models.detr.image_processing_detr.binary_mask_to_rle
-def binary_mask_to_rle(mask):
- """
- Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format.
-
- Args:
- mask (`torch.Tensor` or `numpy.array`):
- A binary mask tensor of shape `(height, width)` where 0 denotes background and 1 denotes the target
- segment_id or class_id.
- Returns:
- `List`: Run-length encoded list of the binary mask. Refer to COCO API for more information about the RLE
- format.
- """
- if is_torch_tensor(mask):
- mask = mask.numpy()
-
- pixels = mask.flatten()
- pixels = np.concatenate([[0], pixels, [0]])
- runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
- runs[1::2] -= runs[::2]
- return list(runs)
-
-
-# Copied from transformers.models.detr.image_processing_detr.convert_segmentation_to_rle
-def convert_segmentation_to_rle(segmentation):
- """
- Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format.
-
- Args:
- segmentation (`torch.Tensor` or `numpy.array`):
- A segmentation map of shape `(height, width)` where each value denotes a segment or class id.
- Returns:
- `List[List]`: A list of lists, where each list is the run-length encoding of a segment / class id.
- """
- segment_ids = torch.unique(segmentation)
-
- run_length_encodings = []
- for idx in segment_ids:
- mask = torch.where(segmentation == idx, 1, 0)
- rle = binary_mask_to_rle(mask)
- run_length_encodings.append(rle)
-
- return run_length_encodings
-
-
-# Copied from transformers.models.detr.image_processing_detr.remove_low_and_no_objects
-def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels):
- """
- Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and
- `labels`.
-
- Args:
- masks (`torch.Tensor`):
- A tensor of shape `(num_queries, height, width)`.
- scores (`torch.Tensor`):
- A tensor of shape `(num_queries)`.
- labels (`torch.Tensor`):
- A tensor of shape `(num_queries)`.
- object_mask_threshold (`float`):
- A number between 0 and 1 used to binarize the masks.
- Raises:
- `ValueError`: Raised when the first dimension doesn't match in all input tensors.
- Returns:
- `Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region
- < `object_mask_threshold`.
- """
- if not (masks.shape[0] == scores.shape[0] == labels.shape[0]):
- raise ValueError("mask, scores and labels must have the same shape!")
-
- to_keep = labels.ne(num_labels) & (scores > object_mask_threshold)
-
- return masks[to_keep], scores[to_keep], labels[to_keep]
-
-
-# Copied from transformers.models.detr.image_processing_detr.check_segment_validity
-def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8):
- # Get the mask associated with the k class
- mask_k = mask_labels == k
- mask_k_area = mask_k.sum()
-
- # Compute the area of all the stuff in query k
- original_area = (mask_probs[k] >= mask_threshold).sum()
- mask_exists = mask_k_area > 0 and original_area > 0
-
- # Eliminate disconnected tiny segments
- if mask_exists:
- area_ratio = mask_k_area / original_area
- if not area_ratio.item() > overlap_mask_area_threshold:
- mask_exists = False
-
- return mask_exists, mask_k
-
-
-# Copied from transformers.models.detr.image_processing_detr.compute_segments
-def compute_segments(
- mask_probs,
- pred_scores,
- pred_labels,
- mask_threshold: float = 0.5,
- overlap_mask_area_threshold: float = 0.8,
- label_ids_to_fuse: Optional[Set[int]] = None,
- target_size: Tuple[int, int] = None,
-):
- height = mask_probs.shape[1] if target_size is None else target_size[0]
- width = mask_probs.shape[2] if target_size is None else target_size[1]
-
- segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device)
- segments: List[Dict] = []
-
- if target_size is not None:
- mask_probs = nn.functional.interpolate(
- mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False
- )[0]
-
- current_segment_id = 0
-
- # Weigh each mask by its prediction score
- mask_probs *= pred_scores.view(-1, 1, 1)
- mask_labels = mask_probs.argmax(0) # [height, width]
-
- # Keep track of instances of each class
- stuff_memory_list: Dict[str, int] = {}
- for k in range(pred_labels.shape[0]):
- pred_class = pred_labels[k].item()
- should_fuse = pred_class in label_ids_to_fuse
-
- # Check if mask exists and large enough to be a segment
- mask_exists, mask_k = check_segment_validity(
- mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold
- )
-
- if mask_exists:
- if pred_class in stuff_memory_list:
- current_segment_id = stuff_memory_list[pred_class]
- else:
- current_segment_id += 1
-
- # Add current object segment to final segmentation map
- segmentation[mask_k] = current_segment_id
- segment_score = round(pred_scores[k].item(), 6)
- segments.append(
- {
- "id": current_segment_id,
- "label_id": pred_class,
- "was_fused": should_fuse,
- "score": segment_score,
- }
- )
- if should_fuse:
- stuff_memory_list[pred_class] = current_segment_id
-
- return segmentation, segments
-
-
-class DeformableDetrImageProcessor(BaseImageProcessor):
- r"""
- Constructs a Deformable DETR image processor.
-
- Args:
- format (`str`, *optional*, defaults to `"coco_detection"`):
- Data format of the annotations. One of "coco_detection" or "coco_panoptic".
- do_resize (`bool`, *optional*, defaults to `True`):
- Controls whether to resize the image's (height, width) dimensions to the specified `size`. Can be
- overridden by the `do_resize` parameter in the `preprocess` method.
- size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 800, "longest_edge": 1333}`):
- Size of the image's (height, width) dimensions after resizing. Can be overridden by the `size` parameter in
- the `preprocess` method.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Controls whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the
- `do_rescale` parameter in the `preprocess` method.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the
- `preprocess` method.
- do_normalize:
- Controls whether to normalize the image. Can be overridden by the `do_normalize` parameter in the
- `preprocess` method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_DEFAULT_MEAN`):
- Mean values to use when normalizing the image. Can be a single value or a list of values, one for each
- channel. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_DEFAULT_STD`):
- Standard deviation values to use when normalizing the image. Can be a single value or a list of values, one
- for each channel. Can be overridden by the `image_std` parameter in the `preprocess` method.
- do_convert_annotations (`bool`, *optional*, defaults to `True`):
- Controls whether to convert the annotations to the format expected by the DETR model. Converts the
- bounding boxes to the format `(center_x, center_y, width, height)` and in the range `[0, 1]`.
- Can be overridden by the `do_convert_annotations` parameter in the `preprocess` method.
- do_pad (`bool`, *optional*, defaults to `True`):
- Controls whether to pad the image. Can be overridden by the `do_pad` parameter in the `preprocess`
- method. If `True` will pad the images in the batch to the largest height and width in the batch.
- Padding will be applied to the bottom and right of the image with zeros.
- """
-
- model_input_names = ["pixel_values", "pixel_mask"]
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.__init__
- def __init__(
- self,
- format: Union[str, AnnotationFormat] = AnnotationFormat.COCO_DETECTION,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- do_normalize: bool = True,
- image_mean: Union[float, List[float]] = None,
- image_std: Union[float, List[float]] = None,
- do_convert_annotations: Optional[bool] = None,
- do_pad: bool = True,
- **kwargs,
- ) -> None:
- if "pad_and_return_pixel_mask" in kwargs:
- do_pad = kwargs.pop("pad_and_return_pixel_mask")
-
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` parameter is deprecated and will be removed in v4.26. "
- "Please specify in `size['longest_edge'] instead`.",
- )
- max_size = kwargs.pop("max_size")
- else:
- max_size = None if size is None else 1333
-
- size = size if size is not None else {"shortest_edge": 800, "longest_edge": 1333}
- size = get_size_dict(size, max_size=max_size, default_to_square=False)
-
- # Backwards compatibility
- if do_convert_annotations is None:
- do_convert_annotations = do_normalize
-
- super().__init__(**kwargs)
- self.format = format
- self.do_resize = do_resize
- self.size = size
- self.resample = resample
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_normalize = do_normalize
- self.do_convert_annotations = do_convert_annotations
- self.image_mean = image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
- self.do_pad = do_pad
- self._valid_processor_keys = [
- "images",
- "annotations",
- "return_segmentation_masks",
- "masks_path",
- "do_resize",
- "size",
- "resample",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "do_convert_annotations",
- "image_mean",
- "image_std",
- "do_pad",
- "format",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- @classmethod
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.from_dict with Detr->DeformableDetr
- def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
- """
- Overrides the `from_dict` method from the base class to make sure parameters are updated if image processor is
- created using from_dict and kwargs e.g. `DeformableDetrImageProcessor.from_pretrained(checkpoint, size=600,
- max_size=800)`
- """
- image_processor_dict = image_processor_dict.copy()
- if "max_size" in kwargs:
- image_processor_dict["max_size"] = kwargs.pop("max_size")
- if "pad_and_return_pixel_mask" in kwargs:
- image_processor_dict["pad_and_return_pixel_mask"] = kwargs.pop("pad_and_return_pixel_mask")
- return super().from_dict(image_processor_dict, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_annotation with DETR->DeformableDetr
- def prepare_annotation(
- self,
- image: np.ndarray,
- target: Dict,
- format: Optional[AnnotationFormat] = None,
- return_segmentation_masks: bool = None,
- masks_path: Optional[Union[str, pathlib.Path]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> Dict:
- """
- Prepare an annotation for feeding into DeformableDetr model.
- """
- format = format if format is not None else self.format
-
- if format == AnnotationFormat.COCO_DETECTION:
- return_segmentation_masks = False if return_segmentation_masks is None else return_segmentation_masks
- target = prepare_coco_detection_annotation(
- image, target, return_segmentation_masks, input_data_format=input_data_format
- )
- elif format == AnnotationFormat.COCO_PANOPTIC:
- return_segmentation_masks = True if return_segmentation_masks is None else return_segmentation_masks
- target = prepare_coco_panoptic_annotation(
- image,
- target,
- masks_path=masks_path,
- return_masks=return_segmentation_masks,
- input_data_format=input_data_format,
- )
- else:
- raise ValueError(f"Format {format} is not supported.")
- return target
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare
- def prepare(self, image, target, return_segmentation_masks=None, masks_path=None):
- logger.warning_once(
- "The `prepare` method is deprecated and will be removed in a v4.33. "
- "Please use `prepare_annotation` instead. Note: the `prepare_annotation` method "
- "does not return the image anymore.",
- )
- target = self.prepare_annotation(image, target, return_segmentation_masks, masks_path, self.format)
- return image, target
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.convert_coco_poly_to_mask
- def convert_coco_poly_to_mask(self, *args, **kwargs):
- logger.warning_once("The `convert_coco_poly_to_mask` method is deprecated and will be removed in v4.33. ")
- return convert_coco_poly_to_mask(*args, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_coco_detection
- def prepare_coco_detection(self, *args, **kwargs):
- logger.warning_once("The `prepare_coco_detection` method is deprecated and will be removed in v4.33. ")
- return prepare_coco_detection_annotation(*args, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_coco_panoptic
- def prepare_coco_panoptic(self, *args, **kwargs):
- logger.warning_once("The `prepare_coco_panoptic` method is deprecated and will be removed in v4.33. ")
- return prepare_coco_panoptic_annotation(*args, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.resize
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize the image to the given size. Size can be `min_size` (scalar) or `(height, width)` tuple. If size is an
- int, smaller edge of the image will be matched to this number.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary containing the size to resize to. Can contain the keys `shortest_edge` and `longest_edge` or
- `height` and `width`.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` parameter is deprecated and will be removed in v4.26. "
- "Please specify in `size['longest_edge'] instead`.",
- )
- max_size = kwargs.pop("max_size")
- else:
- max_size = None
- size = get_size_dict(size, max_size=max_size, default_to_square=False)
- if "shortest_edge" in size and "longest_edge" in size:
- size = get_resize_output_image_size(
- image, size["shortest_edge"], size["longest_edge"], input_data_format=input_data_format
- )
- elif "height" in size and "width" in size:
- size = (size["height"], size["width"])
- else:
- raise ValueError(
- "Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got"
- f" {size.keys()}."
- )
- image = resize(
- image, size=size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs
- )
- return image
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.resize_annotation
- def resize_annotation(
- self,
- annotation,
- orig_size,
- size,
- resample: PILImageResampling = PILImageResampling.NEAREST,
- ) -> Dict:
- """
- Resize the annotation to match the resized image. If size is an int, smaller edge of the mask will be matched
- to this number.
- """
- return resize_annotation(annotation, orig_size=orig_size, target_size=size, resample=resample)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.rescale
- def rescale(
- self,
- image: np.ndarray,
- rescale_factor: float,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.ndarray:
- """
- Rescale the image by the given factor. image = image * rescale_factor.
-
- Args:
- image (`np.ndarray`):
- Image to rescale.
- rescale_factor (`float`):
- The value to use for rescaling.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the input image. If unset, is inferred from the input image. Can be
- one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- """
- return rescale(image, rescale_factor, data_format=data_format, input_data_format=input_data_format)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.normalize_annotation
- def normalize_annotation(self, annotation: Dict, image_size: Tuple[int, int]) -> Dict:
- """
- Normalize the boxes in the annotation from `[top_left_x, top_left_y, bottom_right_x, bottom_right_y]` to
- `[center_x, center_y, width, height]` format and from absolute to relative pixel values.
- """
- return normalize_annotation(annotation, image_size=image_size)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor._update_annotation_for_padded_image
- def _update_annotation_for_padded_image(
- self,
- annotation: Dict,
- input_image_size: Tuple[int, int],
- output_image_size: Tuple[int, int],
- padding,
- update_bboxes,
- ) -> Dict:
- """
- Update the annotation for a padded image.
- """
- new_annotation = {}
- new_annotation["size"] = output_image_size
-
- for key, value in annotation.items():
- if key == "masks":
- masks = value
- masks = pad(
- masks,
- padding,
- mode=PaddingMode.CONSTANT,
- constant_values=0,
- input_data_format=ChannelDimension.FIRST,
- )
- masks = safe_squeeze(masks, 1)
- new_annotation["masks"] = masks
- elif key == "boxes" and update_bboxes:
- boxes = value
- boxes *= np.asarray(
- [
- input_image_size[1] / output_image_size[1],
- input_image_size[0] / output_image_size[0],
- input_image_size[1] / output_image_size[1],
- input_image_size[0] / output_image_size[0],
- ]
- )
- new_annotation["boxes"] = boxes
- elif key == "size":
- new_annotation["size"] = output_image_size
- else:
- new_annotation[key] = value
- return new_annotation
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor._pad_image
- def _pad_image(
- self,
- image: np.ndarray,
- output_size: Tuple[int, int],
- annotation: Optional[Dict[str, Any]] = None,
- constant_values: Union[float, Iterable[float]] = 0,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- update_bboxes: bool = True,
- ) -> np.ndarray:
- """
- Pad an image with zeros to the given size.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- output_height, output_width = output_size
-
- pad_bottom = output_height - input_height
- pad_right = output_width - input_width
- padding = ((0, pad_bottom), (0, pad_right))
- padded_image = pad(
- image,
- padding,
- mode=PaddingMode.CONSTANT,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- )
- if annotation is not None:
- annotation = self._update_annotation_for_padded_image(
- annotation, (input_height, input_width), (output_height, output_width), padding, update_bboxes
- )
- return padded_image, annotation
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.pad
- def pad(
- self,
- images: List[np.ndarray],
- annotations: Optional[Union[AnnotationType, List[AnnotationType]]] = None,
- constant_values: Union[float, Iterable[float]] = 0,
- return_pixel_mask: bool = True,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- update_bboxes: bool = True,
- ) -> BatchFeature:
- """
- Pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width
- in the batch and optionally returns their corresponding pixel mask.
-
- Args:
- images (List[`np.ndarray`]):
- Images to pad.
- annotations (`AnnotationType` or `List[AnnotationType]`, *optional*):
- Annotations to transform according to the padding that is applied to the images.
- constant_values (`float` or `Iterable[float]`, *optional*):
- The value to use for the padding if `mode` is `"constant"`.
- return_pixel_mask (`bool`, *optional*, defaults to `True`):
- Whether to return a pixel mask.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the image. If not provided, it will be the same as the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- update_bboxes (`bool`, *optional*, defaults to `True`):
- Whether to update the bounding boxes in the annotations to match the padded images. If the
- bounding boxes have not been converted to relative coordinates and `(centre_x, centre_y, width, height)`
- format, the bounding boxes will not be updated.
- """
- pad_size = get_max_height_width(images, input_data_format=input_data_format)
-
- annotation_list = annotations if annotations is not None else [None] * len(images)
- padded_images = []
- padded_annotations = []
- for image, annotation in zip(images, annotation_list):
- padded_image, padded_annotation = self._pad_image(
- image,
- pad_size,
- annotation,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- update_bboxes=update_bboxes,
- )
- padded_images.append(padded_image)
- padded_annotations.append(padded_annotation)
-
- data = {"pixel_values": padded_images}
-
- if return_pixel_mask:
- masks = [
- make_pixel_mask(image=image, output_size=pad_size, input_data_format=input_data_format)
- for image in images
- ]
- data["pixel_mask"] = masks
-
- encoded_inputs = BatchFeature(data=data, tensor_type=return_tensors)
-
- if annotations is not None:
- encoded_inputs["labels"] = [
- BatchFeature(annotation, tensor_type=return_tensors) for annotation in padded_annotations
- ]
-
- return encoded_inputs
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.preprocess
- def preprocess(
- self,
- images: ImageInput,
- annotations: Optional[Union[AnnotationType, List[AnnotationType]]] = None,
- return_segmentation_masks: bool = None,
- masks_path: Optional[Union[str, pathlib.Path]] = None,
- do_resize: Optional[bool] = None,
- size: Optional[Dict[str, int]] = None,
- resample=None, # PILImageResampling
- do_rescale: Optional[bool] = None,
- rescale_factor: Optional[Union[int, float]] = None,
- do_normalize: Optional[bool] = None,
- do_convert_annotations: Optional[bool] = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- do_pad: Optional[bool] = None,
- format: Optional[Union[str, AnnotationFormat]] = None,
- return_tensors: Optional[Union[TensorType, str]] = None,
- data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> BatchFeature:
- """
- Preprocess an image or a batch of images so that it can be used by the model.
-
- Args:
- images (`ImageInput`):
- Image or batch of images to preprocess. Expects a single or batch of images with pixel values ranging
- from 0 to 255. If passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- annotations (`AnnotationType` or `List[AnnotationType]`, *optional*):
- List of annotations associated with the image or batch of images. If annotation is for object
- detection, the annotations should be a dictionary with the following keys:
- - "image_id" (`int`): The image id.
- - "annotations" (`List[Dict]`): List of annotations for an image. Each annotation should be a
- dictionary. An image can have no annotations, in which case the list should be empty.
- If annotation is for segmentation, the annotations should be a dictionary with the following keys:
- - "image_id" (`int`): The image id.
- - "segments_info" (`List[Dict]`): List of segments for an image. Each segment should be a dictionary.
- An image can have no segments, in which case the list should be empty.
- - "file_name" (`str`): The file name of the image.
- return_segmentation_masks (`bool`, *optional*, defaults to self.return_segmentation_masks):
- Whether to return segmentation masks.
- masks_path (`str` or `pathlib.Path`, *optional*):
- Path to the directory containing the segmentation masks.
- do_resize (`bool`, *optional*, defaults to self.do_resize):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to self.size):
- Size of the image after resizing.
- resample (`PILImageResampling`, *optional*, defaults to self.resample):
- Resampling filter to use when resizing the image.
- do_rescale (`bool`, *optional*, defaults to self.do_rescale):
- Whether to rescale the image.
- rescale_factor (`float`, *optional*, defaults to self.rescale_factor):
- Rescale factor to use when rescaling the image.
- do_normalize (`bool`, *optional*, defaults to self.do_normalize):
- Whether to normalize the image.
- do_convert_annotations (`bool`, *optional*, defaults to self.do_convert_annotations):
- Whether to convert the annotations to the format expected by the model. Converts the bounding
- boxes from the format `(top_left_x, top_left_y, width, height)` to `(center_x, center_y, width, height)`
- and in relative coordinates.
- image_mean (`float` or `List[float]`, *optional*, defaults to self.image_mean):
- Mean to use when normalizing the image.
- image_std (`float` or `List[float]`, *optional*, defaults to self.image_std):
- Standard deviation to use when normalizing the image.
- do_pad (`bool`, *optional*, defaults to self.do_pad):
- Whether to pad the image. If `True` will pad the images in the batch to the largest image in the batch
- and create a pixel mask. Padding will be applied to the bottom and right of the image with zeros.
- format (`str` or `AnnotationFormat`, *optional*, defaults to self.format):
- Format of the annotations.
- return_tensors (`str` or `TensorType`, *optional*, defaults to self.return_tensors):
- Type of tensors to return. If `None`, will return the list of images.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - Unset: Use the channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- if "pad_and_return_pixel_mask" in kwargs:
- logger.warning_once(
- "The `pad_and_return_pixel_mask` argument is deprecated and will be removed in a future version, "
- "use `do_pad` instead."
- )
- do_pad = kwargs.pop("pad_and_return_pixel_mask")
-
- max_size = None
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` argument is deprecated and will be removed in a future version, use"
- " `size['longest_edge']` instead."
- )
- size = kwargs.pop("max_size")
-
- do_resize = self.do_resize if do_resize is None else do_resize
- size = self.size if size is None else size
- size = get_size_dict(size=size, max_size=max_size, default_to_square=False)
- resample = self.resample if resample is None else resample
- do_rescale = self.do_rescale if do_rescale is None else do_rescale
- rescale_factor = self.rescale_factor if rescale_factor is None else rescale_factor
- do_normalize = self.do_normalize if do_normalize is None else do_normalize
- image_mean = self.image_mean if image_mean is None else image_mean
- image_std = self.image_std if image_std is None else image_std
- do_convert_annotations = (
- self.do_convert_annotations if do_convert_annotations is None else do_convert_annotations
- )
- do_pad = self.do_pad if do_pad is None else do_pad
- format = self.format if format is None else format
-
- images = make_list_of_images(images)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- # Here, the pad() method pads to the maximum of (width, height). It does not need to be validated.
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
-
- if annotations is not None and isinstance(annotations, dict):
- annotations = [annotations]
-
- if annotations is not None and len(images) != len(annotations):
- raise ValueError(
- f"The number of images ({len(images)}) and annotations ({len(annotations)}) do not match."
- )
-
- format = AnnotationFormat(format)
- if annotations is not None:
- validate_annotations(format, SUPPORTED_ANNOTATION_FORMATS, annotations)
-
- if (
- masks_path is not None
- and format == AnnotationFormat.COCO_PANOPTIC
- and not isinstance(masks_path, (pathlib.Path, str))
- ):
- raise ValueError(
- "The path to the directory containing the mask PNG files should be provided as a"
- f" `pathlib.Path` or string object, but is {type(masks_path)} instead."
- )
-
- # All transformations expect numpy arrays
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- # prepare (COCO annotations as a list of Dict -> DETR target as a single Dict per image)
- if annotations is not None:
- prepared_images = []
- prepared_annotations = []
- for image, target in zip(images, annotations):
- target = self.prepare_annotation(
- image,
- target,
- format,
- return_segmentation_masks=return_segmentation_masks,
- masks_path=masks_path,
- input_data_format=input_data_format,
- )
- prepared_images.append(image)
- prepared_annotations.append(target)
- images = prepared_images
- annotations = prepared_annotations
- del prepared_images, prepared_annotations
-
- # transformations
- if do_resize:
- if annotations is not None:
- resized_images, resized_annotations = [], []
- for image, target in zip(images, annotations):
- orig_size = get_image_size(image, input_data_format)
- resized_image = self.resize(
- image, size=size, max_size=max_size, resample=resample, input_data_format=input_data_format
- )
- resized_annotation = self.resize_annotation(
- target, orig_size, get_image_size(resized_image, input_data_format)
- )
- resized_images.append(resized_image)
- resized_annotations.append(resized_annotation)
- images = resized_images
- annotations = resized_annotations
- del resized_images, resized_annotations
- else:
- images = [
- self.resize(image, size=size, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_rescale:
- images = [self.rescale(image, rescale_factor, input_data_format=input_data_format) for image in images]
-
- if do_normalize:
- images = [
- self.normalize(image, image_mean, image_std, input_data_format=input_data_format) for image in images
- ]
-
- if do_convert_annotations and annotations is not None:
- annotations = [
- self.normalize_annotation(annotation, get_image_size(image, input_data_format))
- for annotation, image in zip(annotations, images)
- ]
-
- if do_pad:
- # Pads images and returns their mask: {'pixel_values': ..., 'pixel_mask': ...}
- encoded_inputs = self.pad(
- images,
- annotations=annotations,
- return_pixel_mask=True,
- data_format=data_format,
- input_data_format=input_data_format,
- update_bboxes=do_convert_annotations,
- return_tensors=return_tensors,
- )
- else:
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
- for image in images
- ]
- encoded_inputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)
- if annotations is not None:
- encoded_inputs["labels"] = [
- BatchFeature(annotation, tensor_type=return_tensors) for annotation in annotations
- ]
-
- return encoded_inputs
-
- # POSTPROCESSING METHODS - TODO: add support for other frameworks
- def post_process(self, outputs, target_sizes):
- """
- Converts the raw output of [`DeformableDetrForObjectDetection`] into final bounding boxes in (top_left_x,
- top_left_y, bottom_right_x, bottom_right_y) format. Only supports PyTorch.
-
- Args:
- outputs ([`DeformableDetrObjectDetectionOutput`]):
- Raw outputs of the model.
- target_sizes (`torch.Tensor` of shape `(batch_size, 2)`):
- Tensor containing the size (height, width) of each image of the batch. For evaluation, this must be the
- original image size (before any data augmentation). For visualization, this should be the image size
- after data augment, but before padding.
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
- in the batch as predicted by the model.
- """
- logger.warning_once(
- "`post_process` is deprecated and will be removed in v5 of Transformers, please use"
- " `post_process_object_detection` instead, with `threshold=0.` for equivalent results.",
- )
-
- out_logits, out_bbox = outputs.logits, outputs.pred_boxes
-
- if len(out_logits) != len(target_sizes):
- raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the logits")
- if target_sizes.shape[1] != 2:
- raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch")
-
- prob = out_logits.sigmoid()
- topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), 100, dim=1)
- scores = topk_values
- topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode="floor")
- labels = topk_indexes % out_logits.shape[2]
- boxes = center_to_corners_format(out_bbox)
- boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
-
- # and from relative [0, 1] to absolute [0, height] coordinates
- img_h, img_w = target_sizes.unbind(1)
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
- boxes = boxes * scale_fct[:, None, :]
-
- results = [{"scores": s, "labels": l, "boxes": b} for s, l, b in zip(scores, labels, boxes)]
-
- return results
-
- def post_process_object_detection(
- self, outputs, threshold: float = 0.5, target_sizes: Union[TensorType, List[Tuple]] = None, top_k: int = 100
- ):
- """
- Converts the raw output of [`DeformableDetrForObjectDetection`] into final bounding boxes in (top_left_x,
- top_left_y, bottom_right_x, bottom_right_y) format. Only supports PyTorch.
-
- Args:
- outputs ([`DetrObjectDetectionOutput`]):
- Raw outputs of the model.
- threshold (`float`, *optional*):
- Score threshold to keep object detection predictions.
- target_sizes (`torch.Tensor` or `List[Tuple[int, int]]`, *optional*):
- Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
- (height, width) of each image in the batch. If left to None, predictions will not be resized.
- top_k (`int`, *optional*, defaults to 100):
- Keep only top k bounding boxes before filtering by thresholding.
-
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
- in the batch as predicted by the model.
- """
- out_logits, out_bbox = outputs.logits, outputs.pred_boxes
-
- if target_sizes is not None:
- if len(out_logits) != len(target_sizes):
- raise ValueError(
- "Make sure that you pass in as many target sizes as the batch dimension of the logits"
- )
-
- prob = out_logits.sigmoid()
- prob = prob.view(out_logits.shape[0], -1)
- k_value = min(top_k, prob.size(1))
- topk_values, topk_indexes = torch.topk(prob, k_value, dim=1)
- scores = topk_values
- topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode="floor")
- labels = topk_indexes % out_logits.shape[2]
- boxes = center_to_corners_format(out_bbox)
- boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
-
- # and from relative [0, 1] to absolute [0, height] coordinates
- if target_sizes is not None:
- if isinstance(target_sizes, List):
- img_h = torch.Tensor([i[0] for i in target_sizes])
- img_w = torch.Tensor([i[1] for i in target_sizes])
- else:
- img_h, img_w = target_sizes.unbind(1)
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device)
- boxes = boxes * scale_fct[:, None, :]
-
- results = []
- for s, l, b in zip(scores, labels, boxes):
- score = s[s > threshold]
- label = l[s > threshold]
- box = b[s > threshold]
- results.append({"scores": score, "labels": label, "boxes": box})
-
- return results
diff --git a/transformers/models/deformable_detr/load_custom.py b/transformers/models/deformable_detr/load_custom.py
deleted file mode 100644
index c3a822e2764170c24c7098956e81788856385451..0000000000000000000000000000000000000000
--- a/transformers/models/deformable_detr/load_custom.py
+++ /dev/null
@@ -1,49 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Loading of Deformable DETR's CUDA kernels"""
-import os
-from pathlib import Path
-
-
-def load_cuda_kernels():
- from torch.utils.cpp_extension import load
-
- root = Path(__file__).resolve().parent.parent.parent / "kernels" / "deformable_detr"
- src_files = [
- root / filename
- for filename in [
- "vision.cpp",
- os.path.join("cpu", "ms_deform_attn_cpu.cpp"),
- os.path.join("cuda", "ms_deform_attn_cuda.cu"),
- ]
- ]
-
- load(
- "MultiScaleDeformableAttention",
- src_files,
- with_cuda=True,
- extra_include_paths=[str(root)],
- extra_cflags=["-DWITH_CUDA=1"],
- extra_cuda_cflags=[
- "-DCUDA_HAS_FP16=1",
- "-D__CUDA_NO_HALF_OPERATORS__",
- "-D__CUDA_NO_HALF_CONVERSIONS__",
- "-D__CUDA_NO_HALF2_OPERATORS__",
- ],
- )
-
- import MultiScaleDeformableAttention as MSDA
-
- return MSDA
diff --git a/transformers/models/deformable_detr/modeling_deformable_detr.py b/transformers/models/deformable_detr/modeling_deformable_detr.py
deleted file mode 100644
index c0ac7cffc7ab44578ab33463071d409bf628264f..0000000000000000000000000000000000000000
--- a/transformers/models/deformable_detr/modeling_deformable_detr.py
+++ /dev/null
@@ -1,2536 +0,0 @@
-# coding=utf-8
-# Copyright 2022 SenseTime and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Deformable DETR model."""
-
-
-import copy
-import math
-import os
-import warnings
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor, nn
-from torch.autograd import Function
-from torch.autograd.function import once_differentiable
-
-from ...activations import ACT2FN
-from ...file_utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_scipy_available,
- is_timm_available,
- is_torch_cuda_available,
- is_vision_available,
- replace_return_docstrings,
- requires_backends,
-)
-from ...modeling_attn_mask_utils import _prepare_4d_attention_mask
-from ...modeling_outputs import BaseModelOutput
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import meshgrid
-from ...utils import is_accelerate_available, is_ninja_available, logging
-from ...utils.backbone_utils import load_backbone
-from .configuration_deformable_detr import DeformableDetrConfig
-
-
-logger = logging.get_logger(__name__)
-
-MultiScaleDeformableAttention = None
-
-
-def load_cuda_kernels():
- from torch.utils.cpp_extension import load
-
- global MultiScaleDeformableAttention
-
- root = Path(__file__).resolve().parent.parent.parent / "kernels" / "deformable_detr"
- src_files = [
- root / filename
- for filename in [
- "vision.cpp",
- os.path.join("cpu", "ms_deform_attn_cpu.cpp"),
- os.path.join("cuda", "ms_deform_attn_cuda.cu"),
- ]
- ]
-
- MultiScaleDeformableAttention = load(
- "MultiScaleDeformableAttention",
- src_files,
- with_cuda=True,
- extra_include_paths=[str(root)],
- extra_cflags=["-DWITH_CUDA=1"],
- extra_cuda_cflags=[
- "-DCUDA_HAS_FP16=1",
- "-D__CUDA_NO_HALF_OPERATORS__",
- "-D__CUDA_NO_HALF_CONVERSIONS__",
- "-D__CUDA_NO_HALF2_OPERATORS__",
- ],
- )
-
-
-if is_vision_available():
- from transformers.image_transforms import center_to_corners_format
-
-if is_accelerate_available():
- from accelerate import PartialState
- from accelerate.utils import reduce
-
-
-class MultiScaleDeformableAttentionFunction(Function):
- @staticmethod
- def forward(
- context,
- value,
- value_spatial_shapes,
- value_level_start_index,
- sampling_locations,
- attention_weights,
- im2col_step,
- ):
- context.im2col_step = im2col_step
- output = MultiScaleDeformableAttention.ms_deform_attn_forward(
- value,
- value_spatial_shapes,
- value_level_start_index,
- sampling_locations,
- attention_weights,
- context.im2col_step,
- )
- context.save_for_backward(
- value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights
- )
- return output
-
- @staticmethod
- @once_differentiable
- def backward(context, grad_output):
- (
- value,
- value_spatial_shapes,
- value_level_start_index,
- sampling_locations,
- attention_weights,
- ) = context.saved_tensors
- grad_value, grad_sampling_loc, grad_attn_weight = MultiScaleDeformableAttention.ms_deform_attn_backward(
- value,
- value_spatial_shapes,
- value_level_start_index,
- sampling_locations,
- attention_weights,
- grad_output,
- context.im2col_step,
- )
-
- return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None
-
-
-if is_scipy_available():
- from scipy.optimize import linear_sum_assignment
-
-if is_timm_available():
- from timm import create_model
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "DeformableDetrConfig"
-_CHECKPOINT_FOR_DOC = "sensetime/deformable-detr"
-
-
-from ..deprecated._archive_maps import DEFORMABLE_DETR_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class DeformableDetrDecoderOutput(ModelOutput):
- """
- Base class for outputs of the DeformableDetrDecoder. This class adds two attributes to
- BaseModelOutputWithCrossAttentions, namely:
- - a stacked tensor of intermediate decoder hidden states (i.e. the output of each decoder layer)
- - a stacked tensor of intermediate reference points.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
- Stacked intermediate hidden states (output of each layer of the decoder).
- intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, hidden_size)`):
- Stacked intermediate reference points (reference points of each layer of the decoder).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
- the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- """
-
- last_hidden_state: torch.FloatTensor = None
- intermediate_hidden_states: torch.FloatTensor = None
- intermediate_reference_points: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class DeformableDetrModelOutput(ModelOutput):
- """
- Base class for outputs of the Deformable DETR encoder-decoder model.
-
- Args:
- init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Initial reference points sent through the Transformer decoder.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
- Stacked intermediate hidden states (output of each layer of the decoder).
- intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
- Stacked intermediate reference points (reference points of each layer of the decoder).
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, num_queries, hidden_size)`. Hidden-states of the decoder at the output of each layer
- plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, num_queries,
- num_queries)`. Attentions weights of the decoder, after the attention softmax, used to compute the weighted
- average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_queries, num_heads, 4, 4)`.
- Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
- weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_queries, num_heads, 4, 4)`.
- Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
- self-attention heads.
- enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
- Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
- picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
- foreground and background).
- enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
- Logits of predicted bounding boxes coordinates in the first stage.
- """
-
- init_reference_points: torch.FloatTensor = None
- last_hidden_state: torch.FloatTensor = None
- intermediate_hidden_states: torch.FloatTensor = None
- intermediate_reference_points: torch.FloatTensor = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- enc_outputs_class: Optional[torch.FloatTensor] = None
- enc_outputs_coord_logits: Optional[torch.FloatTensor] = None
-
-
-@dataclass
-class DeformableDetrObjectDetectionOutput(ModelOutput):
- """
- Output type of [`DeformableDetrForObjectDetection`].
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
- Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
- bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
- scale-invariant IoU loss.
- loss_dict (`Dict`, *optional*):
- A dictionary containing the individual losses. Useful for logging.
- logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
- Classification logits (including no-object) for all queries.
- pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
- values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
- possible padding). You can use [`~DeformableDetrProcessor.post_process_object_detection`] to retrieve the
- unnormalized bounding boxes.
- auxiliary_outputs (`list[Dict]`, *optional*):
- Optional, only returned when auxilary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
- and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
- `pred_boxes`) for each decoder layer.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, num_queries, hidden_size)`. Hidden-states of the decoder at the output of each layer
- plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, num_queries,
- num_queries)`. Attentions weights of the decoder, after the attention softmax, used to compute the weighted
- average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_queries, num_heads, 4, 4)`.
- Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
- weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_heads, 4,
- 4)`. Attentions weights of the encoder, after the attention softmax, used to compute the weighted average
- in the self-attention heads.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
- Stacked intermediate hidden states (output of each layer of the decoder).
- intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
- Stacked intermediate reference points (reference points of each layer of the decoder).
- init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Initial reference points sent through the Transformer decoder.
- enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
- Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
- picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
- foreground and background).
- enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
- Logits of predicted bounding boxes coordinates in the first stage.
- """
-
- loss: Optional[torch.FloatTensor] = None
- loss_dict: Optional[Dict] = None
- logits: torch.FloatTensor = None
- pred_boxes: torch.FloatTensor = None
- auxiliary_outputs: Optional[List[Dict]] = None
- init_reference_points: Optional[torch.FloatTensor] = None
- last_hidden_state: Optional[torch.FloatTensor] = None
- intermediate_hidden_states: Optional[torch.FloatTensor] = None
- intermediate_reference_points: Optional[torch.FloatTensor] = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- enc_outputs_class: Optional = None
- enc_outputs_coord_logits: Optional = None
-
-
-def _get_clones(module, N):
- return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
-
-
-def inverse_sigmoid(x, eps=1e-5):
- x = x.clamp(min=0, max=1)
- x1 = x.clamp(min=eps)
- x2 = (1 - x).clamp(min=eps)
- return torch.log(x1 / x2)
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrFrozenBatchNorm2d with Detr->DeformableDetr
-class DeformableDetrFrozenBatchNorm2d(nn.Module):
- """
- BatchNorm2d where the batch statistics and the affine parameters are fixed.
-
- Copy-paste from torchvision.misc.ops with added eps before rqsrt, without which any other models than
- torchvision.models.resnet[18,34,50,101] produce nans.
- """
-
- def __init__(self, n):
- super().__init__()
- self.register_buffer("weight", torch.ones(n))
- self.register_buffer("bias", torch.zeros(n))
- self.register_buffer("running_mean", torch.zeros(n))
- self.register_buffer("running_var", torch.ones(n))
-
- def _load_from_state_dict(
- self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
- ):
- num_batches_tracked_key = prefix + "num_batches_tracked"
- if num_batches_tracked_key in state_dict:
- del state_dict[num_batches_tracked_key]
-
- super()._load_from_state_dict(
- state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
- )
-
- def forward(self, x):
- # move reshapes to the beginning
- # to make it user-friendly
- weight = self.weight.reshape(1, -1, 1, 1)
- bias = self.bias.reshape(1, -1, 1, 1)
- running_var = self.running_var.reshape(1, -1, 1, 1)
- running_mean = self.running_mean.reshape(1, -1, 1, 1)
- epsilon = 1e-5
- scale = weight * (running_var + epsilon).rsqrt()
- bias = bias - running_mean * scale
- return x * scale + bias
-
-
-# Copied from transformers.models.detr.modeling_detr.replace_batch_norm with Detr->DeformableDetr
-def replace_batch_norm(model):
- r"""
- Recursively replace all `torch.nn.BatchNorm2d` with `DeformableDetrFrozenBatchNorm2d`.
-
- Args:
- model (torch.nn.Module):
- input model
- """
- for name, module in model.named_children():
- if isinstance(module, nn.BatchNorm2d):
- new_module = DeformableDetrFrozenBatchNorm2d(module.num_features)
-
- if not module.weight.device == torch.device("meta"):
- new_module.weight.data.copy_(module.weight)
- new_module.bias.data.copy_(module.bias)
- new_module.running_mean.data.copy_(module.running_mean)
- new_module.running_var.data.copy_(module.running_var)
-
- model._modules[name] = new_module
-
- if len(list(module.children())) > 0:
- replace_batch_norm(module)
-
-
-class DeformableDetrConvEncoder(nn.Module):
- """
- Convolutional backbone, using either the AutoBackbone API or one from the timm library.
-
- nn.BatchNorm2d layers are replaced by DeformableDetrFrozenBatchNorm2d as defined above.
-
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.config = config
-
- if config.use_timm_backbone:
- requires_backends(self, ["timm"])
- kwargs = {}
- if config.dilation:
- kwargs["output_stride"] = 16
- backbone = create_model(
- config.backbone,
- pretrained=config.use_pretrained_backbone,
- features_only=True,
- out_indices=(2, 3, 4) if config.num_feature_levels > 1 else (4,),
- in_chans=config.num_channels,
- **kwargs,
- )
- else:
- backbone = load_backbone(config)
-
- # replace batch norm by frozen batch norm
- with torch.no_grad():
- replace_batch_norm(backbone)
- self.model = backbone
- self.intermediate_channel_sizes = (
- self.model.feature_info.channels() if config.use_timm_backbone else self.model.channels
- )
-
- backbone_model_type = config.backbone if config.use_timm_backbone else config.backbone_config.model_type
- if "resnet" in backbone_model_type:
- for name, parameter in self.model.named_parameters():
- if config.use_timm_backbone:
- if "layer2" not in name and "layer3" not in name and "layer4" not in name:
- parameter.requires_grad_(False)
- else:
- if "stage.1" not in name and "stage.2" not in name and "stage.3" not in name:
- parameter.requires_grad_(False)
-
- # Copied from transformers.models.detr.modeling_detr.DetrConvEncoder.forward with Detr->DeformableDetr
- def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
- # send pixel_values through the model to get list of feature maps
- features = self.model(pixel_values) if self.config.use_timm_backbone else self.model(pixel_values).feature_maps
-
- out = []
- for feature_map in features:
- # downsample pixel_mask to match shape of corresponding feature_map
- mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
- out.append((feature_map, mask))
- return out
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrConvModel with Detr->DeformableDetr
-class DeformableDetrConvModel(nn.Module):
- """
- This module adds 2D position embeddings to all intermediate feature maps of the convolutional encoder.
- """
-
- def __init__(self, conv_encoder, position_embedding):
- super().__init__()
- self.conv_encoder = conv_encoder
- self.position_embedding = position_embedding
-
- def forward(self, pixel_values, pixel_mask):
- # send pixel_values and pixel_mask through backbone to get list of (feature_map, pixel_mask) tuples
- out = self.conv_encoder(pixel_values, pixel_mask)
- pos = []
- for feature_map, mask in out:
- # position encoding
- pos.append(self.position_embedding(feature_map, mask).to(feature_map.dtype))
-
- return out, pos
-
-
-class DeformableDetrSinePositionEmbedding(nn.Module):
- """
- This is a more standard version of the position embedding, very similar to the one used by the Attention is all you
- need paper, generalized to work on images.
- """
-
- def __init__(self, embedding_dim=64, temperature=10000, normalize=False, scale=None):
- super().__init__()
- self.embedding_dim = embedding_dim
- self.temperature = temperature
- self.normalize = normalize
- if scale is not None and normalize is False:
- raise ValueError("normalize should be True if scale is passed")
- if scale is None:
- scale = 2 * math.pi
- self.scale = scale
-
- def forward(self, pixel_values, pixel_mask):
- if pixel_mask is None:
- raise ValueError("No pixel mask provided")
- y_embed = pixel_mask.cumsum(1, dtype=torch.float32)
- x_embed = pixel_mask.cumsum(2, dtype=torch.float32)
- if self.normalize:
- eps = 1e-6
- y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale
- x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale
-
- dim_t = torch.arange(self.embedding_dim, dtype=torch.int64, device=pixel_values.device).float()
- dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.embedding_dim)
-
- pos_x = x_embed[:, :, :, None] / dim_t
- pos_y = y_embed[:, :, :, None] / dim_t
- pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
- return pos
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrLearnedPositionEmbedding
-class DeformableDetrLearnedPositionEmbedding(nn.Module):
- """
- This module learns positional embeddings up to a fixed maximum size.
- """
-
- def __init__(self, embedding_dim=256):
- super().__init__()
- self.row_embeddings = nn.Embedding(50, embedding_dim)
- self.column_embeddings = nn.Embedding(50, embedding_dim)
-
- def forward(self, pixel_values, pixel_mask=None):
- height, width = pixel_values.shape[-2:]
- width_values = torch.arange(width, device=pixel_values.device)
- height_values = torch.arange(height, device=pixel_values.device)
- x_emb = self.column_embeddings(width_values)
- y_emb = self.row_embeddings(height_values)
- pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1)
- pos = pos.permute(2, 0, 1)
- pos = pos.unsqueeze(0)
- pos = pos.repeat(pixel_values.shape[0], 1, 1, 1)
- return pos
-
-
-# Copied from transformers.models.detr.modeling_detr.build_position_encoding with Detr->DeformableDetr
-def build_position_encoding(config):
- n_steps = config.d_model // 2
- if config.position_embedding_type == "sine":
- # TODO find a better way of exposing other arguments
- position_embedding = DeformableDetrSinePositionEmbedding(n_steps, normalize=True)
- elif config.position_embedding_type == "learned":
- position_embedding = DeformableDetrLearnedPositionEmbedding(n_steps)
- else:
- raise ValueError(f"Not supported {config.position_embedding_type}")
-
- return position_embedding
-
-
-def multi_scale_deformable_attention(
- value: Tensor, value_spatial_shapes: Tensor, sampling_locations: Tensor, attention_weights: Tensor
-) -> Tensor:
- batch_size, _, num_heads, hidden_dim = value.shape
- _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
- value_list = value.split([height.item() * width.item() for height, width in value_spatial_shapes], dim=1)
- sampling_grids = 2 * sampling_locations - 1
- sampling_value_list = []
- for level_id, (height, width) in enumerate(value_spatial_shapes):
- # batch_size, height*width, num_heads, hidden_dim
- # -> batch_size, height*width, num_heads*hidden_dim
- # -> batch_size, num_heads*hidden_dim, height*width
- # -> batch_size*num_heads, hidden_dim, height, width
- value_l_ = (
- value_list[level_id].flatten(2).transpose(1, 2).reshape(batch_size * num_heads, hidden_dim, height, width)
- )
- # batch_size, num_queries, num_heads, num_points, 2
- # -> batch_size, num_heads, num_queries, num_points, 2
- # -> batch_size*num_heads, num_queries, num_points, 2
- sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1)
- # batch_size*num_heads, hidden_dim, num_queries, num_points
- sampling_value_l_ = nn.functional.grid_sample(
- value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False
- )
- sampling_value_list.append(sampling_value_l_)
- # (batch_size, num_queries, num_heads, num_levels, num_points)
- # -> (batch_size, num_heads, num_queries, num_levels, num_points)
- # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points)
- attention_weights = attention_weights.transpose(1, 2).reshape(
- batch_size * num_heads, 1, num_queries, num_levels * num_points
- )
- output = (
- (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights)
- .sum(-1)
- .view(batch_size, num_heads * hidden_dim, num_queries)
- )
- return output.transpose(1, 2).contiguous()
-
-
-class DeformableDetrMultiscaleDeformableAttention(nn.Module):
- """
- Multiscale deformable attention as proposed in Deformable DETR.
- """
-
- def __init__(self, config: DeformableDetrConfig, num_heads: int, n_points: int):
- super().__init__()
-
- kernel_loaded = MultiScaleDeformableAttention is not None
- if is_torch_cuda_available() and is_ninja_available() and not kernel_loaded:
- try:
- load_cuda_kernels()
- except Exception as e:
- logger.warning(f"Could not load the custom kernel for multi-scale deformable attention: {e}")
-
- if config.d_model % num_heads != 0:
- raise ValueError(
- f"embed_dim (d_model) must be divisible by num_heads, but got {config.d_model} and {num_heads}"
- )
- dim_per_head = config.d_model // num_heads
- # check if dim_per_head is power of 2
- if not ((dim_per_head & (dim_per_head - 1) == 0) and dim_per_head != 0):
- warnings.warn(
- "You'd better set embed_dim (d_model) in DeformableDetrMultiscaleDeformableAttention to make the"
- " dimension of each attention head a power of 2 which is more efficient in the authors' CUDA"
- " implementation."
- )
-
- self.im2col_step = 64
-
- self.d_model = config.d_model
- self.n_levels = config.num_feature_levels
- self.n_heads = num_heads
- self.n_points = n_points
-
- self.sampling_offsets = nn.Linear(config.d_model, num_heads * self.n_levels * n_points * 2)
- self.attention_weights = nn.Linear(config.d_model, num_heads * self.n_levels * n_points)
- self.value_proj = nn.Linear(config.d_model, config.d_model)
- self.output_proj = nn.Linear(config.d_model, config.d_model)
-
- self.disable_custom_kernels = config.disable_custom_kernels
-
- self._reset_parameters()
-
- def _reset_parameters(self):
- nn.init.constant_(self.sampling_offsets.weight.data, 0.0)
- default_dtype = torch.get_default_dtype()
- thetas = torch.arange(self.n_heads, dtype=torch.int64).to(default_dtype) * (2.0 * math.pi / self.n_heads)
- grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
- grid_init = (
- (grid_init / grid_init.abs().max(-1, keepdim=True)[0])
- .view(self.n_heads, 1, 1, 2)
- .repeat(1, self.n_levels, self.n_points, 1)
- )
- for i in range(self.n_points):
- grid_init[:, :, i, :] *= i + 1
- with torch.no_grad():
- self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
- nn.init.constant_(self.attention_weights.weight.data, 0.0)
- nn.init.constant_(self.attention_weights.bias.data, 0.0)
- nn.init.xavier_uniform_(self.value_proj.weight.data)
- nn.init.constant_(self.value_proj.bias.data, 0.0)
- nn.init.xavier_uniform_(self.output_proj.weight.data)
- nn.init.constant_(self.output_proj.bias.data, 0.0)
-
- def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Optional[Tensor]):
- return tensor if position_embeddings is None else tensor + position_embeddings
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- position_embeddings: Optional[torch.Tensor] = None,
- reference_points=None,
- spatial_shapes=None,
- level_start_index=None,
- output_attentions: bool = False,
- ):
- # add position embeddings to the hidden states before projecting to queries and keys
- if position_embeddings is not None:
- hidden_states = self.with_pos_embed(hidden_states, position_embeddings)
-
- batch_size, num_queries, _ = hidden_states.shape
- batch_size, sequence_length, _ = encoder_hidden_states.shape
- if (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() != sequence_length:
- raise ValueError(
- "Make sure to align the spatial shapes with the sequence length of the encoder hidden states"
- )
-
- value = self.value_proj(encoder_hidden_states)
- if attention_mask is not None:
- # we invert the attention_mask
- value = value.masked_fill(~attention_mask[..., None], float(0))
- value = value.view(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads)
- sampling_offsets = self.sampling_offsets(hidden_states).view(
- batch_size, num_queries, self.n_heads, self.n_levels, self.n_points, 2
- )
- attention_weights = self.attention_weights(hidden_states).view(
- batch_size, num_queries, self.n_heads, self.n_levels * self.n_points
- )
- attention_weights = F.softmax(attention_weights, -1).view(
- batch_size, num_queries, self.n_heads, self.n_levels, self.n_points
- )
- # batch_size, num_queries, n_heads, n_levels, n_points, 2
- num_coordinates = reference_points.shape[-1]
- if num_coordinates == 2:
- offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1)
- sampling_locations = (
- reference_points[:, :, None, :, None, :]
- + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
- )
- elif num_coordinates == 4:
- sampling_locations = (
- reference_points[:, :, None, :, None, :2]
- + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
- )
- else:
- raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}")
-
- if self.disable_custom_kernels:
- # PyTorch implementation
- output = multi_scale_deformable_attention(value, spatial_shapes, sampling_locations, attention_weights)
- else:
- try:
- # custom kernel
- output = MultiScaleDeformableAttentionFunction.apply(
- value,
- spatial_shapes,
- level_start_index,
- sampling_locations,
- attention_weights,
- self.im2col_step,
- )
- except Exception:
- # PyTorch implementation
- output = multi_scale_deformable_attention(value, spatial_shapes, sampling_locations, attention_weights)
- output = self.output_proj(output)
-
- return output, attention_weights
-
-
-class DeformableDetrMultiheadAttention(nn.Module):
- """
- Multi-headed attention from 'Attention Is All You Need' paper.
-
- Here, we add position embeddings to the queries and keys (as explained in the Deformable DETR paper).
- """
-
- def __init__(
- self,
- embed_dim: int,
- num_heads: int,
- dropout: float = 0.0,
- bias: bool = True,
- ):
- super().__init__()
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.dropout = dropout
- self.head_dim = embed_dim // num_heads
- if self.head_dim * num_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {num_heads})."
- )
- self.scaling = self.head_dim**-0.5
-
- self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
-
- def _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int):
- return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Optional[Tensor]):
- return tensor if position_embeddings is None else tensor + position_embeddings
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_embeddings: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel"""
-
- batch_size, target_len, embed_dim = hidden_states.size()
- # add position embeddings to the hidden states before projecting to queries and keys
- if position_embeddings is not None:
- hidden_states_original = hidden_states
- hidden_states = self.with_pos_embed(hidden_states, position_embeddings)
-
- # get queries, keys and values
- query_states = self.q_proj(hidden_states) * self.scaling
- key_states = self._shape(self.k_proj(hidden_states), -1, batch_size)
- value_states = self._shape(self.v_proj(hidden_states_original), -1, batch_size)
-
- proj_shape = (batch_size * self.num_heads, -1, self.head_dim)
- query_states = self._shape(query_states, target_len, batch_size).view(*proj_shape)
- key_states = key_states.view(*proj_shape)
- value_states = value_states.view(*proj_shape)
-
- source_len = key_states.size(1)
-
- attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-
- if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len):
- raise ValueError(
- f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- # expand attention_mask
- if attention_mask is not None:
- # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len]
- attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
-
- if attention_mask is not None:
- if attention_mask.size() != (batch_size, 1, target_len, source_len):
- raise ValueError(
- f"Attention mask should be of size {(batch_size, 1, target_len, source_len)}, but is"
- f" {attention_mask.size()}"
- )
- attn_weights = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attention_mask
- attn_weights = attn_weights.view(batch_size * self.num_heads, target_len, source_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if output_attentions:
- # this operation is a bit awkward, but it's required to
- # make sure that attn_weights keeps its gradient.
- # In order to do so, attn_weights have to reshaped
- # twice and have to be reused in the following
- attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len)
- attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-
- attn_output = torch.bmm(attn_probs, value_states)
-
- if attn_output.size() != (batch_size * self.num_heads, target_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.head_dim)
- attn_output = attn_output.transpose(1, 2)
- attn_output = attn_output.reshape(batch_size, target_len, embed_dim)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped
-
-
-class DeformableDetrEncoderLayer(nn.Module):
- def __init__(self, config: DeformableDetrConfig):
- super().__init__()
- self.embed_dim = config.d_model
- self.self_attn = DeformableDetrMultiscaleDeformableAttention(
- config, num_heads=config.encoder_attention_heads, n_points=config.encoder_n_points
- )
- self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
- self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
- self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: torch.Tensor,
- position_embeddings: torch.Tensor = None,
- reference_points=None,
- spatial_shapes=None,
- level_start_index=None,
- output_attentions: bool = False,
- ):
- """
- Args:
- hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Input to the layer.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
- Attention mask.
- position_embeddings (`torch.FloatTensor`, *optional*):
- Position embeddings, to be added to `hidden_states`.
- reference_points (`torch.FloatTensor`, *optional*):
- Reference points.
- spatial_shapes (`torch.LongTensor`, *optional*):
- Spatial shapes of the backbone feature maps.
- level_start_index (`torch.LongTensor`, *optional*):
- Level start index.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- residual = hidden_states
-
- # Apply Multi-scale Deformable Attention Module on the multi-scale feature maps.
- hidden_states, attn_weights = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- encoder_hidden_states=hidden_states,
- encoder_attention_mask=attention_mask,
- position_embeddings=position_embeddings,
- reference_points=reference_points,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- residual = hidden_states
- hidden_states = self.activation_fn(self.fc1(hidden_states))
- hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
-
- hidden_states = self.fc2(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
- hidden_states = residual + hidden_states
- hidden_states = self.final_layer_norm(hidden_states)
-
- if self.training:
- if torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any():
- clamp_value = torch.finfo(hidden_states.dtype).max - 1000
- hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
-
-class DeformableDetrDecoderLayer(nn.Module):
- def __init__(self, config: DeformableDetrConfig):
- super().__init__()
- self.embed_dim = config.d_model
-
- # self-attention
- self.self_attn = DeformableDetrMultiheadAttention(
- embed_dim=self.embed_dim,
- num_heads=config.decoder_attention_heads,
- dropout=config.attention_dropout,
- )
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
-
- self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- # cross-attention
- self.encoder_attn = DeformableDetrMultiscaleDeformableAttention(
- config,
- num_heads=config.decoder_attention_heads,
- n_points=config.decoder_n_points,
- )
- self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- # feedforward neural networks
- self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
- self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- position_embeddings: Optional[torch.Tensor] = None,
- reference_points=None,
- spatial_shapes=None,
- level_start_index=None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- ):
- """
- Args:
- hidden_states (`torch.FloatTensor`):
- Input to the layer of shape `(seq_len, batch, embed_dim)`.
- position_embeddings (`torch.FloatTensor`, *optional*):
- Position embeddings that are added to the queries and keys in the self-attention layer.
- reference_points (`torch.FloatTensor`, *optional*):
- Reference points.
- spatial_shapes (`torch.LongTensor`, *optional*):
- Spatial shapes.
- level_start_index (`torch.LongTensor`, *optional*):
- Level start index.
- encoder_hidden_states (`torch.FloatTensor`):
- cross attention input to the layer of shape `(seq_len, batch, embed_dim)`
- encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
- `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
- values.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- residual = hidden_states
-
- # Self Attention
- hidden_states, self_attn_weights = self.self_attn(
- hidden_states=hidden_states,
- position_embeddings=position_embeddings,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- second_residual = hidden_states
-
- # Cross-Attention
- cross_attn_weights = None
- hidden_states, cross_attn_weights = self.encoder_attn(
- hidden_states=hidden_states,
- attention_mask=encoder_attention_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- position_embeddings=position_embeddings,
- reference_points=reference_points,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = second_residual + hidden_states
-
- hidden_states = self.encoder_attn_layer_norm(hidden_states)
-
- # Fully Connected
- residual = hidden_states
- hidden_states = self.activation_fn(self.fc1(hidden_states))
- hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
- hidden_states = self.fc2(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.final_layer_norm(hidden_states)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights, cross_attn_weights)
-
- return outputs
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrClassificationHead
-class DeformableDetrClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, input_dim: int, inner_dim: int, num_classes: int, pooler_dropout: float):
- super().__init__()
- self.dense = nn.Linear(input_dim, inner_dim)
- self.dropout = nn.Dropout(p=pooler_dropout)
- self.out_proj = nn.Linear(inner_dim, num_classes)
-
- def forward(self, hidden_states: torch.Tensor):
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.dense(hidden_states)
- hidden_states = torch.tanh(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.out_proj(hidden_states)
- return hidden_states
-
-
-class DeformableDetrPreTrainedModel(PreTrainedModel):
- config_class = DeformableDetrConfig
- base_model_prefix = "model"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
- _no_split_modules = [r"DeformableDetrConvEncoder", r"DeformableDetrEncoderLayer", r"DeformableDetrDecoderLayer"]
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- std = self.config.init_std
-
- if isinstance(module, DeformableDetrLearnedPositionEmbedding):
- nn.init.uniform_(module.row_embeddings.weight)
- nn.init.uniform_(module.column_embeddings.weight)
- elif isinstance(module, DeformableDetrMultiscaleDeformableAttention):
- module._reset_parameters()
- elif isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- if hasattr(module, "reference_points") and not self.config.two_stage:
- nn.init.xavier_uniform_(module.reference_points.weight.data, gain=1.0)
- nn.init.constant_(module.reference_points.bias.data, 0.0)
- if hasattr(module, "level_embed"):
- nn.init.normal_(module.level_embed)
-
-
-DEFORMABLE_DETR_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`DeformableDetrConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DEFORMABLE_DETR_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Padding will be ignored by default should you provide it.
-
- Pixel values can be obtained using [`AutoImageProcessor`]. See [`DeformableDetrImageProcessor.__call__`]
- for details.
-
- pixel_mask (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
- Mask to avoid performing attention on padding pixel values. Mask values selected in `[0, 1]`:
-
- - 1 for pixels that are real (i.e. **not masked**),
- - 0 for pixels that are padding (i.e. **masked**).
-
- [What are attention masks?](../glossary#attention-mask)
-
- decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*):
- Not used by default. Can be used to mask object queries.
- encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*):
- Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
- `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of
- hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you
- can choose to directly pass a flattened representation of an image.
- decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
- Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an
- embedded representation.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class DeformableDetrEncoder(DeformableDetrPreTrainedModel):
- """
- Transformer encoder consisting of *config.encoder_layers* deformable attention layers. Each layer is a
- [`DeformableDetrEncoderLayer`].
-
- The encoder updates the flattened multi-scale feature maps through multiple deformable attention layers.
-
- Args:
- config: DeformableDetrConfig
- """
-
- def __init__(self, config: DeformableDetrConfig):
- super().__init__(config)
- self.gradient_checkpointing = False
-
- self.dropout = config.dropout
- self.layers = nn.ModuleList([DeformableDetrEncoderLayer(config) for _ in range(config.encoder_layers)])
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @staticmethod
- def get_reference_points(spatial_shapes, valid_ratios, device):
- """
- Get reference points for each feature map. Used in decoder.
-
- Args:
- spatial_shapes (`torch.LongTensor` of shape `(num_feature_levels, 2)`):
- Spatial shapes of each feature map.
- valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`):
- Valid ratios of each feature map.
- device (`torch.device`):
- Device on which to create the tensors.
- Returns:
- `torch.FloatTensor` of shape `(batch_size, num_queries, num_feature_levels, 2)`
- """
- reference_points_list = []
- for level, (height, width) in enumerate(spatial_shapes):
- ref_y, ref_x = meshgrid(
- torch.linspace(0.5, height - 0.5, height, dtype=valid_ratios.dtype, device=device),
- torch.linspace(0.5, width - 0.5, width, dtype=valid_ratios.dtype, device=device),
- indexing="ij",
- )
- # TODO: valid_ratios could be useless here. check https://github.com/fundamentalvision/Deformable-DETR/issues/36
- ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, level, 1] * height)
- ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, level, 0] * width)
- ref = torch.stack((ref_x, ref_y), -1)
- reference_points_list.append(ref)
- reference_points = torch.cat(reference_points_list, 1)
- reference_points = reference_points[:, :, None] * valid_ratios[:, None]
- return reference_points
-
- def forward(
- self,
- inputs_embeds=None,
- attention_mask=None,
- position_embeddings=None,
- spatial_shapes=None,
- level_start_index=None,
- valid_ratios=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- ):
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Flattened feature map (output of the backbone + projection layer) that is passed to the encoder.
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`:
- - 1 for pixel features that are real (i.e. **not masked**),
- - 0 for pixel features that are padding (i.e. **masked**).
- [What are attention masks?](../glossary#attention-mask)
- position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Position embeddings that are added to the queries and keys in each self-attention layer.
- spatial_shapes (`torch.LongTensor` of shape `(num_feature_levels, 2)`):
- Spatial shapes of each feature map.
- level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`):
- Starting index of each feature map.
- valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`):
- Ratio of valid area in each feature level.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- hidden_states = inputs_embeds
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
- reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=inputs_embeds.device)
-
- encoder_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
- for i, encoder_layer in enumerate(self.layers):
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- encoder_layer.__call__,
- hidden_states,
- attention_mask,
- position_embeddings,
- reference_points,
- spatial_shapes,
- level_start_index,
- output_attentions,
- )
- else:
- layer_outputs = encoder_layer(
- hidden_states,
- attention_mask,
- position_embeddings=position_embeddings,
- reference_points=reference_points,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
- )
-
-
-class DeformableDetrDecoder(DeformableDetrPreTrainedModel):
- """
- Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DeformableDetrDecoderLayer`].
-
- The decoder updates the query embeddings through multiple self-attention and cross-attention layers.
-
- Some tweaks for Deformable DETR:
-
- - `position_embeddings`, `reference_points`, `spatial_shapes` and `valid_ratios` are added to the forward pass.
- - it also returns a stack of intermediate outputs and reference points from all decoding layers.
-
- Args:
- config: DeformableDetrConfig
- """
-
- def __init__(self, config: DeformableDetrConfig):
- super().__init__(config)
-
- self.dropout = config.dropout
- self.layers = nn.ModuleList([DeformableDetrDecoderLayer(config) for _ in range(config.decoder_layers)])
- self.gradient_checkpointing = False
-
- # hack implementation for iterative bounding box refinement and two-stage Deformable DETR
- self.bbox_embed = None
- self.class_embed = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def forward(
- self,
- inputs_embeds=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- position_embeddings=None,
- reference_points=None,
- spatial_shapes=None,
- level_start_index=None,
- valid_ratios=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- ):
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
- The query embeddings that are passed into the decoder.
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
- of the decoder.
- encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing cross-attention on padding pixel_values of the encoder. Mask values selected
- in `[0, 1]`:
- - 1 for pixels that are real (i.e. **not masked**),
- - 0 for pixels that are padding (i.e. **masked**).
- position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
- Position embeddings that are added to the queries and keys in each self-attention layer.
- reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)` is `as_two_stage` else `(batch_size, num_queries, 2)` or , *optional*):
- Reference point in range `[0, 1]`, top-left (0,0), bottom-right (1, 1), including padding area.
- spatial_shapes (`torch.FloatTensor` of shape `(num_feature_levels, 2)`):
- Spatial shapes of the feature maps.
- level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`, *optional*):
- Indexes for the start of each feature level. In range `[0, sequence_length]`.
- valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`, *optional*):
- Ratio of valid area in each feature level.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if inputs_embeds is not None:
- hidden_states = inputs_embeds
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
- intermediate = ()
- intermediate_reference_points = ()
-
- for idx, decoder_layer in enumerate(self.layers):
- num_coordinates = reference_points.shape[-1]
- if num_coordinates == 4:
- reference_points_input = (
- reference_points[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[:, None]
- )
- elif reference_points.shape[-1] == 2:
- reference_points_input = reference_points[:, :, None] * valid_ratios[:, None]
- else:
- raise ValueError("Reference points' last dimension must be of size 2")
-
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- decoder_layer.__call__,
- hidden_states,
- position_embeddings,
- reference_points_input,
- spatial_shapes,
- level_start_index,
- encoder_hidden_states,
- encoder_attention_mask,
- output_attentions,
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- position_embeddings=position_embeddings,
- encoder_hidden_states=encoder_hidden_states,
- reference_points=reference_points_input,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- # hack implementation for iterative bounding box refinement
- if self.bbox_embed is not None:
- tmp = self.bbox_embed[idx](hidden_states)
- num_coordinates = reference_points.shape[-1]
- if num_coordinates == 4:
- new_reference_points = tmp + inverse_sigmoid(reference_points)
- new_reference_points = new_reference_points.sigmoid()
- elif num_coordinates == 2:
- new_reference_points = tmp
- new_reference_points[..., :2] = tmp[..., :2] + inverse_sigmoid(reference_points)
- new_reference_points = new_reference_points.sigmoid()
- else:
- raise ValueError(
- f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}"
- )
- reference_points = new_reference_points.detach()
-
- intermediate += (hidden_states,)
- intermediate_reference_points += (reference_points,)
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- if encoder_hidden_states is not None:
- all_cross_attentions += (layer_outputs[2],)
-
- # Keep batch_size as first dimension
- intermediate = torch.stack(intermediate, dim=1)
- intermediate_reference_points = torch.stack(intermediate_reference_points, dim=1)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- intermediate,
- intermediate_reference_points,
- all_hidden_states,
- all_self_attns,
- all_cross_attentions,
- ]
- if v is not None
- )
- return DeformableDetrDecoderOutput(
- last_hidden_state=hidden_states,
- intermediate_hidden_states=intermediate,
- intermediate_reference_points=intermediate_reference_points,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- cross_attentions=all_cross_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The bare Deformable DETR Model (consisting of a backbone and encoder-decoder Transformer) outputting raw
- hidden-states without any specific head on top.
- """,
- DEFORMABLE_DETR_START_DOCSTRING,
-)
-class DeformableDetrModel(DeformableDetrPreTrainedModel):
- def __init__(self, config: DeformableDetrConfig):
- super().__init__(config)
-
- # Create backbone + positional encoding
- backbone = DeformableDetrConvEncoder(config)
- position_embeddings = build_position_encoding(config)
- self.backbone = DeformableDetrConvModel(backbone, position_embeddings)
-
- # Create input projection layers
- if config.num_feature_levels > 1:
- num_backbone_outs = len(backbone.intermediate_channel_sizes)
- input_proj_list = []
- for _ in range(num_backbone_outs):
- in_channels = backbone.intermediate_channel_sizes[_]
- input_proj_list.append(
- nn.Sequential(
- nn.Conv2d(in_channels, config.d_model, kernel_size=1),
- nn.GroupNorm(32, config.d_model),
- )
- )
- for _ in range(config.num_feature_levels - num_backbone_outs):
- input_proj_list.append(
- nn.Sequential(
- nn.Conv2d(in_channels, config.d_model, kernel_size=3, stride=2, padding=1),
- nn.GroupNorm(32, config.d_model),
- )
- )
- in_channels = config.d_model
- self.input_proj = nn.ModuleList(input_proj_list)
- else:
- self.input_proj = nn.ModuleList(
- [
- nn.Sequential(
- nn.Conv2d(backbone.intermediate_channel_sizes[-1], config.d_model, kernel_size=1),
- nn.GroupNorm(32, config.d_model),
- )
- ]
- )
-
- if not config.two_stage:
- self.query_position_embeddings = nn.Embedding(config.num_queries, config.d_model * 2)
-
- self.encoder = DeformableDetrEncoder(config)
- self.decoder = DeformableDetrDecoder(config)
-
- self.level_embed = nn.Parameter(torch.Tensor(config.num_feature_levels, config.d_model))
-
- if config.two_stage:
- self.enc_output = nn.Linear(config.d_model, config.d_model)
- self.enc_output_norm = nn.LayerNorm(config.d_model)
- self.pos_trans = nn.Linear(config.d_model * 2, config.d_model * 2)
- self.pos_trans_norm = nn.LayerNorm(config.d_model * 2)
- else:
- self.reference_points = nn.Linear(config.d_model, 2)
-
- self.post_init()
-
- def get_encoder(self):
- return self.encoder
-
- def get_decoder(self):
- return self.decoder
-
- def freeze_backbone(self):
- for name, param in self.backbone.conv_encoder.model.named_parameters():
- param.requires_grad_(False)
-
- def unfreeze_backbone(self):
- for name, param in self.backbone.conv_encoder.model.named_parameters():
- param.requires_grad_(True)
-
- def get_valid_ratio(self, mask, dtype=torch.float32):
- """Get the valid ratio of all feature maps."""
-
- _, height, width = mask.shape
- valid_height = torch.sum(mask[:, :, 0], 1)
- valid_width = torch.sum(mask[:, 0, :], 1)
- valid_ratio_height = valid_height.to(dtype) / height
- valid_ratio_width = valid_width.to(dtype) / width
- valid_ratio = torch.stack([valid_ratio_width, valid_ratio_height], -1)
- return valid_ratio
-
- def get_proposal_pos_embed(self, proposals):
- """Get the position embedding of the proposals."""
-
- num_pos_feats = self.config.d_model // 2
- temperature = 10000
- scale = 2 * math.pi
-
- dim_t = torch.arange(num_pos_feats, dtype=torch.int64, device=proposals.device).float()
- dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats)
- # batch_size, num_queries, 4
- proposals = proposals.sigmoid() * scale
- # batch_size, num_queries, 4, 128
- pos = proposals[:, :, :, None] / dim_t
- # batch_size, num_queries, 4, 64, 2 -> batch_size, num_queries, 512
- pos = torch.stack((pos[:, :, :, 0::2].sin(), pos[:, :, :, 1::2].cos()), dim=4).flatten(2)
- return pos
-
- def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes):
- """Generate the encoder output proposals from encoded enc_output.
-
- Args:
- enc_output (Tensor[batch_size, sequence_length, hidden_size]): Output of the encoder.
- padding_mask (Tensor[batch_size, sequence_length]): Padding mask for `enc_output`.
- spatial_shapes (Tensor[num_feature_levels, 2]): Spatial shapes of the feature maps.
-
- Returns:
- `tuple(torch.FloatTensor)`: A tuple of feature map and bbox prediction.
- - object_query (Tensor[batch_size, sequence_length, hidden_size]): Object query features. Later used to
- directly predict a bounding box. (without the need of a decoder)
- - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals, after an inverse
- sigmoid.
- """
- batch_size = enc_output.shape[0]
- proposals = []
- _cur = 0
- for level, (height, width) in enumerate(spatial_shapes):
- mask_flatten_ = padding_mask[:, _cur : (_cur + height * width)].view(batch_size, height, width, 1)
- valid_height = torch.sum(~mask_flatten_[:, :, 0, 0], 1)
- valid_width = torch.sum(~mask_flatten_[:, 0, :, 0], 1)
-
- grid_y, grid_x = meshgrid(
- torch.linspace(0, height - 1, height, dtype=torch.float32, device=enc_output.device),
- torch.linspace(0, width - 1, width, dtype=torch.float32, device=enc_output.device),
- indexing="ij",
- )
- grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1)
-
- scale = torch.cat([valid_width.unsqueeze(-1), valid_height.unsqueeze(-1)], 1).view(batch_size, 1, 1, 2)
- grid = (grid.unsqueeze(0).expand(batch_size, -1, -1, -1) + 0.5) / scale
- width_heigth = torch.ones_like(grid) * 0.05 * (2.0**level)
- proposal = torch.cat((grid, width_heigth), -1).view(batch_size, -1, 4)
- proposals.append(proposal)
- _cur += height * width
- output_proposals = torch.cat(proposals, 1)
- output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True)
- output_proposals = torch.log(output_proposals / (1 - output_proposals)) # inverse sigmoid
- output_proposals = output_proposals.masked_fill(padding_mask.unsqueeze(-1), float("inf"))
- output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf"))
-
- # assign each pixel as an object query
- object_query = enc_output
- object_query = object_query.masked_fill(padding_mask.unsqueeze(-1), float(0))
- object_query = object_query.masked_fill(~output_proposals_valid, float(0))
- object_query = self.enc_output_norm(self.enc_output(object_query))
- return object_query, output_proposals
-
- @add_start_docstrings_to_model_forward(DEFORMABLE_DETR_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DeformableDetrModelOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.FloatTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], DeformableDetrModelOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, DeformableDetrModel
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("SenseTime/deformable-detr")
- >>> model = DeformableDetrModel.from_pretrained("SenseTime/deformable-detr")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> outputs = model(**inputs)
-
- >>> last_hidden_states = outputs.last_hidden_state
- >>> list(last_hidden_states.shape)
- [1, 300, 256]
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- batch_size, num_channels, height, width = pixel_values.shape
- device = pixel_values.device
-
- if pixel_mask is None:
- pixel_mask = torch.ones(((batch_size, height, width)), dtype=torch.long, device=device)
-
- # Extract multi-scale feature maps of same resolution `config.d_model` (cf Figure 4 in paper)
- # First, sent pixel_values + pixel_mask through Backbone to obtain the features
- # which is a list of tuples
- features, position_embeddings_list = self.backbone(pixel_values, pixel_mask)
-
- # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
- sources = []
- masks = []
- for level, (source, mask) in enumerate(features):
- sources.append(self.input_proj[level](source))
- masks.append(mask)
- if mask is None:
- raise ValueError("No attention mask was provided")
-
- # Lowest resolution feature maps are obtained via 3x3 stride 2 convolutions on the final stage
- if self.config.num_feature_levels > len(sources):
- _len_sources = len(sources)
- for level in range(_len_sources, self.config.num_feature_levels):
- if level == _len_sources:
- source = self.input_proj[level](features[-1][0])
- else:
- source = self.input_proj[level](sources[-1])
- mask = nn.functional.interpolate(pixel_mask[None].float(), size=source.shape[-2:]).to(torch.bool)[0]
- pos_l = self.backbone.position_embedding(source, mask).to(source.dtype)
- sources.append(source)
- masks.append(mask)
- position_embeddings_list.append(pos_l)
-
- # Create queries
- query_embeds = None
- if not self.config.two_stage:
- query_embeds = self.query_position_embeddings.weight
-
- # Prepare encoder inputs (by flattening)
- source_flatten = []
- mask_flatten = []
- lvl_pos_embed_flatten = []
- spatial_shapes = []
- for level, (source, mask, pos_embed) in enumerate(zip(sources, masks, position_embeddings_list)):
- batch_size, num_channels, height, width = source.shape
- spatial_shape = (height, width)
- spatial_shapes.append(spatial_shape)
- source = source.flatten(2).transpose(1, 2)
- mask = mask.flatten(1)
- pos_embed = pos_embed.flatten(2).transpose(1, 2)
- lvl_pos_embed = pos_embed + self.level_embed[level].view(1, 1, -1)
- lvl_pos_embed_flatten.append(lvl_pos_embed)
- source_flatten.append(source)
- mask_flatten.append(mask)
- source_flatten = torch.cat(source_flatten, 1)
- mask_flatten = torch.cat(mask_flatten, 1)
- lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
- spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=source_flatten.device)
- level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
- valid_ratios = torch.stack([self.get_valid_ratio(m, dtype=source_flatten.dtype) for m in masks], 1)
-
- # Fourth, sent source_flatten + mask_flatten + lvl_pos_embed_flatten (backbone + proj layer output) through encoder
- # Also provide spatial_shapes, level_start_index and valid_ratios
- if encoder_outputs is None:
- encoder_outputs = self.encoder(
- inputs_embeds=source_flatten,
- attention_mask=mask_flatten,
- position_embeddings=lvl_pos_embed_flatten,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- valid_ratios=valid_ratios,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
- elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
- encoder_outputs = BaseModelOutput(
- last_hidden_state=encoder_outputs[0],
- hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
- attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
- )
-
- # Fifth, prepare decoder inputs
- batch_size, _, num_channels = encoder_outputs[0].shape
- enc_outputs_class = None
- enc_outputs_coord_logits = None
- if self.config.two_stage:
- object_query_embedding, output_proposals = self.gen_encoder_output_proposals(
- encoder_outputs[0], ~mask_flatten, spatial_shapes
- )
-
- # hack implementation for two-stage Deformable DETR
- # apply a detection head to each pixel (A.4 in paper)
- # linear projection for bounding box binary classification (i.e. foreground and background)
- enc_outputs_class = self.decoder.class_embed[-1](object_query_embedding)
- # 3-layer FFN to predict bounding boxes coordinates (bbox regression branch)
- delta_bbox = self.decoder.bbox_embed[-1](object_query_embedding)
- enc_outputs_coord_logits = delta_bbox + output_proposals
-
- # only keep top scoring `config.two_stage_num_proposals` proposals
- topk = self.config.two_stage_num_proposals
- topk_proposals = torch.topk(enc_outputs_class[..., 0], topk, dim=1)[1]
- topk_coords_logits = torch.gather(
- enc_outputs_coord_logits, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4)
- )
-
- topk_coords_logits = topk_coords_logits.detach()
- reference_points = topk_coords_logits.sigmoid()
- init_reference_points = reference_points
- pos_trans_out = self.pos_trans_norm(self.pos_trans(self.get_proposal_pos_embed(topk_coords_logits)))
- query_embed, target = torch.split(pos_trans_out, num_channels, dim=2)
- else:
- query_embed, target = torch.split(query_embeds, num_channels, dim=1)
- query_embed = query_embed.unsqueeze(0).expand(batch_size, -1, -1)
- target = target.unsqueeze(0).expand(batch_size, -1, -1)
- reference_points = self.reference_points(query_embed).sigmoid()
- init_reference_points = reference_points
-
- decoder_outputs = self.decoder(
- inputs_embeds=target,
- position_embeddings=query_embed,
- encoder_hidden_states=encoder_outputs[0],
- encoder_attention_mask=mask_flatten,
- reference_points=reference_points,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- valid_ratios=valid_ratios,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- enc_outputs = tuple(value for value in [enc_outputs_class, enc_outputs_coord_logits] if value is not None)
- tuple_outputs = (init_reference_points,) + decoder_outputs + encoder_outputs + enc_outputs
-
- return tuple_outputs
-
- return DeformableDetrModelOutput(
- init_reference_points=init_reference_points,
- last_hidden_state=decoder_outputs.last_hidden_state,
- intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
- intermediate_reference_points=decoder_outputs.intermediate_reference_points,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- enc_outputs_class=enc_outputs_class,
- enc_outputs_coord_logits=enc_outputs_coord_logits,
- )
-
-
-@add_start_docstrings(
- """
- Deformable DETR Model (consisting of a backbone and encoder-decoder Transformer) with object detection heads on
- top, for tasks such as COCO detection.
- """,
- DEFORMABLE_DETR_START_DOCSTRING,
-)
-class DeformableDetrForObjectDetection(DeformableDetrPreTrainedModel):
- # When using clones, all layers > 0 will be clones, but layer 0 *is* required
- _tied_weights_keys = [r"bbox_embed\.[1-9]\d*", r"class_embed\.[1-9]\d*"]
- # We can't initialize the model on meta device as some weights are modified during the initialization
- _no_split_modules = None
-
- def __init__(self, config: DeformableDetrConfig):
- super().__init__(config)
-
- # Deformable DETR encoder-decoder model
- self.model = DeformableDetrModel(config)
-
- # Detection heads on top
- self.class_embed = nn.Linear(config.d_model, config.num_labels)
- self.bbox_embed = DeformableDetrMLPPredictionHead(
- input_dim=config.d_model, hidden_dim=config.d_model, output_dim=4, num_layers=3
- )
-
- prior_prob = 0.01
- bias_value = -math.log((1 - prior_prob) / prior_prob)
- self.class_embed.bias.data = torch.ones(config.num_labels) * bias_value
- nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
- nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)
-
- # if two-stage, the last class_embed and bbox_embed is for region proposal generation
- num_pred = (config.decoder_layers + 1) if config.two_stage else config.decoder_layers
- if config.with_box_refine:
- self.class_embed = _get_clones(self.class_embed, num_pred)
- self.bbox_embed = _get_clones(self.bbox_embed, num_pred)
- nn.init.constant_(self.bbox_embed[0].layers[-1].bias.data[2:], -2.0)
- # hack implementation for iterative bounding box refinement
- self.model.decoder.bbox_embed = self.bbox_embed
- else:
- nn.init.constant_(self.bbox_embed.layers[-1].bias.data[2:], -2.0)
- self.class_embed = nn.ModuleList([self.class_embed for _ in range(num_pred)])
- self.bbox_embed = nn.ModuleList([self.bbox_embed for _ in range(num_pred)])
- self.model.decoder.bbox_embed = None
- if config.two_stage:
- # hack implementation for two-stage
- self.model.decoder.class_embed = self.class_embed
- for box_embed in self.bbox_embed:
- nn.init.constant_(box_embed.layers[-1].bias.data[2:], 0.0)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- # taken from https://github.com/facebookresearch/detr/blob/master/models/detr.py
- @torch.jit.unused
- def _set_aux_loss(self, outputs_class, outputs_coord):
- # this is a workaround to make torchscript happy, as torchscript
- # doesn't support dictionary with non-homogeneous values, such
- # as a dict having both a Tensor and a list.
- return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
-
- @add_start_docstrings_to_model_forward(DEFORMABLE_DETR_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DeformableDetrObjectDetectionOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.FloatTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[List[dict]] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], DeformableDetrObjectDetectionOutput]:
- r"""
- labels (`List[Dict]` of len `(batch_size,)`, *optional*):
- Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
- following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
- respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
- in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, DeformableDetrForObjectDetection
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("SenseTime/deformable-detr")
- >>> model = DeformableDetrForObjectDetection.from_pretrained("SenseTime/deformable-detr")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
- >>> target_sizes = torch.tensor([image.size[::-1]])
- >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[
- ... 0
- ... ]
- >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
- ... box = [round(i, 2) for i in box.tolist()]
- ... print(
- ... f"Detected {model.config.id2label[label.item()]} with confidence "
- ... f"{round(score.item(), 3)} at location {box}"
- ... )
- Detected cat with confidence 0.8 at location [16.5, 52.84, 318.25, 470.78]
- Detected cat with confidence 0.789 at location [342.19, 24.3, 640.02, 372.25]
- Detected remote with confidence 0.633 at location [40.79, 72.78, 176.76, 117.25]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # First, sent images through DETR base model to obtain encoder + decoder outputs
- outputs = self.model(
- pixel_values,
- pixel_mask=pixel_mask,
- decoder_attention_mask=decoder_attention_mask,
- encoder_outputs=encoder_outputs,
- inputs_embeds=inputs_embeds,
- decoder_inputs_embeds=decoder_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs.intermediate_hidden_states if return_dict else outputs[2]
- init_reference = outputs.init_reference_points if return_dict else outputs[0]
- inter_references = outputs.intermediate_reference_points if return_dict else outputs[3]
-
- # class logits + predicted bounding boxes
- outputs_classes = []
- outputs_coords = []
-
- for level in range(hidden_states.shape[1]):
- if level == 0:
- reference = init_reference
- else:
- reference = inter_references[:, level - 1]
- reference = inverse_sigmoid(reference)
- outputs_class = self.class_embed[level](hidden_states[:, level])
- delta_bbox = self.bbox_embed[level](hidden_states[:, level])
- if reference.shape[-1] == 4:
- outputs_coord_logits = delta_bbox + reference
- elif reference.shape[-1] == 2:
- delta_bbox[..., :2] += reference
- outputs_coord_logits = delta_bbox
- else:
- raise ValueError(f"reference.shape[-1] should be 4 or 2, but got {reference.shape[-1]}")
- outputs_coord = outputs_coord_logits.sigmoid()
- outputs_classes.append(outputs_class)
- outputs_coords.append(outputs_coord)
- outputs_class = torch.stack(outputs_classes)
- outputs_coord = torch.stack(outputs_coords)
-
- logits = outputs_class[-1]
- pred_boxes = outputs_coord[-1]
-
- loss, loss_dict, auxiliary_outputs = None, None, None
- if labels is not None:
- # First: create the matcher
- matcher = DeformableDetrHungarianMatcher(
- class_cost=self.config.class_cost, bbox_cost=self.config.bbox_cost, giou_cost=self.config.giou_cost
- )
- # Second: create the criterion
- losses = ["labels", "boxes", "cardinality"]
- criterion = DeformableDetrLoss(
- matcher=matcher,
- num_classes=self.config.num_labels,
- focal_alpha=self.config.focal_alpha,
- losses=losses,
- )
- criterion.to(self.device)
- # Third: compute the losses, based on outputs and labels
- outputs_loss = {}
- outputs_loss["logits"] = logits
- outputs_loss["pred_boxes"] = pred_boxes
- if self.config.auxiliary_loss:
- auxiliary_outputs = self._set_aux_loss(outputs_class, outputs_coord)
- outputs_loss["auxiliary_outputs"] = auxiliary_outputs
- if self.config.two_stage:
- enc_outputs_coord = outputs.enc_outputs_coord_logits.sigmoid()
- outputs_loss["enc_outputs"] = {"logits": outputs.enc_outputs_class, "pred_boxes": enc_outputs_coord}
-
- loss_dict = criterion(outputs_loss, labels)
- # Fourth: compute total loss, as a weighted sum of the various losses
- weight_dict = {"loss_ce": 1, "loss_bbox": self.config.bbox_loss_coefficient}
- weight_dict["loss_giou"] = self.config.giou_loss_coefficient
- if self.config.auxiliary_loss:
- aux_weight_dict = {}
- for i in range(self.config.decoder_layers - 1):
- aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
- weight_dict.update(aux_weight_dict)
- loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
-
- if not return_dict:
- if auxiliary_outputs is not None:
- output = (logits, pred_boxes) + auxiliary_outputs + outputs
- else:
- output = (logits, pred_boxes) + outputs
- tuple_outputs = ((loss, loss_dict) + output) if loss is not None else output
-
- return tuple_outputs
-
- dict_outputs = DeformableDetrObjectDetectionOutput(
- loss=loss,
- loss_dict=loss_dict,
- logits=logits,
- pred_boxes=pred_boxes,
- auxiliary_outputs=auxiliary_outputs,
- last_hidden_state=outputs.last_hidden_state,
- decoder_hidden_states=outputs.decoder_hidden_states,
- decoder_attentions=outputs.decoder_attentions,
- cross_attentions=outputs.cross_attentions,
- encoder_last_hidden_state=outputs.encoder_last_hidden_state,
- encoder_hidden_states=outputs.encoder_hidden_states,
- encoder_attentions=outputs.encoder_attentions,
- intermediate_hidden_states=outputs.intermediate_hidden_states,
- intermediate_reference_points=outputs.intermediate_reference_points,
- init_reference_points=outputs.init_reference_points,
- enc_outputs_class=outputs.enc_outputs_class,
- enc_outputs_coord_logits=outputs.enc_outputs_coord_logits,
- )
-
- return dict_outputs
-
-
-# Copied from transformers.models.detr.modeling_detr.dice_loss
-def dice_loss(inputs, targets, num_boxes):
- """
- Compute the DICE loss, similar to generalized IOU for masks
-
- Args:
- inputs: A float tensor of arbitrary shape.
- The predictions for each example.
- targets: A float tensor with the same shape as inputs. Stores the binary
- classification label for each element in inputs (0 for the negative class and 1 for the positive
- class).
- """
- inputs = inputs.sigmoid()
- inputs = inputs.flatten(1)
- numerator = 2 * (inputs * targets).sum(1)
- denominator = inputs.sum(-1) + targets.sum(-1)
- loss = 1 - (numerator + 1) / (denominator + 1)
- return loss.sum() / num_boxes
-
-
-# Copied from transformers.models.detr.modeling_detr.sigmoid_focal_loss
-def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2):
- """
- Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
-
- Args:
- inputs (`torch.FloatTensor` of arbitrary shape):
- The predictions for each example.
- targets (`torch.FloatTensor` with the same shape as `inputs`)
- A tensor storing the binary classification label for each element in the `inputs` (0 for the negative class
- and 1 for the positive class).
- alpha (`float`, *optional*, defaults to `0.25`):
- Optional weighting factor in the range (0,1) to balance positive vs. negative examples.
- gamma (`int`, *optional*, defaults to `2`):
- Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples.
-
- Returns:
- Loss tensor
- """
- prob = inputs.sigmoid()
- ce_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
- # add modulating factor
- p_t = prob * targets + (1 - prob) * (1 - targets)
- loss = ce_loss * ((1 - p_t) ** gamma)
-
- if alpha >= 0:
- alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
- loss = alpha_t * loss
-
- return loss.mean(1).sum() / num_boxes
-
-
-class DeformableDetrLoss(nn.Module):
- """
- This class computes the losses for `DeformableDetrForObjectDetection`. The process happens in two steps: 1) we
- compute hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair of
- matched ground-truth / prediction (supervise class and box).
-
- Args:
- matcher (`DeformableDetrHungarianMatcher`):
- Module able to compute a matching between targets and proposals.
- num_classes (`int`):
- Number of object categories, omitting the special no-object category.
- focal_alpha (`float`):
- Alpha parameter in focal loss.
- losses (`List[str]`):
- List of all the losses to be applied. See `get_loss` for a list of all available losses.
- """
-
- def __init__(self, matcher, num_classes, focal_alpha, losses):
- super().__init__()
- self.matcher = matcher
- self.num_classes = num_classes
- self.focal_alpha = focal_alpha
- self.losses = losses
-
- # removed logging parameter, which was part of the original implementation
- def loss_labels(self, outputs, targets, indices, num_boxes):
- """
- Classification loss (Binary focal loss) targets dicts must contain the key "class_labels" containing a tensor
- of dim [nb_target_boxes]
- """
- if "logits" not in outputs:
- raise KeyError("No logits were found in the outputs")
- source_logits = outputs["logits"]
-
- idx = self._get_source_permutation_idx(indices)
- target_classes_o = torch.cat([t["class_labels"][J] for t, (_, J) in zip(targets, indices)])
- target_classes = torch.full(
- source_logits.shape[:2], self.num_classes, dtype=torch.int64, device=source_logits.device
- )
- target_classes[idx] = target_classes_o
-
- target_classes_onehot = torch.zeros(
- [source_logits.shape[0], source_logits.shape[1], source_logits.shape[2] + 1],
- dtype=source_logits.dtype,
- layout=source_logits.layout,
- device=source_logits.device,
- )
- target_classes_onehot.scatter_(2, target_classes.unsqueeze(-1), 1)
-
- target_classes_onehot = target_classes_onehot[:, :, :-1]
- loss_ce = (
- sigmoid_focal_loss(source_logits, target_classes_onehot, num_boxes, alpha=self.focal_alpha, gamma=2)
- * source_logits.shape[1]
- )
- losses = {"loss_ce": loss_ce}
-
- return losses
-
- @torch.no_grad()
- # Copied from transformers.models.detr.modeling_detr.DetrLoss.loss_cardinality
- def loss_cardinality(self, outputs, targets, indices, num_boxes):
- """
- Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes.
-
- This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients.
- """
- logits = outputs["logits"]
- device = logits.device
- target_lengths = torch.as_tensor([len(v["class_labels"]) for v in targets], device=device)
- # Count the number of predictions that are NOT "no-object" (which is the last class)
- card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1)
- card_err = nn.functional.l1_loss(card_pred.float(), target_lengths.float())
- losses = {"cardinality_error": card_err}
- return losses
-
- # Copied from transformers.models.detr.modeling_detr.DetrLoss.loss_boxes
- def loss_boxes(self, outputs, targets, indices, num_boxes):
- """
- Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss.
-
- Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes
- are expected in format (center_x, center_y, w, h), normalized by the image size.
- """
- if "pred_boxes" not in outputs:
- raise KeyError("No predicted boxes found in outputs")
- idx = self._get_source_permutation_idx(indices)
- source_boxes = outputs["pred_boxes"][idx]
- target_boxes = torch.cat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0)
-
- loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none")
-
- losses = {}
- losses["loss_bbox"] = loss_bbox.sum() / num_boxes
-
- loss_giou = 1 - torch.diag(
- generalized_box_iou(center_to_corners_format(source_boxes), center_to_corners_format(target_boxes))
- )
- losses["loss_giou"] = loss_giou.sum() / num_boxes
- return losses
-
- # Copied from transformers.models.detr.modeling_detr.DetrLoss._get_source_permutation_idx
- def _get_source_permutation_idx(self, indices):
- # permute predictions following indices
- batch_idx = torch.cat([torch.full_like(source, i) for i, (source, _) in enumerate(indices)])
- source_idx = torch.cat([source for (source, _) in indices])
- return batch_idx, source_idx
-
- # Copied from transformers.models.detr.modeling_detr.DetrLoss._get_target_permutation_idx
- def _get_target_permutation_idx(self, indices):
- # permute targets following indices
- batch_idx = torch.cat([torch.full_like(target, i) for i, (_, target) in enumerate(indices)])
- target_idx = torch.cat([target for (_, target) in indices])
- return batch_idx, target_idx
-
- def get_loss(self, loss, outputs, targets, indices, num_boxes):
- loss_map = {
- "labels": self.loss_labels,
- "cardinality": self.loss_cardinality,
- "boxes": self.loss_boxes,
- }
- if loss not in loss_map:
- raise ValueError(f"Loss {loss} not supported")
- return loss_map[loss](outputs, targets, indices, num_boxes)
-
- def forward(self, outputs, targets):
- """
- This performs the loss computation.
-
- Args:
- outputs (`dict`, *optional*):
- Dictionary of tensors, see the output specification of the model for the format.
- targets (`List[dict]`, *optional*):
- List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the
- losses applied, see each loss' doc.
- """
- outputs_without_aux = {k: v for k, v in outputs.items() if k != "auxiliary_outputs" and k != "enc_outputs"}
-
- # Retrieve the matching between the outputs of the last layer and the targets
- indices = self.matcher(outputs_without_aux, targets)
-
- # Compute the average number of target boxes accross all nodes, for normalization purposes
- num_boxes = sum(len(t["class_labels"]) for t in targets)
- num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
- world_size = 1
- if is_accelerate_available():
- if PartialState._shared_state != {}:
- num_boxes = reduce(num_boxes)
- world_size = PartialState().num_processes
- num_boxes = torch.clamp(num_boxes / world_size, min=1).item()
-
- # Compute all the requested losses
- losses = {}
- for loss in self.losses:
- losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))
-
- # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
- if "auxiliary_outputs" in outputs:
- for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]):
- indices = self.matcher(auxiliary_outputs, targets)
- for loss in self.losses:
- l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes)
- l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
- losses.update(l_dict)
-
- if "enc_outputs" in outputs:
- enc_outputs = outputs["enc_outputs"]
- bin_targets = copy.deepcopy(targets)
- for bt in bin_targets:
- bt["class_labels"] = torch.zeros_like(bt["class_labels"])
- indices = self.matcher(enc_outputs, bin_targets)
- for loss in self.losses:
- l_dict = self.get_loss(loss, enc_outputs, bin_targets, indices, num_boxes)
- l_dict = {k + "_enc": v for k, v in l_dict.items()}
- losses.update(l_dict)
-
- return losses
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrMLPPredictionHead
-class DeformableDetrMLPPredictionHead(nn.Module):
- """
- Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates,
- height and width of a bounding box w.r.t. an image.
-
- Copied from https://github.com/facebookresearch/detr/blob/master/models/detr.py
-
- """
-
- def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
- super().__init__()
- self.num_layers = num_layers
- h = [hidden_dim] * (num_layers - 1)
- self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
-
- def forward(self, x):
- for i, layer in enumerate(self.layers):
- x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
- return x
-
-
-class DeformableDetrHungarianMatcher(nn.Module):
- """
- This class computes an assignment between the targets and the predictions of the network.
-
- For efficiency reasons, the targets don't include the no_object. Because of this, in general, there are more
- predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, while the others are
- un-matched (and thus treated as non-objects).
-
- Args:
- class_cost:
- The relative weight of the classification error in the matching cost.
- bbox_cost:
- The relative weight of the L1 error of the bounding box coordinates in the matching cost.
- giou_cost:
- The relative weight of the giou loss of the bounding box in the matching cost.
- """
-
- def __init__(self, class_cost: float = 1, bbox_cost: float = 1, giou_cost: float = 1):
- super().__init__()
- requires_backends(self, ["scipy"])
-
- self.class_cost = class_cost
- self.bbox_cost = bbox_cost
- self.giou_cost = giou_cost
- if class_cost == 0 and bbox_cost == 0 and giou_cost == 0:
- raise ValueError("All costs of the Matcher can't be 0")
-
- @torch.no_grad()
- def forward(self, outputs, targets):
- """
- Args:
- outputs (`dict`):
- A dictionary that contains at least these entries:
- * "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
- * "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates.
- targets (`List[dict]`):
- A list of targets (len(targets) = batch_size), where each target is a dict containing:
- * "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of
- ground-truth
- objects in the target) containing the class labels
- * "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates.
-
- Returns:
- `List[Tuple]`: A list of size `batch_size`, containing tuples of (index_i, index_j) where:
- - index_i is the indices of the selected predictions (in order)
- - index_j is the indices of the corresponding selected targets (in order)
- For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
- """
- batch_size, num_queries = outputs["logits"].shape[:2]
-
- # We flatten to compute the cost matrices in a batch
- out_prob = outputs["logits"].flatten(0, 1).sigmoid() # [batch_size * num_queries, num_classes]
- out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4]
-
- # Also concat the target labels and boxes
- target_ids = torch.cat([v["class_labels"] for v in targets])
- target_bbox = torch.cat([v["boxes"] for v in targets])
-
- # Compute the classification cost.
- alpha = 0.25
- gamma = 2.0
- neg_cost_class = (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log())
- pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log())
- class_cost = pos_cost_class[:, target_ids] - neg_cost_class[:, target_ids]
-
- # Compute the L1 cost between boxes
- bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)
-
- # Compute the giou cost between boxes
- giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))
-
- # Final cost matrix
- cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
- cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()
-
- sizes = [len(v["boxes"]) for v in targets]
- indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]
- return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
-
-
-# Copied from transformers.models.detr.modeling_detr._upcast
-def _upcast(t: Tensor) -> Tensor:
- # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
- if t.is_floating_point():
- return t if t.dtype in (torch.float32, torch.float64) else t.float()
- else:
- return t if t.dtype in (torch.int32, torch.int64) else t.int()
-
-
-# Copied from transformers.models.detr.modeling_detr.box_area
-def box_area(boxes: Tensor) -> Tensor:
- """
- Computes the area of a set of bounding boxes, which are specified by its (x1, y1, x2, y2) coordinates.
-
- Args:
- boxes (`torch.FloatTensor` of shape `(number_of_boxes, 4)`):
- Boxes for which the area will be computed. They are expected to be in (x1, y1, x2, y2) format with `0 <= x1
- < x2` and `0 <= y1 < y2`.
-
- Returns:
- `torch.FloatTensor`: a tensor containing the area for each box.
- """
- boxes = _upcast(boxes)
- return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
-
-
-# Copied from transformers.models.detr.modeling_detr.box_iou
-def box_iou(boxes1, boxes2):
- area1 = box_area(boxes1)
- area2 = box_area(boxes2)
-
- left_top = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
- right_bottom = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
-
- width_height = (right_bottom - left_top).clamp(min=0) # [N,M,2]
- inter = width_height[:, :, 0] * width_height[:, :, 1] # [N,M]
-
- union = area1[:, None] + area2 - inter
-
- iou = inter / union
- return iou, union
-
-
-# Copied from transformers.models.detr.modeling_detr.generalized_box_iou
-def generalized_box_iou(boxes1, boxes2):
- """
- Generalized IoU from https://giou.stanford.edu/. The boxes should be in [x0, y0, x1, y1] (corner) format.
-
- Returns:
- `torch.FloatTensor`: a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2)
- """
- # degenerate boxes gives inf / nan results
- # so do an early check
- if not (boxes1[:, 2:] >= boxes1[:, :2]).all():
- raise ValueError(f"boxes1 must be in [x0, y0, x1, y1] (corner) format, but got {boxes1}")
- if not (boxes2[:, 2:] >= boxes2[:, :2]).all():
- raise ValueError(f"boxes2 must be in [x0, y0, x1, y1] (corner) format, but got {boxes2}")
- iou, union = box_iou(boxes1, boxes2)
-
- top_left = torch.min(boxes1[:, None, :2], boxes2[:, :2])
- bottom_right = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
-
- width_height = (bottom_right - top_left).clamp(min=0) # [N,M,2]
- area = width_height[:, :, 0] * width_height[:, :, 1]
-
- return iou - (area - union) / area
-
-
-# Copied from transformers.models.detr.modeling_detr._max_by_axis
-def _max_by_axis(the_list):
- # type: (List[List[int]]) -> List[int]
- maxes = the_list[0]
- for sublist in the_list[1:]:
- for index, item in enumerate(sublist):
- maxes[index] = max(maxes[index], item)
- return maxes
-
-
-# Copied from transformers.models.detr.modeling_detr.NestedTensor
-class NestedTensor(object):
- def __init__(self, tensors, mask: Optional[Tensor]):
- self.tensors = tensors
- self.mask = mask
-
- def to(self, device):
- cast_tensor = self.tensors.to(device)
- mask = self.mask
- if mask is not None:
- cast_mask = mask.to(device)
- else:
- cast_mask = None
- return NestedTensor(cast_tensor, cast_mask)
-
- def decompose(self):
- return self.tensors, self.mask
-
- def __repr__(self):
- return str(self.tensors)
-
-
-# Copied from transformers.models.detr.modeling_detr.nested_tensor_from_tensor_list
-def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
- if tensor_list[0].ndim == 3:
- max_size = _max_by_axis([list(img.shape) for img in tensor_list])
- batch_shape = [len(tensor_list)] + max_size
- batch_size, num_channels, height, width = batch_shape
- dtype = tensor_list[0].dtype
- device = tensor_list[0].device
- tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
- mask = torch.ones((batch_size, height, width), dtype=torch.bool, device=device)
- for img, pad_img, m in zip(tensor_list, tensor, mask):
- pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
- m[: img.shape[1], : img.shape[2]] = False
- else:
- raise ValueError("Only 3-dimensional tensors are supported")
- return NestedTensor(tensor, mask)
diff --git a/transformers/models/deit/__init__.py b/transformers/models/deit/__init__.py
deleted file mode 100644
index a0b44186efbc05bef9faed3a47057fcfe3610862..0000000000000000000000000000000000000000
--- a/transformers/models/deit/__init__.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# Copyright 2021 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_tf_available,
- is_torch_available,
- is_vision_available,
-)
-
-
-_import_structure = {"configuration_deit": ["DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP", "DeiTConfig", "DeiTOnnxConfig"]}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_deit"] = ["DeiTFeatureExtractor"]
- _import_structure["image_processing_deit"] = ["DeiTImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_deit"] = [
- "DEIT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DeiTForImageClassification",
- "DeiTForImageClassificationWithTeacher",
- "DeiTForMaskedImageModeling",
- "DeiTModel",
- "DeiTPreTrainedModel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_deit"] = [
- "TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFDeiTForImageClassification",
- "TFDeiTForImageClassificationWithTeacher",
- "TFDeiTForMaskedImageModeling",
- "TFDeiTModel",
- "TFDeiTPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_deit import DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP, DeiTConfig, DeiTOnnxConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_deit import DeiTFeatureExtractor
- from .image_processing_deit import DeiTImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_deit import (
- DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
- DeiTForImageClassification,
- DeiTForImageClassificationWithTeacher,
- DeiTForMaskedImageModeling,
- DeiTModel,
- DeiTPreTrainedModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_deit import (
- TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFDeiTForImageClassification,
- TFDeiTForImageClassificationWithTeacher,
- TFDeiTForMaskedImageModeling,
- TFDeiTModel,
- TFDeiTPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deit/__pycache__/__init__.cpython-310.pyc b/transformers/models/deit/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 0c6266067fa88e2650c49add2d4e0828f68ce1ee..0000000000000000000000000000000000000000
Binary files a/transformers/models/deit/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deit/__pycache__/configuration_deit.cpython-310.pyc b/transformers/models/deit/__pycache__/configuration_deit.cpython-310.pyc
deleted file mode 100644
index e1bad38e6be82abba0341763d6aade5160138713..0000000000000000000000000000000000000000
Binary files a/transformers/models/deit/__pycache__/configuration_deit.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deit/__pycache__/convert_deit_timm_to_pytorch.cpython-310.pyc b/transformers/models/deit/__pycache__/convert_deit_timm_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 9e3e5139fb3567079f2050b12e11ae29546e529d..0000000000000000000000000000000000000000
Binary files a/transformers/models/deit/__pycache__/convert_deit_timm_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deit/__pycache__/feature_extraction_deit.cpython-310.pyc b/transformers/models/deit/__pycache__/feature_extraction_deit.cpython-310.pyc
deleted file mode 100644
index 55a7e9a7f486b7d5b6aabd781e3c2f180b5f9b5d..0000000000000000000000000000000000000000
Binary files a/transformers/models/deit/__pycache__/feature_extraction_deit.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deit/__pycache__/image_processing_deit.cpython-310.pyc b/transformers/models/deit/__pycache__/image_processing_deit.cpython-310.pyc
deleted file mode 100644
index 4dbb1006937f8e8c8d6ab3ee8640006127a1271c..0000000000000000000000000000000000000000
Binary files a/transformers/models/deit/__pycache__/image_processing_deit.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deit/__pycache__/modeling_deit.cpython-310.pyc b/transformers/models/deit/__pycache__/modeling_deit.cpython-310.pyc
deleted file mode 100644
index d36f64d6e60080463aa07422da4d1be5a09e9153..0000000000000000000000000000000000000000
Binary files a/transformers/models/deit/__pycache__/modeling_deit.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deit/__pycache__/modeling_tf_deit.cpython-310.pyc b/transformers/models/deit/__pycache__/modeling_tf_deit.cpython-310.pyc
deleted file mode 100644
index a243fcbeaf61c15bb402e17ad16ccfe0df753432..0000000000000000000000000000000000000000
Binary files a/transformers/models/deit/__pycache__/modeling_tf_deit.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deit/configuration_deit.py b/transformers/models/deit/configuration_deit.py
deleted file mode 100644
index 394c6ff93704ccedad80c1a20c7f8a9aa2e5a04d..0000000000000000000000000000000000000000
--- a/transformers/models/deit/configuration_deit.py
+++ /dev/null
@@ -1,142 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Facebook AI Research (FAIR) and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DeiT model configuration"""
-
-from collections import OrderedDict
-from typing import Mapping
-
-from packaging import version
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DeiTConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DeiTModel`]. It is used to instantiate an DeiT
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the DeiT
- [facebook/deit-base-distilled-patch16-224](https://huggingface.co/facebook/deit-base-distilled-patch16-224)
- architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 16):
- The size (resolution) of each patch.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- qkv_bias (`bool`, *optional*, defaults to `True`):
- Whether to add a bias to the queries, keys and values.
- encoder_stride (`int`, *optional*, defaults to 16):
- Factor to increase the spatial resolution by in the decoder head for masked image modeling.
-
- Example:
-
- ```python
- >>> from transformers import DeiTConfig, DeiTModel
-
- >>> # Initializing a DeiT deit-base-distilled-patch16-224 style configuration
- >>> configuration = DeiTConfig()
-
- >>> # Initializing a model (with random weights) from the deit-base-distilled-patch16-224 style configuration
- >>> model = DeiTModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "deit"
-
- def __init__(
- self,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.0,
- attention_probs_dropout_prob=0.0,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- image_size=224,
- patch_size=16,
- num_channels=3,
- qkv_bias=True,
- encoder_stride=16,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.qkv_bias = qkv_bias
- self.encoder_stride = encoder_stride
-
-
-class DeiTOnnxConfig(OnnxConfig):
- torch_onnx_minimum_version = version.parse("1.11")
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- return OrderedDict(
- [
- ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
- ]
- )
-
- @property
- def atol_for_validation(self) -> float:
- return 1e-4
diff --git a/transformers/models/deit/convert_deit_timm_to_pytorch.py b/transformers/models/deit/convert_deit_timm_to_pytorch.py
deleted file mode 100644
index 2b5c795ff2d2ab6d8b3e6ce6f8a0150ff3911f33..0000000000000000000000000000000000000000
--- a/transformers/models/deit/convert_deit_timm_to_pytorch.py
+++ /dev/null
@@ -1,219 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DeiT distilled checkpoints from the timm library."""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import timm
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from transformers import DeiTConfig, DeiTForImageClassificationWithTeacher, DeiTImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys(config, base_model=False):
- rename_keys = []
- for i in range(config.num_hidden_layers):
- # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
- rename_keys.append((f"blocks.{i}.norm1.weight", f"deit.encoder.layer.{i}.layernorm_before.weight"))
- rename_keys.append((f"blocks.{i}.norm1.bias", f"deit.encoder.layer.{i}.layernorm_before.bias"))
- rename_keys.append((f"blocks.{i}.attn.proj.weight", f"deit.encoder.layer.{i}.attention.output.dense.weight"))
- rename_keys.append((f"blocks.{i}.attn.proj.bias", f"deit.encoder.layer.{i}.attention.output.dense.bias"))
- rename_keys.append((f"blocks.{i}.norm2.weight", f"deit.encoder.layer.{i}.layernorm_after.weight"))
- rename_keys.append((f"blocks.{i}.norm2.bias", f"deit.encoder.layer.{i}.layernorm_after.bias"))
- rename_keys.append((f"blocks.{i}.mlp.fc1.weight", f"deit.encoder.layer.{i}.intermediate.dense.weight"))
- rename_keys.append((f"blocks.{i}.mlp.fc1.bias", f"deit.encoder.layer.{i}.intermediate.dense.bias"))
- rename_keys.append((f"blocks.{i}.mlp.fc2.weight", f"deit.encoder.layer.{i}.output.dense.weight"))
- rename_keys.append((f"blocks.{i}.mlp.fc2.bias", f"deit.encoder.layer.{i}.output.dense.bias"))
-
- # projection layer + position embeddings
- rename_keys.extend(
- [
- ("cls_token", "deit.embeddings.cls_token"),
- ("dist_token", "deit.embeddings.distillation_token"),
- ("patch_embed.proj.weight", "deit.embeddings.patch_embeddings.projection.weight"),
- ("patch_embed.proj.bias", "deit.embeddings.patch_embeddings.projection.bias"),
- ("pos_embed", "deit.embeddings.position_embeddings"),
- ]
- )
-
- if base_model:
- # layernorm + pooler
- rename_keys.extend(
- [
- ("norm.weight", "layernorm.weight"),
- ("norm.bias", "layernorm.bias"),
- ("pre_logits.fc.weight", "pooler.dense.weight"),
- ("pre_logits.fc.bias", "pooler.dense.bias"),
- ]
- )
-
- # if just the base model, we should remove "deit" from all keys that start with "deit"
- rename_keys = [(pair[0], pair[1][4:]) if pair[1].startswith("deit") else pair for pair in rename_keys]
- else:
- # layernorm + classification heads
- rename_keys.extend(
- [
- ("norm.weight", "deit.layernorm.weight"),
- ("norm.bias", "deit.layernorm.bias"),
- ("head.weight", "cls_classifier.weight"),
- ("head.bias", "cls_classifier.bias"),
- ("head_dist.weight", "distillation_classifier.weight"),
- ("head_dist.bias", "distillation_classifier.bias"),
- ]
- )
-
- return rename_keys
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config, base_model=False):
- for i in range(config.num_hidden_layers):
- if base_model:
- prefix = ""
- else:
- prefix = "deit."
- # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"blocks.{i}.attn.qkv.weight")
- in_proj_bias = state_dict.pop(f"blocks.{i}.attn.qkv.bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"{prefix}encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[
- : config.hidden_size, :
- ]
- state_dict[f"{prefix}encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size]
- state_dict[f"{prefix}encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- config.hidden_size : config.hidden_size * 2, :
- ]
- state_dict[f"{prefix}encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
- config.hidden_size : config.hidden_size * 2
- ]
- state_dict[f"{prefix}encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
- -config.hidden_size :, :
- ]
- state_dict[f"{prefix}encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :]
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-@torch.no_grad()
-def convert_deit_checkpoint(deit_name, pytorch_dump_folder_path):
- """
- Copy/paste/tweak model's weights to our DeiT structure.
- """
-
- # define default DeiT configuration
- config = DeiTConfig()
- # all deit models have fine-tuned heads
- base_model = False
- # dataset (fine-tuned on ImageNet 2012), patch_size and image_size
- config.num_labels = 1000
- repo_id = "huggingface/label-files"
- filename = "imagenet-1k-id2label.json"
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
- config.patch_size = int(deit_name[-6:-4])
- config.image_size = int(deit_name[-3:])
- # size of the architecture
- if deit_name[9:].startswith("tiny"):
- config.hidden_size = 192
- config.intermediate_size = 768
- config.num_hidden_layers = 12
- config.num_attention_heads = 3
- elif deit_name[9:].startswith("small"):
- config.hidden_size = 384
- config.intermediate_size = 1536
- config.num_hidden_layers = 12
- config.num_attention_heads = 6
- if deit_name[9:].startswith("base"):
- pass
- elif deit_name[4:].startswith("large"):
- config.hidden_size = 1024
- config.intermediate_size = 4096
- config.num_hidden_layers = 24
- config.num_attention_heads = 16
-
- # load original model from timm
- timm_model = timm.create_model(deit_name, pretrained=True)
- timm_model.eval()
-
- # load state_dict of original model, remove and rename some keys
- state_dict = timm_model.state_dict()
- rename_keys = create_rename_keys(config, base_model)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- read_in_q_k_v(state_dict, config, base_model)
-
- # load HuggingFace model
- model = DeiTForImageClassificationWithTeacher(config).eval()
- model.load_state_dict(state_dict)
-
- # Check outputs on an image, prepared by DeiTImageProcessor
- size = int(
- (256 / 224) * config.image_size
- ) # to maintain same ratio w.r.t. 224 images, see https://github.com/facebookresearch/deit/blob/ab5715372db8c6cad5740714b2216d55aeae052e/datasets.py#L103
- image_processor = DeiTImageProcessor(size=size, crop_size=config.image_size)
- encoding = image_processor(images=prepare_img(), return_tensors="pt")
- pixel_values = encoding["pixel_values"]
- outputs = model(pixel_values)
-
- timm_logits = timm_model(pixel_values)
- assert timm_logits.shape == outputs.logits.shape
- assert torch.allclose(timm_logits, outputs.logits, atol=1e-3)
-
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model {deit_name} to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- print(f"Saving image processor to {pytorch_dump_folder_path}")
- image_processor.save_pretrained(pytorch_dump_folder_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--deit_name",
- default="vit_deit_base_distilled_patch16_224",
- type=str,
- help="Name of the DeiT timm model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
- )
-
- args = parser.parse_args()
- convert_deit_checkpoint(args.deit_name, args.pytorch_dump_folder_path)
diff --git a/transformers/models/deit/feature_extraction_deit.py b/transformers/models/deit/feature_extraction_deit.py
deleted file mode 100644
index b66922ea95753a81b93a3f9c99607119017df3f3..0000000000000000000000000000000000000000
--- a/transformers/models/deit/feature_extraction_deit.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for DeiT."""
-
-import warnings
-
-from ...utils import logging
-from .image_processing_deit import DeiTImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-class DeiTFeatureExtractor(DeiTImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class DeiTFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
- " use DeiTImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/deit/image_processing_deit.py b/transformers/models/deit/image_processing_deit.py
deleted file mode 100644
index 2a8ebb36377854aa80bf8505c7e98b1eb661648a..0000000000000000000000000000000000000000
--- a/transformers/models/deit/image_processing_deit.py
+++ /dev/null
@@ -1,320 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for DeiT."""
-
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
-from ...image_transforms import resize, to_channel_dimension_format
-from ...image_utils import (
- IMAGENET_STANDARD_MEAN,
- IMAGENET_STANDARD_STD,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import TensorType, is_vision_available, logging
-
-
-if is_vision_available():
- import PIL
-
-
-logger = logging.get_logger(__name__)
-
-
-class DeiTImageProcessor(BaseImageProcessor):
- r"""
- Constructs a DeiT image processor.
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by
- `do_resize` in `preprocess`.
- size (`Dict[str, int]` *optional*, defaults to `{"height": 256, "width": 256}`):
- Size of the image after `resize`. Can be overridden by `size` in `preprocess`.
- resample (`PILImageResampling` filter, *optional*, defaults to `Resampling.BICUBIC`):
- Resampling filter to use if resizing the image. Can be overridden by `resample` in `preprocess`.
- do_center_crop (`bool`, *optional*, defaults to `True`):
- Whether to center crop the image. If the input size is smaller than `crop_size` along any edge, the image
- is padded with 0's and then center cropped. Can be overridden by `do_center_crop` in `preprocess`.
- crop_size (`Dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
- Desired output size when applying center-cropping. Can be overridden by `crop_size` in `preprocess`.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the
- `preprocess` method.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`
- parameter in the `preprocess` method.
- do_normalize (`bool`, *optional*, defaults to `True`):
- Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
- method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
- Mean to use if normalizing the image. This is a float or list of floats the length of the number of
- channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
- Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
- number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
- """
-
- model_input_names = ["pixel_values"]
-
- def __init__(
- self,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PIL.Image.BICUBIC,
- do_center_crop: bool = True,
- crop_size: Dict[str, int] = None,
- rescale_factor: Union[int, float] = 1 / 255,
- do_rescale: bool = True,
- do_normalize: bool = True,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- size = size if size is not None else {"height": 256, "width": 256}
- size = get_size_dict(size)
- crop_size = crop_size if crop_size is not None else {"height": 224, "width": 224}
- crop_size = get_size_dict(crop_size, param_name="crop_size")
-
- self.do_resize = do_resize
- self.size = size
- self.resample = resample
- self.do_center_crop = do_center_crop
- self.crop_size = crop_size
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_normalize = do_normalize
- self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size",
- "resample",
- "do_center_crop",
- "crop_size",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "image_mean",
- "image_std",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- # Copied from transformers.models.vit.image_processing_vit.ViTImageProcessor.resize with PILImageResampling.BILINEAR->PILImageResampling.BICUBIC
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize an image to `(size["height"], size["width"])`.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
- `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BICUBIC`.
- data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
-
- Returns:
- `np.ndarray`: The resized image.
- """
- size = get_size_dict(size)
- if "height" not in size or "width" not in size:
- raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}")
- output_size = (size["height"], size["width"])
- return resize(
- image,
- size=output_size,
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
-
- def preprocess(
- self,
- images: ImageInput,
- do_resize: bool = None,
- size: Dict[str, int] = None,
- resample=None,
- do_center_crop: bool = None,
- crop_size: Dict[str, int] = None,
- do_rescale: bool = None,
- rescale_factor: float = None,
- do_normalize: bool = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: ChannelDimension = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> PIL.Image.Image:
- """
- Preprocess an image or batch of images.
-
- Args:
- images (`ImageInput`):
- Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
- passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
- Size of the image after `resize`.
- resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
- PILImageResampling filter to use if resizing the image Only has an effect if `do_resize` is set to
- `True`.
- do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`):
- Whether to center crop the image.
- crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
- Size of the image after center crop. If one edge the image is smaller than `crop_size`, it will be
- padded with zeros and then cropped
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether to rescale the image values between [0 - 1].
- rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
- Rescale factor to rescale the image by if `do_rescale` is set to `True`.
- do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
- Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
- Image mean.
- image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
- Image standard deviation.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - `None`: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- do_resize = do_resize if do_resize is not None else self.do_resize
- resample = resample if resample is not None else self.resample
- do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- do_normalize = do_normalize if do_normalize is not None else self.do_normalize
- image_mean = image_mean if image_mean is not None else self.image_mean
- image_std = image_std if image_std is not None else self.image_std
-
- size = size if size is not None else self.size
- size = get_size_dict(size)
- crop_size = crop_size if crop_size is not None else self.crop_size
- crop_size = get_size_dict(crop_size, param_name="crop_size")
-
- images = make_list_of_images(images)
-
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_center_crop=do_center_crop,
- crop_size=crop_size,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
- # All transformations expect numpy arrays.
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- if do_resize:
- images = [
- self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_center_crop:
- images = [
- self.center_crop(image=image, size=crop_size, input_data_format=input_data_format) for image in images
- ]
-
- if do_rescale:
- images = [
- self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_normalize:
- images = [
- self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
- for image in images
- ]
-
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
- ]
-
- data = {"pixel_values": images}
- return BatchFeature(data=data, tensor_type=return_tensors)
diff --git a/transformers/models/deit/modeling_deit.py b/transformers/models/deit/modeling_deit.py
deleted file mode 100644
index 5efcc95d503da407426611b2b67d5c76ddde73e8..0000000000000000000000000000000000000000
--- a/transformers/models/deit/modeling_deit.py
+++ /dev/null
@@ -1,891 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Facebook AI Research (FAIR), Ross Wightman, The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DeiT model."""
-
-
-import collections.abc
-import math
-from dataclasses import dataclass
-from typing import Optional, Set, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutput,
- BaseModelOutputWithPooling,
- ImageClassifierOutput,
- MaskedImageModelingOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_deit import DeiTConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "DeiTConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/deit-base-distilled-patch16-224"
-_EXPECTED_OUTPUT_SHAPE = [1, 198, 768]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "facebook/deit-base-distilled-patch16-224"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class DeiTEmbeddings(nn.Module):
- """
- Construct the CLS token, distillation token, position and patch embeddings. Optionally, also the mask token.
- """
-
- def __init__(self, config: DeiTConfig, use_mask_token: bool = False) -> None:
- super().__init__()
-
- self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
- self.distillation_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
- self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) if use_mask_token else None
- self.patch_embeddings = DeiTPatchEmbeddings(config)
- num_patches = self.patch_embeddings.num_patches
- self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 2, config.hidden_size))
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.BoolTensor] = None) -> torch.Tensor:
- embeddings = self.patch_embeddings(pixel_values)
- batch_size, seq_length, _ = embeddings.size()
-
- if bool_masked_pos is not None:
- mask_tokens = self.mask_token.expand(batch_size, seq_length, -1)
- # replace the masked visual tokens by mask_tokens
- mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
- embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
-
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- distillation_tokens = self.distillation_token.expand(batch_size, -1, -1)
- embeddings = torch.cat((cls_tokens, distillation_tokens, embeddings), dim=1)
- embeddings = embeddings + self.position_embeddings
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-class DeiTPatchEmbeddings(nn.Module):
- """
- This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
- `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
- Transformer.
- """
-
- def __init__(self, config):
- super().__init__()
- image_size, patch_size = config.image_size, config.patch_size
- num_channels, hidden_size = config.num_channels, config.hidden_size
-
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.num_patches = num_patches
-
- self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
-
- def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
- batch_size, num_channels, height, width = pixel_values.shape
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- if height != self.image_size[0] or width != self.image_size[1]:
- raise ValueError(
- f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."
- )
- x = self.projection(pixel_values).flatten(2).transpose(1, 2)
- return x
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->DeiT
-class DeiTSelfAttention(nn.Module):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
- f"heads {config.num_attention_heads}."
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
- self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
- self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self, hidden_states, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- mixed_query_layer = self.query(hidden_states)
-
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->DeiT
-class DeiTSelfOutput(nn.Module):
- """
- The residual connection is defined in DeiTLayer instead of here (as is the case with other models), due to the
- layernorm applied before each block.
- """
-
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->DeiT
-class DeiTAttention(nn.Module):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__()
- self.attention = DeiTSelfAttention(config)
- self.output = DeiTSelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads: Set[int]) -> None:
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.attention.query = prune_linear_layer(self.attention.query, index)
- self.attention.key = prune_linear_layer(self.attention.key, index)
- self.attention.value = prune_linear_layer(self.attention.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
- self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- self_outputs = self.attention(hidden_states, head_mask, output_attentions)
-
- attention_output = self.output(self_outputs[0], hidden_states)
-
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTIntermediate with ViT->DeiT
-class DeiTIntermediate(nn.Module):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTOutput with ViT->DeiT
-class DeiTOutput(nn.Module):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- hidden_states = hidden_states + input_tensor
-
- return hidden_states
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTLayer with ViT->DeiT
-class DeiTLayer(nn.Module):
- """This corresponds to the Block class in the timm implementation."""
-
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = DeiTAttention(config)
- self.intermediate = DeiTIntermediate(config)
- self.output = DeiTOutput(config)
- self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- self_attention_outputs = self.attention(
- self.layernorm_before(hidden_states), # in DeiT, layernorm is applied before self-attention
- head_mask,
- output_attentions=output_attentions,
- )
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- # first residual connection
- hidden_states = attention_output + hidden_states
-
- # in DeiT, layernorm is also applied after self-attention
- layer_output = self.layernorm_after(hidden_states)
- layer_output = self.intermediate(layer_output)
-
- # second residual connection is done here
- layer_output = self.output(layer_output, hidden_states)
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTEncoder with ViT->DeiT
-class DeiTEncoder(nn.Module):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([DeiTLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[tuple, BaseModelOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- layer_head_mask,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions)
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-class DeiTPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DeiTConfig
- base_model_prefix = "deit"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
- _no_split_modules = ["DeiTLayer"]
-
- def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid
- # `trunc_normal_cpu` not implemented in `half` issues
- module.weight.data = nn.init.trunc_normal_(
- module.weight.data.to(torch.float32), mean=0.0, std=self.config.initializer_range
- ).to(module.weight.dtype)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-DEIT_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`DeiTConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DEIT_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`DeiTImageProcessor.__call__`] for details.
-
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DeiT Model transformer outputting raw hidden-states without any specific head on top.",
- DEIT_START_DOCSTRING,
-)
-class DeiTModel(DeiTPreTrainedModel):
- def __init__(self, config: DeiTConfig, add_pooling_layer: bool = True, use_mask_token: bool = False) -> None:
- super().__init__(config)
- self.config = config
-
- self.embeddings = DeiTEmbeddings(config, use_mask_token=use_mask_token)
- self.encoder = DeiTEncoder(config)
-
- self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.pooler = DeiTPooler(config) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self) -> DeiTPatchEmbeddings:
- return self.embeddings.patch_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(DEIT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPooling,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPooling]:
- r"""
- bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- # TODO: maybe have a cleaner way to cast the input (from `ImageProcessor` side?)
- expected_dtype = self.embeddings.patch_embeddings.projection.weight.dtype
- if pixel_values.dtype != expected_dtype:
- pixel_values = pixel_values.to(expected_dtype)
-
- embedding_output = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
-
- encoder_outputs = self.encoder(
- embedding_output,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,)
- return head_outputs + encoder_outputs[1:]
-
- return BaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTPooler with ViT->DeiT
-class DeiTPooler(nn.Module):
- def __init__(self, config: DeiTConfig):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = nn.Tanh()
-
- def forward(self, hidden_states):
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(first_token_tensor)
- pooled_output = self.activation(pooled_output)
- return pooled_output
-
-
-@add_start_docstrings(
- """DeiT Model with a decoder on top for masked image modeling, as proposed in [SimMIM](https://arxiv.org/abs/2111.09886).
-
-
-
- Note that we provide a script to pre-train this model on custom data in our [examples
- directory](https://github.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining).
-
-
- """,
- DEIT_START_DOCSTRING,
-)
-class DeiTForMaskedImageModeling(DeiTPreTrainedModel):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__(config)
-
- self.deit = DeiTModel(config, add_pooling_layer=False, use_mask_token=True)
-
- self.decoder = nn.Sequential(
- nn.Conv2d(
- in_channels=config.hidden_size,
- out_channels=config.encoder_stride**2 * config.num_channels,
- kernel_size=1,
- ),
- nn.PixelShuffle(config.encoder_stride),
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DEIT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=MaskedImageModelingOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, MaskedImageModelingOutput]:
- r"""
- bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
-
- Returns:
-
- Examples:
- ```python
- >>> from transformers import AutoImageProcessor, DeiTForMaskedImageModeling
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224")
- >>> model = DeiTForMaskedImageModeling.from_pretrained("facebook/deit-base-distilled-patch16-224")
-
- >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2
- >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
- >>> # create random boolean mask of shape (batch_size, num_patches)
- >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
-
- >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
- >>> loss, reconstructed_pixel_values = outputs.loss, outputs.reconstruction
- >>> list(reconstructed_pixel_values.shape)
- [1, 3, 224, 224]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deit(
- pixel_values,
- bool_masked_pos=bool_masked_pos,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- # Reshape to (batch_size, num_channels, height, width)
- sequence_output = sequence_output[:, 1:-1]
- batch_size, sequence_length, num_channels = sequence_output.shape
- height = width = int(sequence_length**0.5)
- sequence_output = sequence_output.permute(0, 2, 1).reshape(batch_size, num_channels, height, width)
-
- # Reconstruct pixel values
- reconstructed_pixel_values = self.decoder(sequence_output)
-
- masked_im_loss = None
- if bool_masked_pos is not None:
- size = self.config.image_size // self.config.patch_size
- bool_masked_pos = bool_masked_pos.reshape(-1, size, size)
- mask = (
- bool_masked_pos.repeat_interleave(self.config.patch_size, 1)
- .repeat_interleave(self.config.patch_size, 2)
- .unsqueeze(1)
- .contiguous()
- )
- reconstruction_loss = nn.functional.l1_loss(pixel_values, reconstructed_pixel_values, reduction="none")
- masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / self.config.num_channels
-
- if not return_dict:
- output = (reconstructed_pixel_values,) + outputs[1:]
- return ((masked_im_loss,) + output) if masked_im_loss is not None else output
-
- return MaskedImageModelingOutput(
- loss=masked_im_loss,
- reconstruction=reconstructed_pixel_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DeiT Model transformer with an image classification head on top (a linear layer on top of the final hidden state of
- the [CLS] token) e.g. for ImageNet.
- """,
- DEIT_START_DOCSTRING,
-)
-class DeiTForImageClassification(DeiTPreTrainedModel):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.deit = DeiTModel(config, add_pooling_layer=False)
-
- # Classifier head
- self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DEIT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=ImageClassifierOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, ImageClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, DeiTForImageClassification
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> torch.manual_seed(3) # doctest: +IGNORE_RESULT
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> # note: we are loading a DeiTForImageClassificationWithTeacher from the hub here,
- >>> # so the head will be randomly initialized, hence the predictions will be random
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224")
- >>> model = DeiTForImageClassification.from_pretrained("facebook/deit-base-distilled-patch16-224")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
- >>> outputs = model(**inputs)
- >>> logits = outputs.logits
- >>> # model predicts one of the 1000 ImageNet classes
- >>> predicted_class_idx = logits.argmax(-1).item()
- >>> print("Predicted class:", model.config.id2label[predicted_class_idx])
- Predicted class: Polaroid camera, Polaroid Land camera
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deit(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.classifier(sequence_output[:, 0, :])
- # we don't use the distillation token
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@dataclass
-class DeiTForImageClassificationWithTeacherOutput(ModelOutput):
- """
- Output type of [`DeiTForImageClassificationWithTeacher`].
-
- Args:
- logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores as the average of the cls_logits and distillation logits.
- cls_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores of the classification head (i.e. the linear layer on top of the final hidden state of the
- class token).
- distillation_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores of the distillation head (i.e. the linear layer on top of the final hidden state of the
- distillation token).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
- the self-attention heads.
- """
-
- logits: torch.FloatTensor = None
- cls_logits: torch.FloatTensor = None
- distillation_logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@add_start_docstrings(
- """
- DeiT Model transformer with image classification heads on top (a linear layer on top of the final hidden state of
- the [CLS] token and a linear layer on top of the final hidden state of the distillation token) e.g. for ImageNet.
-
- .. warning::
-
- This model supports inference-only. Fine-tuning with distillation (i.e. with a teacher) is not yet
- supported.
- """,
- DEIT_START_DOCSTRING,
-)
-class DeiTForImageClassificationWithTeacher(DeiTPreTrainedModel):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.deit = DeiTModel(config, add_pooling_layer=False)
-
- # Classifier heads
- self.cls_classifier = (
- nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
- self.distillation_classifier = (
- nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DEIT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=DeiTForImageClassificationWithTeacherOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, DeiTForImageClassificationWithTeacherOutput]:
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deit(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- cls_logits = self.cls_classifier(sequence_output[:, 0, :])
- distillation_logits = self.distillation_classifier(sequence_output[:, 1, :])
-
- # during inference, return the average of both classifier predictions
- logits = (cls_logits + distillation_logits) / 2
-
- if not return_dict:
- output = (logits, cls_logits, distillation_logits) + outputs[1:]
- return output
-
- return DeiTForImageClassificationWithTeacherOutput(
- logits=logits,
- cls_logits=cls_logits,
- distillation_logits=distillation_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/deit/modeling_tf_deit.py b/transformers/models/deit/modeling_tf_deit.py
deleted file mode 100644
index aec5f6df95922aac77bd8e210b74c37cc17beebe..0000000000000000000000000000000000000000
--- a/transformers/models/deit/modeling_tf_deit.py
+++ /dev/null
@@ -1,1178 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Facebook AI Research (FAIR) and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TensorFlow DeiT model."""
-
-
-from __future__ import annotations
-
-import collections.abc
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFBaseModelOutputWithPooling,
- TFImageClassifierOutput,
- TFMaskedImageModelingOutput,
-)
-from ...modeling_tf_utils import (
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import shape_list, stable_softmax
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_deit import DeiTConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "DeiTConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/deit-base-distilled-patch16-224"
-_EXPECTED_OUTPUT_SHAPE = [1, 198, 768]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "facebook/deit-base-distilled-patch16-224"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class TFDeiTForImageClassificationWithTeacherOutput(ModelOutput):
- """
- Output type of [`DeiTForImageClassificationWithTeacher`].
-
- Args:
- logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores as the average of the cls_logits and distillation logits.
- cls_logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores of the classification head (i.e. the linear layer on top of the final hidden state of the
- class token).
- distillation_logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores of the distillation head (i.e. the linear layer on top of the final hidden state of the
- distillation token).
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus
- the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
- the self-attention heads.
- """
-
- logits: tf.Tensor = None
- cls_logits: tf.Tensor = None
- distillation_logits: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-class TFDeiTEmbeddings(keras.layers.Layer):
- """
- Construct the CLS token, distillation token, position and patch embeddings. Optionally, also the mask token.
- """
-
- def __init__(self, config: DeiTConfig, use_mask_token: bool = False, **kwargs) -> None:
- super().__init__(**kwargs)
- self.config = config
- self.use_mask_token = use_mask_token
- self.patch_embeddings = TFDeiTPatchEmbeddings(config=config, name="patch_embeddings")
- self.dropout = keras.layers.Dropout(config.hidden_dropout_prob, name="dropout")
-
- def build(self, input_shape=None):
- self.cls_token = self.add_weight(
- shape=(1, 1, self.config.hidden_size),
- initializer=keras.initializers.zeros(),
- trainable=True,
- name="cls_token",
- )
- self.distillation_token = self.add_weight(
- shape=(1, 1, self.config.hidden_size),
- initializer=keras.initializers.zeros(),
- trainable=True,
- name="distillation_token",
- )
- self.mask_token = None
- if self.use_mask_token:
- self.mask_token = self.add_weight(
- shape=(1, 1, self.config.hidden_size),
- initializer=keras.initializers.zeros(),
- trainable=True,
- name="mask_token",
- )
- num_patches = self.patch_embeddings.num_patches
- self.position_embeddings = self.add_weight(
- shape=(1, num_patches + 2, self.config.hidden_size),
- initializer=keras.initializers.zeros(),
- trainable=True,
- name="position_embeddings",
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "patch_embeddings", None) is not None:
- with tf.name_scope(self.patch_embeddings.name):
- self.patch_embeddings.build(None)
- if getattr(self, "dropout", None) is not None:
- with tf.name_scope(self.dropout.name):
- self.dropout.build(None)
-
- def call(
- self, pixel_values: tf.Tensor, bool_masked_pos: tf.Tensor | None = None, training: bool = False
- ) -> tf.Tensor:
- embeddings = self.patch_embeddings(pixel_values)
- batch_size, seq_length, _ = shape_list(embeddings)
-
- if bool_masked_pos is not None:
- mask_tokens = tf.tile(self.mask_token, [batch_size, seq_length, 1])
- # replace the masked visual tokens by mask_tokens
- mask = tf.expand_dims(bool_masked_pos, axis=-1)
- mask = tf.cast(mask, dtype=mask_tokens.dtype)
- embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
-
- cls_tokens = tf.repeat(self.cls_token, repeats=batch_size, axis=0)
- distillation_tokens = tf.repeat(self.distillation_token, repeats=batch_size, axis=0)
- embeddings = tf.concat((cls_tokens, distillation_tokens, embeddings), axis=1)
- embeddings = embeddings + self.position_embeddings
- embeddings = self.dropout(embeddings, training=training)
- return embeddings
-
-
-class TFDeiTPatchEmbeddings(keras.layers.Layer):
- """
- This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
- `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
- Transformer.
- """
-
- def __init__(self, config: DeiTConfig, **kwargs) -> None:
- super().__init__(**kwargs)
- image_size, patch_size = config.image_size, config.patch_size
- num_channels, hidden_size = config.num_channels, config.hidden_size
-
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.num_patches = num_patches
-
- self.projection = keras.layers.Conv2D(
- hidden_size, kernel_size=patch_size, strides=patch_size, name="projection"
- )
-
- def call(self, pixel_values: tf.Tensor) -> tf.Tensor:
- batch_size, height, width, num_channels = shape_list(pixel_values)
- if tf.executing_eagerly() and num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- if tf.executing_eagerly() and (height != self.image_size[0] or width != self.image_size[1]):
- raise ValueError(
- f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."
- )
- x = self.projection(pixel_values)
- batch_size, height, width, num_channels = shape_list(x)
- x = tf.reshape(x, (batch_size, height * width, num_channels))
- return x
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "projection", None) is not None:
- with tf.name_scope(self.projection.name):
- self.projection.build([None, None, None, self.num_channels])
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTSelfAttention with ViT->DeiT
-class TFDeiTSelfAttention(keras.layers.Layer):
- def __init__(self, config: DeiTConfig, **kwargs):
- super().__init__(**kwargs)
-
- if config.hidden_size % config.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number "
- f"of attention heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.sqrt_att_head_size = math.sqrt(self.attention_head_size)
-
- self.query = keras.layers.Dense(
- units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"
- )
- self.key = keras.layers.Dense(
- units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key"
- )
- self.value = keras.layers.Dense(
- units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"
- )
- self.dropout = keras.layers.Dropout(rate=config.attention_probs_dropout_prob)
- self.config = config
-
- def transpose_for_scores(self, tensor: tf.Tensor, batch_size: int) -> tf.Tensor:
- # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]
- tensor = tf.reshape(tensor=tensor, shape=(batch_size, -1, self.num_attention_heads, self.attention_head_size))
-
- # Transpose the tensor from [batch_size, seq_length, num_attention_heads, attention_head_size] to [batch_size, num_attention_heads, seq_length, attention_head_size]
- return tf.transpose(tensor, perm=[0, 2, 1, 3])
-
- def call(
- self,
- hidden_states: tf.Tensor,
- head_mask: tf.Tensor,
- output_attentions: bool,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- batch_size = shape_list(hidden_states)[0]
- mixed_query_layer = self.query(inputs=hidden_states)
- mixed_key_layer = self.key(inputs=hidden_states)
- mixed_value_layer = self.value(inputs=hidden_states)
- query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
- key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
- value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- # (batch size, num_heads, seq_len_q, seq_len_k)
- attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
- dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
- attention_scores = tf.divide(attention_scores, dk)
-
- # Normalize the attention scores to probabilities.
- attention_probs = stable_softmax(logits=attention_scores, axis=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(inputs=attention_probs, training=training)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = tf.multiply(attention_probs, head_mask)
-
- attention_output = tf.matmul(attention_probs, value_layer)
- attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])
-
- # (batch_size, seq_len_q, all_head_size)
- attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
- outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "query", None) is not None:
- with tf.name_scope(self.query.name):
- self.query.build([None, None, self.config.hidden_size])
- if getattr(self, "key", None) is not None:
- with tf.name_scope(self.key.name):
- self.key.build([None, None, self.config.hidden_size])
- if getattr(self, "value", None) is not None:
- with tf.name_scope(self.value.name):
- self.value.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTSelfOutput with ViT->DeiT
-class TFDeiTSelfOutput(keras.layers.Layer):
- """
- The residual connection is defined in TFDeiTLayer instead of here (as is the case with other models), due to the
- layernorm applied before each block.
- """
-
- def __init__(self, config: DeiTConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.dropout(inputs=hidden_states, training=training)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTAttention with ViT->DeiT
-class TFDeiTAttention(keras.layers.Layer):
- def __init__(self, config: DeiTConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.self_attention = TFDeiTSelfAttention(config, name="attention")
- self.dense_output = TFDeiTSelfOutput(config, name="output")
-
- def prune_heads(self, heads):
- raise NotImplementedError
-
- def call(
- self,
- input_tensor: tf.Tensor,
- head_mask: tf.Tensor,
- output_attentions: bool,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- self_outputs = self.self_attention(
- hidden_states=input_tensor, head_mask=head_mask, output_attentions=output_attentions, training=training
- )
- attention_output = self.dense_output(
- hidden_states=self_outputs[0], input_tensor=input_tensor, training=training
- )
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "self_attention", None) is not None:
- with tf.name_scope(self.self_attention.name):
- self.self_attention.build(None)
- if getattr(self, "dense_output", None) is not None:
- with tf.name_scope(self.dense_output.name):
- self.dense_output.build(None)
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTIntermediate with ViT->DeiT
-class TFDeiTIntermediate(keras.layers.Layer):
- def __init__(self, config: DeiTConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
-
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.intermediate_act_fn = config.hidden_act
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTOutput with ViT->DeiT
-class TFDeiTOutput(keras.layers.Layer):
- def __init__(self, config: DeiTConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.dropout(inputs=hidden_states, training=training)
- hidden_states = hidden_states + input_tensor
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.intermediate_size])
-
-
-class TFDeiTLayer(keras.layers.Layer):
- """This corresponds to the Block class in the timm implementation."""
-
- def __init__(self, config: DeiTConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.attention = TFDeiTAttention(config, name="attention")
- self.intermediate = TFDeiTIntermediate(config, name="intermediate")
- self.deit_output = TFDeiTOutput(config, name="output")
-
- self.layernorm_before = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm_before")
- self.layernorm_after = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm_after")
- self.config = config
-
- def call(
- self,
- hidden_states: tf.Tensor,
- head_mask: tf.Tensor,
- output_attentions: bool,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- attention_outputs = self.attention(
- # in DeiT, layernorm is applied before self-attention
- input_tensor=self.layernorm_before(inputs=hidden_states, training=training),
- head_mask=head_mask,
- output_attentions=output_attentions,
- training=training,
- )
- attention_output = attention_outputs[0]
-
- # first residual connection
- hidden_states = attention_output + hidden_states
-
- # in DeiT, layernorm is also applied after self-attention
- layer_output = self.layernorm_after(inputs=hidden_states, training=training)
-
- intermediate_output = self.intermediate(hidden_states=layer_output, training=training)
-
- # second residual connection is done here
- layer_output = self.deit_output(
- hidden_states=intermediate_output, input_tensor=hidden_states, training=training
- )
- outputs = (layer_output,) + attention_outputs[1:] # add attentions if we output them
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "intermediate", None) is not None:
- with tf.name_scope(self.intermediate.name):
- self.intermediate.build(None)
- if getattr(self, "deit_output", None) is not None:
- with tf.name_scope(self.deit_output.name):
- self.deit_output.build(None)
- if getattr(self, "layernorm_before", None) is not None:
- with tf.name_scope(self.layernorm_before.name):
- self.layernorm_before.build([None, None, self.config.hidden_size])
- if getattr(self, "layernorm_after", None) is not None:
- with tf.name_scope(self.layernorm_after.name):
- self.layernorm_after.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTEncoder with ViT->DeiT
-class TFDeiTEncoder(keras.layers.Layer):
- def __init__(self, config: DeiTConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.layer = [TFDeiTLayer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]
-
- def call(
- self,
- hidden_states: tf.Tensor,
- head_mask: tf.Tensor,
- output_attentions: bool,
- output_hidden_states: bool,
- return_dict: bool,
- training: bool = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_outputs = layer_module(
- hidden_states=hidden_states,
- head_mask=head_mask[i],
- output_attentions=output_attentions,
- training=training,
- )
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- # Add last layer
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)
-
- return TFBaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "layer", None) is not None:
- for layer in self.layer:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-@keras_serializable
-class TFDeiTMainLayer(keras.layers.Layer):
- config_class = DeiTConfig
-
- def __init__(
- self, config: DeiTConfig, add_pooling_layer: bool = True, use_mask_token: bool = False, **kwargs
- ) -> None:
- super().__init__(**kwargs)
- self.config = config
-
- self.embeddings = TFDeiTEmbeddings(config, use_mask_token=use_mask_token, name="embeddings")
- self.encoder = TFDeiTEncoder(config, name="encoder")
-
- self.layernorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
- self.pooler = TFDeiTPooler(config, name="pooler") if add_pooling_layer else None
-
- def get_input_embeddings(self) -> TFDeiTPatchEmbeddings:
- return self.embeddings.patch_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- raise NotImplementedError
-
- def get_head_mask(self, head_mask):
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.config.num_hidden_layers
-
- return head_mask
-
- @unpack_inputs
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- bool_masked_pos: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor, ...]]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- # TF 2.0 image layers can't use NCHW format when running on CPU.
- # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
- pixel_values = tf.transpose(pixel_values, (0, 2, 3, 1))
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask)
-
- embedding_output = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos, training=training)
-
- encoder_outputs = self.encoder(
- embedding_output,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output, training=training)
- pooled_output = self.pooler(sequence_output, training=training) if self.pooler is not None else None
-
- if not return_dict:
- head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,)
- return head_outputs + encoder_outputs[1:]
-
- return TFBaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, None, self.config.hidden_size])
- if getattr(self, "pooler", None) is not None:
- with tf.name_scope(self.pooler.name):
- self.pooler.build(None)
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTPreTrainedModel with ViT->DeiT all-casing
-class TFDeiTPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DeiTConfig
- base_model_prefix = "deit"
- main_input_name = "pixel_values"
-
-
-DEIT_START_DOCSTRING = r"""
- This model is a TensorFlow
- [keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer). Use it as a regular
- TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and behavior.
-
- Parameters:
- config ([`DeiTConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DEIT_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`DeiTImageProcessor.__call__`] for details.
-
- head_mask (`tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DeiT Model transformer outputting raw hidden-states without any specific head on top.",
- DEIT_START_DOCSTRING,
-)
-class TFDeiTModel(TFDeiTPreTrainedModel):
- def __init__(
- self, config: DeiTConfig, add_pooling_layer: bool = True, use_mask_token: bool = False, **kwargs
- ) -> None:
- super().__init__(config, **kwargs)
-
- self.deit = TFDeiTMainLayer(
- config, add_pooling_layer=add_pooling_layer, use_mask_token=use_mask_token, name="deit"
- )
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEIT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutputWithPooling,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- bool_masked_pos: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[Tuple, TFBaseModelOutputWithPooling]:
- outputs = self.deit(
- pixel_values=pixel_values,
- bool_masked_pos=bool_masked_pos,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deit", None) is not None:
- with tf.name_scope(self.deit.name):
- self.deit.build(None)
-
-
-# Copied from transformers.models.vit.modeling_tf_vit.TFViTPooler with ViT->DeiT
-class TFDeiTPooler(keras.layers.Layer):
- def __init__(self, config: DeiTConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size,
- kernel_initializer=get_initializer(config.initializer_range),
- activation="tanh",
- name="dense",
- )
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(inputs=first_token_tensor)
-
- return pooled_output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFDeitPixelShuffle(keras.layers.Layer):
- """TF layer implementation of torch.nn.PixelShuffle"""
-
- def __init__(self, upscale_factor: int, **kwargs) -> None:
- super().__init__(**kwargs)
- if not isinstance(upscale_factor, int) or upscale_factor < 2:
- raise ValueError(f"upscale_factor must be an integer value >= 2 got {upscale_factor}")
- self.upscale_factor = upscale_factor
-
- def call(self, x: tf.Tensor) -> tf.Tensor:
- hidden_states = x
- batch_size, _, _, num_input_channels = shape_list(hidden_states)
- block_size_squared = self.upscale_factor**2
- output_depth = int(num_input_channels / block_size_squared)
- # When the number of output channels >= 2, PyTorch's PixelShuffle and
- # TF's depth_to_space differ in their output as the order of channels selected for combining
- # is a permutation of the other c.f.
- # https://stackoverflow.com/questions/68272502/tf-depth-to-space-not-same-as-torchs-pixelshuffle-when-output-channels-1
- permutation = tf.constant(
- [[i + j * block_size_squared for i in range(block_size_squared) for j in range(output_depth)]]
- )
- hidden_states = tf.gather(params=hidden_states, indices=tf.tile(permutation, [batch_size, 1]), batch_dims=-1)
- hidden_states = tf.nn.depth_to_space(hidden_states, block_size=self.upscale_factor, data_format="NHWC")
- return hidden_states
-
-
-class TFDeitDecoder(keras.layers.Layer):
- def __init__(self, config: DeiTConfig, **kwargs) -> None:
- super().__init__(**kwargs)
- self.conv2d = keras.layers.Conv2D(
- filters=config.encoder_stride**2 * config.num_channels, kernel_size=1, name="0"
- )
- self.pixel_shuffle = TFDeitPixelShuffle(config.encoder_stride, name="1")
- self.config = config
-
- def call(self, inputs: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = inputs
- hidden_states = self.conv2d(hidden_states)
- hidden_states = self.pixel_shuffle(hidden_states)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "conv2d", None) is not None:
- with tf.name_scope(self.conv2d.name):
- self.conv2d.build([None, None, None, self.config.hidden_size])
- if getattr(self, "pixel_shuffle", None) is not None:
- with tf.name_scope(self.pixel_shuffle.name):
- self.pixel_shuffle.build(None)
-
-
-@add_start_docstrings(
- "DeiT Model with a decoder on top for masked image modeling, as proposed in"
- " [SimMIM](https://arxiv.org/abs/2111.09886).",
- DEIT_START_DOCSTRING,
-)
-class TFDeiTForMaskedImageModeling(TFDeiTPreTrainedModel):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__(config)
-
- self.deit = TFDeiTMainLayer(config, add_pooling_layer=False, use_mask_token=True, name="deit")
- self.decoder = TFDeitDecoder(config, name="decoder")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEIT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFMaskedImageModelingOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- bool_masked_pos: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[tuple, TFMaskedImageModelingOutput]:
- r"""
- bool_masked_pos (`tf.Tensor` of type bool and shape `(batch_size, num_patches)`):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
-
- Returns:
-
- Examples:
- ```python
- >>> from transformers import AutoImageProcessor, TFDeiTForMaskedImageModeling
- >>> import tensorflow as tf
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224")
- >>> model = TFDeiTForMaskedImageModeling.from_pretrained("facebook/deit-base-distilled-patch16-224")
-
- >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2
- >>> pixel_values = image_processor(images=image, return_tensors="tf").pixel_values
- >>> # create random boolean mask of shape (batch_size, num_patches)
- >>> bool_masked_pos = tf.cast(tf.random.uniform((1, num_patches), minval=0, maxval=2, dtype=tf.int32), tf.bool)
-
- >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
- >>> loss, reconstructed_pixel_values = outputs.loss, outputs.reconstruction
- >>> list(reconstructed_pixel_values.shape)
- [1, 3, 224, 224]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deit(
- pixel_values,
- bool_masked_pos=bool_masked_pos,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = outputs[0]
-
- # Reshape to (batch_size, num_channels, height, width)
- sequence_output = sequence_output[:, 1:-1]
- batch_size, sequence_length, num_channels = shape_list(sequence_output)
- height = width = int(sequence_length**0.5)
- sequence_output = tf.reshape(sequence_output, (batch_size, height, width, num_channels))
-
- # Reconstruct pixel values
- reconstructed_pixel_values = self.decoder(sequence_output, training=training)
- # TF 2.0 image layers can't use NCHW format when running on CPU, so intermediate layers use NHWC,
- # including the decoder. We transpose to compute the loss against the pixel values
- # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
- reconstructed_pixel_values = tf.transpose(reconstructed_pixel_values, (0, 3, 1, 2))
-
- masked_im_loss = None
- if bool_masked_pos is not None:
- size = self.config.image_size // self.config.patch_size
- bool_masked_pos = tf.reshape(bool_masked_pos, (-1, size, size))
- mask = tf.repeat(bool_masked_pos, self.config.patch_size, 1)
- mask = tf.repeat(mask, self.config.patch_size, 2)
- mask = tf.expand_dims(mask, 1)
- mask = tf.cast(mask, tf.float32)
-
- reconstruction_loss = keras.losses.mean_absolute_error(
- # Swap axes as metric calculation reduces over the final dimension
- tf.transpose(pixel_values, (1, 2, 3, 0)),
- tf.transpose(reconstructed_pixel_values, (1, 2, 3, 0)),
- )
- reconstruction_loss = tf.expand_dims(reconstruction_loss, 0)
- total_loss = tf.reduce_sum(reconstruction_loss * mask)
- num_masked_pixels = (tf.reduce_sum(mask) + 1e-5) * self.config.num_channels
- masked_im_loss = total_loss / num_masked_pixels
- masked_im_loss = tf.reshape(masked_im_loss, (1,))
-
- if not return_dict:
- output = (reconstructed_pixel_values,) + outputs[1:]
- return ((masked_im_loss,) + output) if masked_im_loss is not None else output
-
- return TFMaskedImageModelingOutput(
- loss=masked_im_loss,
- reconstruction=reconstructed_pixel_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deit", None) is not None:
- with tf.name_scope(self.deit.name):
- self.deit.build(None)
- if getattr(self, "decoder", None) is not None:
- with tf.name_scope(self.decoder.name):
- self.decoder.build(None)
-
-
-@add_start_docstrings(
- """
- DeiT Model transformer with an image classification head on top (a linear layer on top of the final hidden state of
- the [CLS] token) e.g. for ImageNet.
- """,
- DEIT_START_DOCSTRING,
-)
-class TFDeiTForImageClassification(TFDeiTPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: DeiTConfig):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.deit = TFDeiTMainLayer(config, add_pooling_layer=False, name="deit")
-
- # Classifier head
- self.classifier = (
- keras.layers.Dense(config.num_labels, name="classifier")
- if config.num_labels > 0
- else keras.layers.Activation("linear", name="classifier")
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEIT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFImageClassifierOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- labels: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[tf.Tensor, TFImageClassifierOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, TFDeiTForImageClassification
- >>> import tensorflow as tf
- >>> from PIL import Image
- >>> import requests
-
- >>> keras.utils.set_random_seed(3) # doctest: +IGNORE_RESULT
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> # note: we are loading a TFDeiTForImageClassificationWithTeacher from the hub here,
- >>> # so the head will be randomly initialized, hence the predictions will be random
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224")
- >>> model = TFDeiTForImageClassification.from_pretrained("facebook/deit-base-distilled-patch16-224")
-
- >>> inputs = image_processor(images=image, return_tensors="tf")
- >>> outputs = model(**inputs)
- >>> logits = outputs.logits
- >>> # model predicts one of the 1000 ImageNet classes
- >>> predicted_class_idx = tf.math.argmax(logits, axis=-1)[0]
- >>> print("Predicted class:", model.config.id2label[int(predicted_class_idx)])
- Predicted class: little blue heron, Egretta caerulea
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deit(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = outputs[0]
-
- logits = self.classifier(sequence_output[:, 0, :])
- # we don't use the distillation token
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFImageClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deit", None) is not None:
- with tf.name_scope(self.deit.name):
- self.deit.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- DeiT Model transformer with image classification heads on top (a linear layer on top of the final hidden state of
- the [CLS] token and a linear layer on top of the final hidden state of the distillation token) e.g. for ImageNet.
-
- .. warning::
-
- This model supports inference-only. Fine-tuning with distillation (i.e. with a teacher) is not yet
- supported.
- """,
- DEIT_START_DOCSTRING,
-)
-class TFDeiTForImageClassificationWithTeacher(TFDeiTPreTrainedModel):
- def __init__(self, config: DeiTConfig) -> None:
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.deit = TFDeiTMainLayer(config, add_pooling_layer=False, name="deit")
-
- # Classifier heads
- self.cls_classifier = (
- keras.layers.Dense(config.num_labels, name="cls_classifier")
- if config.num_labels > 0
- else keras.layers.Activation("linear", name="cls_classifier")
- )
- self.distillation_classifier = (
- keras.layers.Dense(config.num_labels, name="distillation_classifier")
- if config.num_labels > 0
- else keras.layers.Activation("linear", name="distillation_classifier")
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DEIT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=TFDeiTForImageClassificationWithTeacherOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def call(
- self,
- pixel_values: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[tuple, TFDeiTForImageClassificationWithTeacherOutput]:
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.deit(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = outputs[0]
-
- cls_logits = self.cls_classifier(sequence_output[:, 0, :])
- distillation_logits = self.distillation_classifier(sequence_output[:, 1, :])
-
- # during inference, return the average of both classifier predictions
- logits = (cls_logits + distillation_logits) / 2
-
- if not return_dict:
- output = (logits, cls_logits, distillation_logits) + outputs[1:]
- return output
-
- return TFDeiTForImageClassificationWithTeacherOutput(
- logits=logits,
- cls_logits=cls_logits,
- distillation_logits=distillation_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "deit", None) is not None:
- with tf.name_scope(self.deit.name):
- self.deit.build(None)
- if getattr(self, "cls_classifier", None) is not None:
- with tf.name_scope(self.cls_classifier.name):
- self.cls_classifier.build([None, None, self.config.hidden_size])
- if getattr(self, "distillation_classifier", None) is not None:
- with tf.name_scope(self.distillation_classifier.name):
- self.distillation_classifier.build([None, None, self.config.hidden_size])
diff --git a/transformers/models/deprecated/__init__.py b/transformers/models/deprecated/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/transformers/models/deprecated/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index b278b142ac2ee3b6b0337a68840024d4f975575e..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/__pycache__/_archive_maps.cpython-310.pyc b/transformers/models/deprecated/__pycache__/_archive_maps.cpython-310.pyc
deleted file mode 100644
index 3e9f2d1acaa62a1a72782dd175dabf077fb6736c..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/__pycache__/_archive_maps.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/_archive_maps.py b/transformers/models/deprecated/_archive_maps.py
deleted file mode 100644
index 256813e0883f4540a391f9c59d2bd121ed8a5356..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/_archive_maps.py
+++ /dev/null
@@ -1,2774 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from collections import OrderedDict
-
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class DeprecatedDict(dict):
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- def __getitem__(self, item):
- logger.warning(
- "Archive maps are deprecated and will be removed in version v4.40.0 as they are no longer relevant. "
- "If looking to get all checkpoints for a given architecture, we recommend using `huggingface_hub` "
- "with the list_models method."
- )
- return self[item]
-
-
-class DeprecatedList(list):
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- def __getitem__(self, item):
- logger.warning_once(
- "Archive maps are deprecated and will be removed in version v4.40.0 as they are no longer relevant. "
- "If looking to get all checkpoints for a given architecture, we recommend using `huggingface_hub` "
- "with the `list_models` method."
- )
- return super().__getitem__(item)
-
-
-ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "albert/albert-base-v1": "https://huggingface.co/albert/albert-base-v1/resolve/main/config.json",
- "albert/albert-large-v1": "https://huggingface.co/albert/albert-large-v1/resolve/main/config.json",
- "albert/albert-xlarge-v1": "https://huggingface.co/albert/albert-xlarge-v1/resolve/main/config.json",
- "albert/albert-xxlarge-v1": "https://huggingface.co/albert/albert-xxlarge-v1/resolve/main/config.json",
- "albert/albert-base-v2": "https://huggingface.co/albert/albert-base-v2/resolve/main/config.json",
- "albert/albert-large-v2": "https://huggingface.co/albert/albert-large-v2/resolve/main/config.json",
- "albert/albert-xlarge-v2": "https://huggingface.co/albert/albert-xlarge-v2/resolve/main/config.json",
- "albert/albert-xxlarge-v2": "https://huggingface.co/albert/albert-xxlarge-v2/resolve/main/config.json",
- }
-)
-
-ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "albert/albert-base-v1",
- "albert/albert-large-v1",
- "albert/albert-xlarge-v1",
- "albert/albert-xxlarge-v1",
- "albert/albert-base-v2",
- "albert/albert-large-v2",
- "albert/albert-xlarge-v2",
- "albert/albert-xxlarge-v2",
- ]
-)
-
-TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "albert/albert-base-v1",
- "albert/albert-large-v1",
- "albert/albert-xlarge-v1",
- "albert/albert-xxlarge-v1",
- "albert/albert-base-v2",
- "albert/albert-large-v2",
- "albert/albert-xlarge-v2",
- "albert/albert-xxlarge-v2",
- ]
-)
-
-ALIGN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"kakaobrain/align-base": "https://huggingface.co/kakaobrain/align-base/resolve/main/config.json"}
-)
-
-ALIGN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["kakaobrain/align-base"])
-
-ALTCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"BAAI/AltCLIP": "https://huggingface.co/BAAI/AltCLIP/resolve/main/config.json"}
-)
-
-ALTCLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["BAAI/AltCLIP"])
-
-AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "MIT/ast-finetuned-audioset-10-10-0.4593": "https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593/resolve/main/config.json"
- }
-)
-
-AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["MIT/ast-finetuned-audioset-10-10-0.4593"]
-)
-
-AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "huggingface/autoformer-tourism-monthly": "https://huggingface.co/huggingface/autoformer-tourism-monthly/resolve/main/config.json"
- }
-)
-
-AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["huggingface/autoformer-tourism-monthly"])
-
-BARK_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["suno/bark-small", "suno/bark"])
-
-BART_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/bart-large"])
-
-BEIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/beit-base-patch16-224-pt22k": "https://huggingface.co/microsoft/beit-base-patch16-224-pt22k/resolve/main/config.json"
- }
-)
-
-BEIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/beit-base-patch16-224"])
-
-BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google-bert/bert-base-uncased": "https://huggingface.co/google-bert/bert-base-uncased/resolve/main/config.json",
- "google-bert/bert-large-uncased": "https://huggingface.co/google-bert/bert-large-uncased/resolve/main/config.json",
- "google-bert/bert-base-cased": "https://huggingface.co/google-bert/bert-base-cased/resolve/main/config.json",
- "google-bert/bert-large-cased": "https://huggingface.co/google-bert/bert-large-cased/resolve/main/config.json",
- "google-bert/bert-base-multilingual-uncased": "https://huggingface.co/google-bert/bert-base-multilingual-uncased/resolve/main/config.json",
- "google-bert/bert-base-multilingual-cased": "https://huggingface.co/google-bert/bert-base-multilingual-cased/resolve/main/config.json",
- "google-bert/bert-base-chinese": "https://huggingface.co/google-bert/bert-base-chinese/resolve/main/config.json",
- "google-bert/bert-base-german-cased": "https://huggingface.co/google-bert/bert-base-german-cased/resolve/main/config.json",
- "google-bert/bert-large-uncased-whole-word-masking": "https://huggingface.co/google-bert/bert-large-uncased-whole-word-masking/resolve/main/config.json",
- "google-bert/bert-large-cased-whole-word-masking": "https://huggingface.co/google-bert/bert-large-cased-whole-word-masking/resolve/main/config.json",
- "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad": "https://huggingface.co/google-bert/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/config.json",
- "google-bert/bert-large-cased-whole-word-masking-finetuned-squad": "https://huggingface.co/google-bert/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/config.json",
- "google-bert/bert-base-cased-finetuned-mrpc": "https://huggingface.co/google-bert/bert-base-cased-finetuned-mrpc/resolve/main/config.json",
- "google-bert/bert-base-german-dbmdz-cased": "https://huggingface.co/google-bert/bert-base-german-dbmdz-cased/resolve/main/config.json",
- "google-bert/bert-base-german-dbmdz-uncased": "https://huggingface.co/google-bert/bert-base-german-dbmdz-uncased/resolve/main/config.json",
- "cl-tohoku/bert-base-japanese": "https://huggingface.co/cl-tohoku/bert-base-japanese/resolve/main/config.json",
- "cl-tohoku/bert-base-japanese-whole-word-masking": "https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking/resolve/main/config.json",
- "cl-tohoku/bert-base-japanese-char": "https://huggingface.co/cl-tohoku/bert-base-japanese-char/resolve/main/config.json",
- "cl-tohoku/bert-base-japanese-char-whole-word-masking": "https://huggingface.co/cl-tohoku/bert-base-japanese-char-whole-word-masking/resolve/main/config.json",
- "TurkuNLP/bert-base-finnish-cased-v1": "https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/config.json",
- "TurkuNLP/bert-base-finnish-uncased-v1": "https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/config.json",
- "wietsedv/bert-base-dutch-cased": "https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/config.json",
- }
-)
-
-BERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google-bert/bert-base-uncased",
- "google-bert/bert-large-uncased",
- "google-bert/bert-base-cased",
- "google-bert/bert-large-cased",
- "google-bert/bert-base-multilingual-uncased",
- "google-bert/bert-base-multilingual-cased",
- "google-bert/bert-base-chinese",
- "google-bert/bert-base-german-cased",
- "google-bert/bert-large-uncased-whole-word-masking",
- "google-bert/bert-large-cased-whole-word-masking",
- "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad",
- "google-bert/bert-large-cased-whole-word-masking-finetuned-squad",
- "google-bert/bert-base-cased-finetuned-mrpc",
- "google-bert/bert-base-german-dbmdz-cased",
- "google-bert/bert-base-german-dbmdz-uncased",
- "cl-tohoku/bert-base-japanese",
- "cl-tohoku/bert-base-japanese-whole-word-masking",
- "cl-tohoku/bert-base-japanese-char",
- "cl-tohoku/bert-base-japanese-char-whole-word-masking",
- "TurkuNLP/bert-base-finnish-cased-v1",
- "TurkuNLP/bert-base-finnish-uncased-v1",
- "wietsedv/bert-base-dutch-cased",
- ]
-)
-
-TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google-bert/bert-base-uncased",
- "google-bert/bert-large-uncased",
- "google-bert/bert-base-cased",
- "google-bert/bert-large-cased",
- "google-bert/bert-base-multilingual-uncased",
- "google-bert/bert-base-multilingual-cased",
- "google-bert/bert-base-chinese",
- "google-bert/bert-base-german-cased",
- "google-bert/bert-large-uncased-whole-word-masking",
- "google-bert/bert-large-cased-whole-word-masking",
- "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad",
- "google-bert/bert-large-cased-whole-word-masking-finetuned-squad",
- "google-bert/bert-base-cased-finetuned-mrpc",
- "cl-tohoku/bert-base-japanese",
- "cl-tohoku/bert-base-japanese-whole-word-masking",
- "cl-tohoku/bert-base-japanese-char",
- "cl-tohoku/bert-base-japanese-char-whole-word-masking",
- "TurkuNLP/bert-base-finnish-cased-v1",
- "TurkuNLP/bert-base-finnish-uncased-v1",
- "wietsedv/bert-base-dutch-cased",
- ]
-)
-
-BIG_BIRD_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/bigbird-roberta-base": "https://huggingface.co/google/bigbird-roberta-base/resolve/main/config.json",
- "google/bigbird-roberta-large": "https://huggingface.co/google/bigbird-roberta-large/resolve/main/config.json",
- "google/bigbird-base-trivia-itc": "https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/config.json",
- }
-)
-
-BIG_BIRD_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["google/bigbird-roberta-base", "google/bigbird-roberta-large", "google/bigbird-base-trivia-itc"]
-)
-
-BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/bigbird-pegasus-large-arxiv": "https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json",
- "google/bigbird-pegasus-large-pubmed": "https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json",
- "google/bigbird-pegasus-large-bigpatent": "https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json",
- }
-)
-
-BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/bigbird-pegasus-large-arxiv",
- "google/bigbird-pegasus-large-pubmed",
- "google/bigbird-pegasus-large-bigpatent",
- ]
-)
-
-BIOGPT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/biogpt": "https://huggingface.co/microsoft/biogpt/resolve/main/config.json"}
-)
-
-BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/biogpt", "microsoft/BioGPT-Large"])
-
-BIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/bit-50": "https://huggingface.co/google/bit-50/resolve/main/config.json"}
-)
-
-BIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/bit-50"])
-
-BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/blenderbot-3B": "https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json"}
-)
-
-BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/blenderbot-3B"])
-
-BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP = {
- "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json",
- # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
-}
-
-BLENDERBOT_SMALL_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/blenderbot_small-90M"])
-
-BLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "Salesforce/blip-vqa-base": "https://huggingface.co/Salesforce/blip-vqa-base/resolve/main/config.json",
- "Salesforce/blip-vqa-capfit-large": "https://huggingface.co/Salesforce/blip-vqa-base-capfit/resolve/main/config.json",
- "Salesforce/blip-image-captioning-base": "https://huggingface.co/Salesforce/blip-image-captioning-base/resolve/main/config.json",
- "Salesforce/blip-image-captioning-large": "https://huggingface.co/Salesforce/blip-image-captioning-large/resolve/main/config.json",
- "Salesforce/blip-itm-base-coco": "https://huggingface.co/Salesforce/blip-itm-base-coco/resolve/main/config.json",
- "Salesforce/blip-itm-large-coco": "https://huggingface.co/Salesforce/blip-itm-large-coco/resolve/main/config.json",
- "Salesforce/blip-itm-base-flikr": "https://huggingface.co/Salesforce/blip-itm-base-flikr/resolve/main/config.json",
- "Salesforce/blip-itm-large-flikr": "https://huggingface.co/Salesforce/blip-itm-large-flikr/resolve/main/config.json",
- }
-)
-
-BLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "Salesforce/blip-vqa-base",
- "Salesforce/blip-vqa-capfilt-large",
- "Salesforce/blip-image-captioning-base",
- "Salesforce/blip-image-captioning-large",
- "Salesforce/blip-itm-base-coco",
- "Salesforce/blip-itm-large-coco",
- "Salesforce/blip-itm-base-flickr",
- "Salesforce/blip-itm-large-flickr",
- ]
-)
-
-TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "Salesforce/blip-vqa-base",
- "Salesforce/blip-vqa-capfilt-large",
- "Salesforce/blip-image-captioning-base",
- "Salesforce/blip-image-captioning-large",
- "Salesforce/blip-itm-base-coco",
- "Salesforce/blip-itm-large-coco",
- "Salesforce/blip-itm-base-flickr",
- "Salesforce/blip-itm-large-flickr",
- ]
-)
-
-BLIP_2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"salesforce/blip2-opt-2.7b": "https://huggingface.co/salesforce/blip2-opt-2.7b/resolve/main/config.json"}
-)
-
-BLIP_2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Salesforce/blip2-opt-2.7b"])
-
-BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "bigscience/bloom": "https://huggingface.co/bigscience/bloom/resolve/main/config.json",
- "bigscience/bloom-560m": "https://huggingface.co/bigscience/bloom-560m/blob/main/config.json",
- "bigscience/bloom-1b1": "https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json",
- "bigscience/bloom-1b7": "https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json",
- "bigscience/bloom-3b": "https://huggingface.co/bigscience/bloom-3b/blob/main/config.json",
- "bigscience/bloom-7b1": "https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json",
- }
-)
-
-BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "bigscience/bigscience-small-testing",
- "bigscience/bloom-560m",
- "bigscience/bloom-1b1",
- "bigscience/bloom-1b7",
- "bigscience/bloom-3b",
- "bigscience/bloom-7b1",
- "bigscience/bloom",
- ]
-)
-
-BRIDGETOWER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "BridgeTower/bridgetower-base": "https://huggingface.co/BridgeTower/bridgetower-base/blob/main/config.json",
- "BridgeTower/bridgetower-base-itm-mlm": "https://huggingface.co/BridgeTower/bridgetower-base-itm-mlm/blob/main/config.json",
- }
-)
-
-BRIDGETOWER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["BridgeTower/bridgetower-base", "BridgeTower/bridgetower-base-itm-mlm"]
-)
-
-BROS_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "jinho8345/bros-base-uncased": "https://huggingface.co/jinho8345/bros-base-uncased/blob/main/config.json",
- "jinho8345/bros-large-uncased": "https://huggingface.co/jinho8345/bros-large-uncased/blob/main/config.json",
- }
-)
-
-BROS_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["jinho8345/bros-base-uncased", "jinho8345/bros-large-uncased"])
-
-CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "almanach/camembert-base": "https://huggingface.co/almanach/camembert-base/resolve/main/config.json",
- "umberto-commoncrawl-cased-v1": "https://huggingface.co/Musixmatch/umberto-commoncrawl-cased-v1/resolve/main/config.json",
- "umberto-wikipedia-uncased-v1": "https://huggingface.co/Musixmatch/umberto-wikipedia-uncased-v1/resolve/main/config.json",
- }
-)
-
-CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["almanach/camembert-base", "Musixmatch/umberto-commoncrawl-cased-v1", "Musixmatch/umberto-wikipedia-uncased-v1"]
-)
-
-TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList([])
-
-CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/canine-s": "https://huggingface.co/google/canine-s/resolve/main/config.json"}
-)
-
-CANINE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/canine-s", "google/canine-r"])
-
-CHINESE_CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "OFA-Sys/chinese-clip-vit-base-patch16": "https://huggingface.co/OFA-Sys/chinese-clip-vit-base-patch16/resolve/main/config.json"
- }
-)
-
-CHINESE_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["OFA-Sys/chinese-clip-vit-base-patch16"])
-
-CLAP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["laion/clap-htsat-fused", "laion/clap-htsat-unfused"])
-
-CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"openai/clip-vit-base-patch32": "https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/config.json"}
-)
-
-CLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["openai/clip-vit-base-patch32"])
-
-TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["openai/clip-vit-base-patch32"])
-
-CLIPSEG_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"CIDAS/clipseg-rd64": "https://huggingface.co/CIDAS/clipseg-rd64/resolve/main/config.json"}
-)
-
-CLIPSEG_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["CIDAS/clipseg-rd64-refined"])
-
-CLVP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"susnato/clvp_dev": "https://huggingface.co/susnato/clvp_dev/resolve/main/config.json"}
-)
-
-CLVP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["susnato/clvp_dev"])
-
-CODEGEN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "Salesforce/codegen-350M-nl": "https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json",
- "Salesforce/codegen-350M-multi": "https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json",
- "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json",
- "Salesforce/codegen-2B-nl": "https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json",
- "Salesforce/codegen-2B-multi": "https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json",
- "Salesforce/codegen-2B-mono": "https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json",
- "Salesforce/codegen-6B-nl": "https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json",
- "Salesforce/codegen-6B-multi": "https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json",
- "Salesforce/codegen-6B-mono": "https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json",
- "Salesforce/codegen-16B-nl": "https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json",
- "Salesforce/codegen-16B-multi": "https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json",
- "Salesforce/codegen-16B-mono": "https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json",
- }
-)
-
-CODEGEN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "Salesforce/codegen-350M-nl",
- "Salesforce/codegen-350M-multi",
- "Salesforce/codegen-350M-mono",
- "Salesforce/codegen-2B-nl",
- "Salesforce/codegen-2B-multi",
- "Salesforce/codegen-2B-mono",
- "Salesforce/codegen-6B-nl",
- "Salesforce/codegen-6B-multi",
- "Salesforce/codegen-6B-mono",
- "Salesforce/codegen-16B-nl",
- "Salesforce/codegen-16B-multi",
- "Salesforce/codegen-16B-mono",
- ]
-)
-
-CONDITIONAL_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/conditional-detr-resnet-50": "https://huggingface.co/microsoft/conditional-detr-resnet-50/resolve/main/config.json"
- }
-)
-
-CONDITIONAL_DETR_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/conditional-detr-resnet-50"])
-
-CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "YituTech/conv-bert-base": "https://huggingface.co/YituTech/conv-bert-base/resolve/main/config.json",
- "YituTech/conv-bert-medium-small": "https://huggingface.co/YituTech/conv-bert-medium-small/resolve/main/config.json",
- "YituTech/conv-bert-small": "https://huggingface.co/YituTech/conv-bert-small/resolve/main/config.json",
- }
-)
-
-CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["YituTech/conv-bert-base", "YituTech/conv-bert-medium-small", "YituTech/conv-bert-small"]
-)
-
-TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["YituTech/conv-bert-base", "YituTech/conv-bert-medium-small", "YituTech/conv-bert-small"]
-)
-
-CONVNEXT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/convnext-tiny-224": "https://huggingface.co/facebook/convnext-tiny-224/resolve/main/config.json"}
-)
-
-CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/convnext-tiny-224"])
-
-CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/convnextv2-tiny-1k-224": "https://huggingface.co/facebook/convnextv2-tiny-1k-224/resolve/main/config.json"
- }
-)
-
-CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/convnextv2-tiny-1k-224"])
-
-CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"openbmb/cpm-ant-10b": "https://huggingface.co/openbmb/cpm-ant-10b/blob/main/config.json"}
-)
-
-CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["openbmb/cpm-ant-10b"])
-
-CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"Salesforce/ctrl": "https://huggingface.co/Salesforce/ctrl/resolve/main/config.json"}
-)
-
-CTRL_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Salesforce/ctrl"])
-
-TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Salesforce/ctrl"])
-
-CVT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/cvt-13": "https://huggingface.co/microsoft/cvt-13/resolve/main/config.json"}
-)
-
-CVT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "microsoft/cvt-13",
- "microsoft/cvt-13-384",
- "microsoft/cvt-13-384-22k",
- "microsoft/cvt-21",
- "microsoft/cvt-21-384",
- "microsoft/cvt-21-384-22k",
- ]
-)
-
-TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "microsoft/cvt-13",
- "microsoft/cvt-13-384",
- "microsoft/cvt-13-384-22k",
- "microsoft/cvt-21",
- "microsoft/cvt-21-384",
- "microsoft/cvt-21-384-22k",
- ]
-)
-
-DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/data2vec-text-base": "https://huggingface.co/data2vec/resolve/main/config.json"}
-)
-
-DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/data2vec-vision-base-ft": "https://huggingface.co/facebook/data2vec-vision-base-ft/resolve/main/config.json"
- }
-)
-
-DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "facebook/data2vec-audio-base",
- "facebook/data2vec-audio-base-10m",
- "facebook/data2vec-audio-base-100h",
- "facebook/data2vec-audio-base-960h",
- ]
-)
-
-DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/data2vec-text-base"])
-
-DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/data2vec-vision-base-ft1k"])
-
-DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/deberta-base": "https://huggingface.co/microsoft/deberta-base/resolve/main/config.json",
- "microsoft/deberta-large": "https://huggingface.co/microsoft/deberta-large/resolve/main/config.json",
- "microsoft/deberta-xlarge": "https://huggingface.co/microsoft/deberta-xlarge/resolve/main/config.json",
- "microsoft/deberta-base-mnli": "https://huggingface.co/microsoft/deberta-base-mnli/resolve/main/config.json",
- "microsoft/deberta-large-mnli": "https://huggingface.co/microsoft/deberta-large-mnli/resolve/main/config.json",
- "microsoft/deberta-xlarge-mnli": "https://huggingface.co/microsoft/deberta-xlarge-mnli/resolve/main/config.json",
- }
-)
-
-DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "microsoft/deberta-base",
- "microsoft/deberta-large",
- "microsoft/deberta-xlarge",
- "microsoft/deberta-base-mnli",
- "microsoft/deberta-large-mnli",
- "microsoft/deberta-xlarge-mnli",
- ]
-)
-
-TF_DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["kamalkraj/deberta-base"])
-
-DEBERTA_V2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/deberta-v2-xlarge": "https://huggingface.co/microsoft/deberta-v2-xlarge/resolve/main/config.json",
- "microsoft/deberta-v2-xxlarge": "https://huggingface.co/microsoft/deberta-v2-xxlarge/resolve/main/config.json",
- "microsoft/deberta-v2-xlarge-mnli": "https://huggingface.co/microsoft/deberta-v2-xlarge-mnli/resolve/main/config.json",
- "microsoft/deberta-v2-xxlarge-mnli": "https://huggingface.co/microsoft/deberta-v2-xxlarge-mnli/resolve/main/config.json",
- }
-)
-
-DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "microsoft/deberta-v2-xlarge",
- "microsoft/deberta-v2-xxlarge",
- "microsoft/deberta-v2-xlarge-mnli",
- "microsoft/deberta-v2-xxlarge-mnli",
- ]
-)
-
-TF_DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["kamalkraj/deberta-v2-xlarge"])
-
-DECISION_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "edbeeching/decision-transformer-gym-hopper-medium": "https://huggingface.co/edbeeching/decision-transformer-gym-hopper-medium/resolve/main/config.json"
- }
-)
-
-DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["edbeeching/decision-transformer-gym-hopper-medium"]
-)
-
-DEFORMABLE_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"SenseTime/deformable-detr": "https://huggingface.co/sensetime/deformable-detr/resolve/main/config.json"}
-)
-
-DEFORMABLE_DETR_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["sensetime/deformable-detr"])
-
-DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/deit-base-distilled-patch16-224": "https://huggingface.co/facebook/deit-base-patch16-224/resolve/main/config.json"
- }
-)
-
-DEIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/deit-base-distilled-patch16-224"])
-
-TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/deit-base-distilled-patch16-224"])
-
-MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"speechbrain/m-ctc-t-large": "https://huggingface.co/speechbrain/m-ctc-t-large/resolve/main/config.json"}
-)
-
-MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["speechbrain/m-ctc-t-large"])
-
-OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"s-JoL/Open-Llama-V1": "https://huggingface.co/s-JoL/Open-Llama-V1/blob/main/config.json"}
-)
-
-RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "yjernite/retribert-base-uncased": "https://huggingface.co/yjernite/retribert-base-uncased/resolve/main/config.json"
- }
-)
-
-RETRIBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["yjernite/retribert-base-uncased"])
-
-TRAJECTORY_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "CarlCochet/trajectory-transformer-halfcheetah-medium-v2": "https://huggingface.co/CarlCochet/trajectory-transformer-halfcheetah-medium-v2/resolve/main/config.json"
- }
-)
-
-TRAJECTORY_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["CarlCochet/trajectory-transformer-halfcheetah-medium-v2"]
-)
-
-TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"transfo-xl/transfo-xl-wt103": "https://huggingface.co/transfo-xl/transfo-xl-wt103/resolve/main/config.json"}
-)
-
-TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["transfo-xl/transfo-xl-wt103"])
-
-TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["transfo-xl/transfo-xl-wt103"])
-
-VAN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "Visual-Attention-Network/van-base": "https://huggingface.co/Visual-Attention-Network/van-base/blob/main/config.json"
- }
-)
-
-VAN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Visual-Attention-Network/van-base"])
-
-DEPTH_ANYTHING_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "LiheYoung/depth-anything-small-hf": "https://huggingface.co/LiheYoung/depth-anything-small-hf/resolve/main/config.json"
- }
-)
-
-DEPTH_ANYTHING_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["LiheYoung/depth-anything-small-hf"])
-
-DETA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"ut/deta": "https://huggingface.co/ut/deta/resolve/main/config.json"}
-)
-
-DETA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["jozhang97/deta-swin-large-o365"])
-
-DETR_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/detr-resnet-50": "https://huggingface.co/facebook/detr-resnet-50/resolve/main/config.json"}
-)
-
-DETR_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/detr-resnet-50"])
-
-DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"shi-labs/dinat-mini-in1k-224": "https://huggingface.co/shi-labs/dinat-mini-in1k-224/resolve/main/config.json"}
-)
-
-DINAT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["shi-labs/dinat-mini-in1k-224"])
-
-DINOV2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/dinov2-base": "https://huggingface.co/facebook/dinov2-base/resolve/main/config.json"}
-)
-
-DINOV2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/dinov2-base"])
-
-DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "distilbert-base-uncased": "https://huggingface.co/distilbert-base-uncased/resolve/main/config.json",
- "distilbert-base-uncased-distilled-squad": "https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json",
- "distilbert-base-cased": "https://huggingface.co/distilbert-base-cased/resolve/main/config.json",
- "distilbert-base-cased-distilled-squad": "https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json",
- "distilbert-base-german-cased": "https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json",
- "distilbert-base-multilingual-cased": "https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json",
- "distilbert-base-uncased-finetuned-sst-2-english": "https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json",
- }
-)
-
-DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "distilbert-base-uncased",
- "distilbert-base-uncased-distilled-squad",
- "distilbert-base-cased",
- "distilbert-base-cased-distilled-squad",
- "distilbert-base-german-cased",
- "distilbert-base-multilingual-cased",
- "distilbert-base-uncased-finetuned-sst-2-english",
- ]
-)
-
-TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "distilbert-base-uncased",
- "distilbert-base-uncased-distilled-squad",
- "distilbert-base-cased",
- "distilbert-base-cased-distilled-squad",
- "distilbert-base-multilingual-cased",
- "distilbert-base-uncased-finetuned-sst-2-english",
- ]
-)
-
-DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"naver-clova-ix/donut-base": "https://huggingface.co/naver-clova-ix/donut-base/resolve/main/config.json"}
-)
-
-DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["naver-clova-ix/donut-base"])
-
-DPR_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/dpr-ctx_encoder-single-nq-base": "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json",
- "facebook/dpr-question_encoder-single-nq-base": "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json",
- "facebook/dpr-reader-single-nq-base": "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json",
- "facebook/dpr-ctx_encoder-multiset-base": "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json",
- "facebook/dpr-question_encoder-multiset-base": "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json",
- "facebook/dpr-reader-multiset-base": "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json",
- }
-)
-
-DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["facebook/dpr-ctx_encoder-single-nq-base", "facebook/dpr-ctx_encoder-multiset-base"]
-)
-
-DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["facebook/dpr-question_encoder-single-nq-base", "facebook/dpr-question_encoder-multiset-base"]
-)
-
-DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["facebook/dpr-reader-single-nq-base", "facebook/dpr-reader-multiset-base"]
-)
-
-TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["facebook/dpr-ctx_encoder-single-nq-base", "facebook/dpr-ctx_encoder-multiset-base"]
-)
-
-TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["facebook/dpr-question_encoder-single-nq-base", "facebook/dpr-question_encoder-multiset-base"]
-)
-
-TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["facebook/dpr-reader-single-nq-base", "facebook/dpr-reader-multiset-base"]
-)
-
-DPT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"Intel/dpt-large": "https://huggingface.co/Intel/dpt-large/resolve/main/config.json"}
-)
-
-DPT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Intel/dpt-large", "Intel/dpt-hybrid-midas"])
-
-EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "snap-research/efficientformer-l1-300": "https://huggingface.co/snap-research/efficientformer-l1-300/resolve/main/config.json"
- }
-)
-
-EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["snap-research/efficientformer-l1-300"])
-
-TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["snap-research/efficientformer-l1-300"])
-
-EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/efficientnet-b7": "https://huggingface.co/google/efficientnet-b7/resolve/main/config.json"}
-)
-
-EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/efficientnet-b7"])
-
-ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/electra-small-generator": "https://huggingface.co/google/electra-small-generator/resolve/main/config.json",
- "google/electra-base-generator": "https://huggingface.co/google/electra-base-generator/resolve/main/config.json",
- "google/electra-large-generator": "https://huggingface.co/google/electra-large-generator/resolve/main/config.json",
- "google/electra-small-discriminator": "https://huggingface.co/google/electra-small-discriminator/resolve/main/config.json",
- "google/electra-base-discriminator": "https://huggingface.co/google/electra-base-discriminator/resolve/main/config.json",
- "google/electra-large-discriminator": "https://huggingface.co/google/electra-large-discriminator/resolve/main/config.json",
- }
-)
-
-ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/electra-small-generator",
- "google/electra-base-generator",
- "google/electra-large-generator",
- "google/electra-small-discriminator",
- "google/electra-base-discriminator",
- "google/electra-large-discriminator",
- ]
-)
-
-TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/electra-small-generator",
- "google/electra-base-generator",
- "google/electra-large-generator",
- "google/electra-small-discriminator",
- "google/electra-base-discriminator",
- "google/electra-large-discriminator",
- ]
-)
-
-ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/encodec_24khz": "https://huggingface.co/facebook/encodec_24khz/resolve/main/config.json",
- "facebook/encodec_48khz": "https://huggingface.co/facebook/encodec_48khz/resolve/main/config.json",
- }
-)
-
-ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/encodec_24khz", "facebook/encodec_48khz"])
-
-ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "nghuyong/ernie-1.0-base-zh": "https://huggingface.co/nghuyong/ernie-1.0-base-zh/resolve/main/config.json",
- "nghuyong/ernie-2.0-base-en": "https://huggingface.co/nghuyong/ernie-2.0-base-en/resolve/main/config.json",
- "nghuyong/ernie-2.0-large-en": "https://huggingface.co/nghuyong/ernie-2.0-large-en/resolve/main/config.json",
- "nghuyong/ernie-3.0-base-zh": "https://huggingface.co/nghuyong/ernie-3.0-base-zh/resolve/main/config.json",
- "nghuyong/ernie-3.0-medium-zh": "https://huggingface.co/nghuyong/ernie-3.0-medium-zh/resolve/main/config.json",
- "nghuyong/ernie-3.0-mini-zh": "https://huggingface.co/nghuyong/ernie-3.0-mini-zh/resolve/main/config.json",
- "nghuyong/ernie-3.0-micro-zh": "https://huggingface.co/nghuyong/ernie-3.0-micro-zh/resolve/main/config.json",
- "nghuyong/ernie-3.0-nano-zh": "https://huggingface.co/nghuyong/ernie-3.0-nano-zh/resolve/main/config.json",
- "nghuyong/ernie-gram-zh": "https://huggingface.co/nghuyong/ernie-gram-zh/resolve/main/config.json",
- "nghuyong/ernie-health-zh": "https://huggingface.co/nghuyong/ernie-health-zh/resolve/main/config.json",
- }
-)
-
-ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "nghuyong/ernie-1.0-base-zh",
- "nghuyong/ernie-2.0-base-en",
- "nghuyong/ernie-2.0-large-en",
- "nghuyong/ernie-3.0-base-zh",
- "nghuyong/ernie-3.0-medium-zh",
- "nghuyong/ernie-3.0-mini-zh",
- "nghuyong/ernie-3.0-micro-zh",
- "nghuyong/ernie-3.0-nano-zh",
- "nghuyong/ernie-gram-zh",
- "nghuyong/ernie-health-zh",
- ]
-)
-
-ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "susnato/ernie-m-base_pytorch": "https://huggingface.co/susnato/ernie-m-base_pytorch/blob/main/config.json",
- "susnato/ernie-m-large_pytorch": "https://huggingface.co/susnato/ernie-m-large_pytorch/blob/main/config.json",
- }
-)
-
-ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["susnato/ernie-m-base_pytorch", "susnato/ernie-m-large_pytorch"]
-)
-
-ESM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/esm-1b": "https://huggingface.co/facebook/esm-1b/resolve/main/config.json"}
-)
-
-ESM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/esm2_t6_8M_UR50D", "facebook/esm2_t12_35M_UR50D"])
-
-FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "tiiuae/falcon-40b": "https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json",
- "tiiuae/falcon-7b": "https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json",
- }
-)
-
-FALCON_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "tiiuae/falcon-40b",
- "tiiuae/falcon-40b-instruct",
- "tiiuae/falcon-7b",
- "tiiuae/falcon-7b-instruct",
- "tiiuae/falcon-rw-7b",
- "tiiuae/falcon-rw-1b",
- ]
-)
-
-FASTSPEECH2_CONFORMER_HIFIGAN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "espnet/fastspeech2_conformer_hifigan": "https://huggingface.co/espnet/fastspeech2_conformer_hifigan/raw/main/config.json"
- }
-)
-
-FASTSPEECH2_CONFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"espnet/fastspeech2_conformer": "https://huggingface.co/espnet/fastspeech2_conformer/raw/main/config.json"}
-)
-
-FASTSPEECH2_CONFORMER_WITH_HIFIGAN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "espnet/fastspeech2_conformer_with_hifigan": "https://huggingface.co/espnet/fastspeech2_conformer_with_hifigan/raw/main/config.json"
- }
-)
-
-FASTSPEECH2_CONFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["espnet/fastspeech2_conformer"])
-
-FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "flaubert/flaubert_small_cased": "https://huggingface.co/flaubert/flaubert_small_cased/resolve/main/config.json",
- "flaubert/flaubert_base_uncased": "https://huggingface.co/flaubert/flaubert_base_uncased/resolve/main/config.json",
- "flaubert/flaubert_base_cased": "https://huggingface.co/flaubert/flaubert_base_cased/resolve/main/config.json",
- "flaubert/flaubert_large_cased": "https://huggingface.co/flaubert/flaubert_large_cased/resolve/main/config.json",
- }
-)
-
-FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "flaubert/flaubert_small_cased",
- "flaubert/flaubert_base_uncased",
- "flaubert/flaubert_base_cased",
- "flaubert/flaubert_large_cased",
- ]
-)
-
-TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList([])
-
-FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/flava-full": "https://huggingface.co/facebook/flava-full/resolve/main/config.json"}
-)
-
-FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/flava-full"])
-
-FNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/fnet-base": "https://huggingface.co/google/fnet-base/resolve/main/config.json",
- "google/fnet-large": "https://huggingface.co/google/fnet-large/resolve/main/config.json",
- }
-)
-
-FNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/fnet-base", "google/fnet-large"])
-
-FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/focalnet-tiny": "https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json"}
-)
-
-FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/focalnet-tiny"])
-
-FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict({})
-
-FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "funnel-transformer/small": "https://huggingface.co/funnel-transformer/small/resolve/main/config.json",
- "funnel-transformer/small-base": "https://huggingface.co/funnel-transformer/small-base/resolve/main/config.json",
- "funnel-transformer/medium": "https://huggingface.co/funnel-transformer/medium/resolve/main/config.json",
- "funnel-transformer/medium-base": "https://huggingface.co/funnel-transformer/medium-base/resolve/main/config.json",
- "funnel-transformer/intermediate": "https://huggingface.co/funnel-transformer/intermediate/resolve/main/config.json",
- "funnel-transformer/intermediate-base": "https://huggingface.co/funnel-transformer/intermediate-base/resolve/main/config.json",
- "funnel-transformer/large": "https://huggingface.co/funnel-transformer/large/resolve/main/config.json",
- "funnel-transformer/large-base": "https://huggingface.co/funnel-transformer/large-base/resolve/main/config.json",
- "funnel-transformer/xlarge": "https://huggingface.co/funnel-transformer/xlarge/resolve/main/config.json",
- "funnel-transformer/xlarge-base": "https://huggingface.co/funnel-transformer/xlarge-base/resolve/main/config.json",
- }
-)
-
-FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "funnel-transformer/small",
- "funnel-transformer/small-base",
- "funnel-transformer/medium",
- "funnel-transformer/medium-base",
- "funnel-transformer/intermediate",
- "funnel-transformer/intermediate-base",
- "funnel-transformer/large",
- "funnel-transformer/large-base",
- "funnel-transformer/xlarge-base",
- "funnel-transformer/xlarge",
- ]
-)
-
-TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "funnel-transformer/small",
- "funnel-transformer/small-base",
- "funnel-transformer/medium",
- "funnel-transformer/medium-base",
- "funnel-transformer/intermediate",
- "funnel-transformer/intermediate-base",
- "funnel-transformer/large",
- "funnel-transformer/large-base",
- "funnel-transformer/xlarge-base",
- "funnel-transformer/xlarge",
- ]
-)
-
-FUYU_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"adept/fuyu-8b": "https://huggingface.co/adept/fuyu-8b/resolve/main/config.json"}
-)
-
-GEMMA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict({})
-
-GIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/git-base": "https://huggingface.co/microsoft/git-base/resolve/main/config.json"}
-)
-
-GIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/git-base"])
-
-GLPN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"vinvino02/glpn-kitti": "https://huggingface.co/vinvino02/glpn-kitti/resolve/main/config.json"}
-)
-
-GLPN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["vinvino02/glpn-kitti"])
-
-GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "openai-community/gpt2": "https://huggingface.co/openai-community/gpt2/resolve/main/config.json",
- "openai-community/gpt2-medium": "https://huggingface.co/openai-community/gpt2-medium/resolve/main/config.json",
- "openai-community/gpt2-large": "https://huggingface.co/openai-community/gpt2-large/resolve/main/config.json",
- "openai-community/gpt2-xl": "https://huggingface.co/openai-community/gpt2-xl/resolve/main/config.json",
- "distilbert/distilgpt2": "https://huggingface.co/distilbert/distilgpt2/resolve/main/config.json",
- }
-)
-
-GPT2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "openai-community/gpt2",
- "openai-community/gpt2-medium",
- "openai-community/gpt2-large",
- "openai-community/gpt2-xl",
- "distilbert/distilgpt2",
- ]
-)
-
-TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "openai-community/gpt2",
- "openai-community/gpt2-medium",
- "openai-community/gpt2-large",
- "openai-community/gpt2-xl",
- "distilbert/distilgpt2",
- ]
-)
-
-GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "bigcode/gpt_bigcode-santacoder": "https://huggingface.co/bigcode/gpt_bigcode-santacoder/resolve/main/config.json"
- }
-)
-
-GPT_BIGCODE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["bigcode/gpt_bigcode-santacoder"])
-
-GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"EleutherAI/gpt-neo-1.3B": "https://huggingface.co/EleutherAI/gpt-neo-1.3B/resolve/main/config.json"}
-)
-
-GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["EleutherAI/gpt-neo-1.3B"])
-
-GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"EleutherAI/gpt-neox-20b": "https://huggingface.co/EleutherAI/gpt-neox-20b/resolve/main/config.json"}
-)
-
-GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["EleutherAI/gpt-neox-20b"])
-
-GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"abeja/gpt-neox-japanese-2.7b": "https://huggingface.co/abeja/gpt-neox-japanese-2.7b/resolve/main/config.json"}
-)
-
-GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["https://huggingface.co/abeja/gpt-neox-japanese-2.7b/resolve/main/config.json"]
-)
-
-GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"EleutherAI/gpt-j-6B": "https://huggingface.co/EleutherAI/gpt-j-6B/resolve/main/config.json"}
-)
-
-GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["EleutherAI/gpt-j-6B"])
-
-GPTSAN_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "tanreinama/GPTSAN-2.8B-spout_is_uniform": "https://huggingface.co/tanreinama/GPTSAN-2.8B-spout_is_uniform/resolve/main/config.json"
- }
-)
-
-GPTSAN_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Tanrei/GPTSAN-japanese"])
-
-GRAPHORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"graphormer-base": "https://huggingface.co/clefourrier/graphormer-base-pcqm4mv2/resolve/main/config.json"}
-)
-
-GRAPHORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["clefourrier/graphormer-base-pcqm4mv1", "clefourrier/graphormer-base-pcqm4mv2"]
-)
-
-GROUPVIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"nvidia/groupvit-gcc-yfcc": "https://huggingface.co/nvidia/groupvit-gcc-yfcc/resolve/main/config.json"}
-)
-
-GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["nvidia/groupvit-gcc-yfcc"])
-
-TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["nvidia/groupvit-gcc-yfcc"])
-
-HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/hubert-base-ls960": "https://huggingface.co/facebook/hubert-base-ls960/resolve/main/config.json"}
-)
-
-HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/hubert-base-ls960"])
-
-TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/hubert-base-ls960"])
-
-IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "kssteven/ibert-roberta-base": "https://huggingface.co/kssteven/ibert-roberta-base/resolve/main/config.json",
- "kssteven/ibert-roberta-large": "https://huggingface.co/kssteven/ibert-roberta-large/resolve/main/config.json",
- "kssteven/ibert-roberta-large-mnli": "https://huggingface.co/kssteven/ibert-roberta-large-mnli/resolve/main/config.json",
- }
-)
-
-IBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["kssteven/ibert-roberta-base", "kssteven/ibert-roberta-large", "kssteven/ibert-roberta-large-mnli"]
-)
-
-IDEFICS_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "HuggingFaceM4/idefics-9b": "https://huggingface.co/HuggingFaceM4/idefics-9b/blob/main/config.json",
- "HuggingFaceM4/idefics-80b": "https://huggingface.co/HuggingFaceM4/idefics-80b/blob/main/config.json",
- }
-)
-
-IDEFICS_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["HuggingFaceM4/idefics-9b", "HuggingFaceM4/idefics-80b"])
-
-IMAGEGPT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"openai/imagegpt-small": "", "openai/imagegpt-medium": "", "openai/imagegpt-large": ""}
-)
-
-IMAGEGPT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["openai/imagegpt-small", "openai/imagegpt-medium", "openai/imagegpt-large"]
-)
-
-INFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "huggingface/informer-tourism-monthly": "https://huggingface.co/huggingface/informer-tourism-monthly/resolve/main/config.json"
- }
-)
-
-INFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["huggingface/informer-tourism-monthly"])
-
-INSTRUCTBLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "Salesforce/instruct-blip-flan-t5": "https://huggingface.co/Salesforce/instruct-blip-flan-t5/resolve/main/config.json"
- }
-)
-
-INSTRUCTBLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Salesforce/instructblip-flan-t5-xl"])
-
-JUKEBOX_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "openai/jukebox-5b-lyrics": "https://huggingface.co/openai/jukebox-5b-lyrics/blob/main/config.json",
- "openai/jukebox-1b-lyrics": "https://huggingface.co/openai/jukebox-1b-lyrics/blob/main/config.json",
- }
-)
-
-JUKEBOX_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["openai/jukebox-1b-lyrics", "openai/jukebox-5b-lyrics"])
-
-KOSMOS2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/kosmos-2-patch14-224": "https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/config.json"
- }
-)
-
-KOSMOS2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/kosmos-2-patch14-224"])
-
-LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/layoutlm-base-uncased": "https://huggingface.co/microsoft/layoutlm-base-uncased/resolve/main/config.json",
- "microsoft/layoutlm-large-uncased": "https://huggingface.co/microsoft/layoutlm-large-uncased/resolve/main/config.json",
- }
-)
-
-LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["layoutlm-base-uncased", "layoutlm-large-uncased"])
-
-TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["microsoft/layoutlm-base-uncased", "microsoft/layoutlm-large-uncased"]
-)
-
-LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "layoutlmv2-base-uncased": "https://huggingface.co/microsoft/layoutlmv2-base-uncased/resolve/main/config.json",
- "layoutlmv2-large-uncased": "https://huggingface.co/microsoft/layoutlmv2-large-uncased/resolve/main/config.json",
- }
-)
-
-LAYOUTLMV2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["microsoft/layoutlmv2-base-uncased", "microsoft/layoutlmv2-large-uncased"]
-)
-
-LAYOUTLMV3_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/layoutlmv3-base": "https://huggingface.co/microsoft/layoutlmv3-base/resolve/main/config.json"}
-)
-
-LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/layoutlmv3-base", "microsoft/layoutlmv3-large"])
-
-TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["microsoft/layoutlmv3-base", "microsoft/layoutlmv3-large"]
-)
-
-LED_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"allenai/led-base-16384": "https://huggingface.co/allenai/led-base-16384/resolve/main/config.json"}
-)
-
-LED_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["allenai/led-base-16384"])
-
-LEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/levit-128S": "https://huggingface.co/facebook/levit-128S/resolve/main/config.json"}
-)
-
-LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/levit-128S"])
-
-LILT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "SCUT-DLVCLab/lilt-roberta-en-base": "https://huggingface.co/SCUT-DLVCLab/lilt-roberta-en-base/resolve/main/config.json"
- }
-)
-
-LILT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["SCUT-DLVCLab/lilt-roberta-en-base"])
-
-LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict({})
-
-LLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"llava-hf/llava-v1.5-7b": "https://huggingface.co/llava-hf/llava-v1.5-7b/resolve/main/config.json"}
-)
-
-LLAVA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["llava-hf/llava-1.5-7b-hf", "llava-hf/llava-1.5-13b-hf", "llava-hf/bakLlava-v1-hf"]
-)
-
-LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "allenai/longformer-base-4096": "https://huggingface.co/allenai/longformer-base-4096/resolve/main/config.json",
- "allenai/longformer-large-4096": "https://huggingface.co/allenai/longformer-large-4096/resolve/main/config.json",
- "allenai/longformer-large-4096-finetuned-triviaqa": "https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/config.json",
- "allenai/longformer-base-4096-extra.pos.embd.only": "https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/config.json",
- "allenai/longformer-large-4096-extra.pos.embd.only": "https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/config.json",
- }
-)
-
-LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "allenai/longformer-base-4096",
- "allenai/longformer-large-4096",
- "allenai/longformer-large-4096-finetuned-triviaqa",
- "allenai/longformer-base-4096-extra.pos.embd.only",
- "allenai/longformer-large-4096-extra.pos.embd.only",
- ]
-)
-
-TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "allenai/longformer-base-4096",
- "allenai/longformer-large-4096",
- "allenai/longformer-large-4096-finetuned-triviaqa",
- "allenai/longformer-base-4096-extra.pos.embd.only",
- "allenai/longformer-large-4096-extra.pos.embd.only",
- ]
-)
-
-LONGT5_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/long-t5-local-base": "https://huggingface.co/google/long-t5-local-base/blob/main/config.json",
- "google/long-t5-local-large": "https://huggingface.co/google/long-t5-local-large/blob/main/config.json",
- "google/long-t5-tglobal-base": "https://huggingface.co/google/long-t5-tglobal-base/blob/main/config.json",
- "google/long-t5-tglobal-large": "https://huggingface.co/google/long-t5-tglobal-large/blob/main/config.json",
- }
-)
-
-LONGT5_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/long-t5-local-base",
- "google/long-t5-local-large",
- "google/long-t5-tglobal-base",
- "google/long-t5-tglobal-large",
- ]
-)
-
-LUKE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "studio-ousia/luke-base": "https://huggingface.co/studio-ousia/luke-base/resolve/main/config.json",
- "studio-ousia/luke-large": "https://huggingface.co/studio-ousia/luke-large/resolve/main/config.json",
- }
-)
-
-LUKE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["studio-ousia/luke-base", "studio-ousia/luke-large"])
-
-LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"unc-nlp/lxmert-base-uncased": "https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/config.json"}
-)
-
-TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["unc-nlp/lxmert-base-uncased"])
-
-M2M_100_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/config.json"}
-)
-
-M2M_100_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/m2m100_418M"])
-
-MAMBA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"state-spaces/mamba-2.8b": "https://huggingface.co/state-spaces/mamba-2.8b/resolve/main/config.json"}
-)
-
-MAMBA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList([])
-
-MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/markuplm-base": "https://huggingface.co/microsoft/markuplm-base/resolve/main/config.json",
- "microsoft/markuplm-large": "https://huggingface.co/microsoft/markuplm-large/resolve/main/config.json",
- }
-)
-
-MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/markuplm-base", "microsoft/markuplm-large"])
-
-MASK2FORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/mask2former-swin-small-coco-instance": "https://huggingface.co/facebook/mask2former-swin-small-coco-instance/blob/main/config.json"
- }
-)
-
-MASK2FORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/mask2former-swin-small-coco-instance"])
-
-MASKFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/maskformer-swin-base-ade": "https://huggingface.co/facebook/maskformer-swin-base-ade/blob/main/config.json"
- }
-)
-
-MASKFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/maskformer-swin-base-ade"])
-
-MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"mnaylor/mega-base-wikitext": "https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json"}
-)
-
-MEGA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["mnaylor/mega-base-wikitext"])
-
-MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict({})
-
-MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["nvidia/megatron-bert-cased-345m"])
-
-MGP_STR_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"alibaba-damo/mgp-str-base": "https://huggingface.co/alibaba-damo/mgp-str-base/resolve/main/config.json"}
-)
-
-MGP_STR_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["alibaba-damo/mgp-str-base"])
-
-MISTRAL_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "mistralai/Mistral-7B-v0.1": "https://huggingface.co/mistralai/Mistral-7B-v0.1/resolve/main/config.json",
- "mistralai/Mistral-7B-Instruct-v0.1": "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1/resolve/main/config.json",
- }
-)
-
-MIXTRAL_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"mistral-ai/Mixtral-8x7B": "https://huggingface.co/mistral-ai/Mixtral-8x7B/resolve/main/config.json"}
-)
-
-MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/mobilebert-uncased": "https://huggingface.co/google/mobilebert-uncased/resolve/main/config.json"}
-)
-
-MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/mobilebert-uncased"])
-
-TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/mobilebert-uncased"])
-
-MOBILENET_V1_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/mobilenet_v1_1.0_224": "https://huggingface.co/google/mobilenet_v1_1.0_224/resolve/main/config.json",
- "google/mobilenet_v1_0.75_192": "https://huggingface.co/google/mobilenet_v1_0.75_192/resolve/main/config.json",
- }
-)
-
-MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["google/mobilenet_v1_1.0_224", "google/mobilenet_v1_0.75_192"]
-)
-
-MOBILENET_V2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/mobilenet_v2_1.4_224": "https://huggingface.co/google/mobilenet_v2_1.4_224/resolve/main/config.json",
- "google/mobilenet_v2_1.0_224": "https://huggingface.co/google/mobilenet_v2_1.0_224/resolve/main/config.json",
- "google/mobilenet_v2_0.75_160": "https://huggingface.co/google/mobilenet_v2_0.75_160/resolve/main/config.json",
- "google/mobilenet_v2_0.35_96": "https://huggingface.co/google/mobilenet_v2_0.35_96/resolve/main/config.json",
- }
-)
-
-MOBILENET_V2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/mobilenet_v2_1.4_224",
- "google/mobilenet_v2_1.0_224",
- "google/mobilenet_v2_0.37_160",
- "google/mobilenet_v2_0.35_96",
- ]
-)
-
-MOBILEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "apple/mobilevit-small": "https://huggingface.co/apple/mobilevit-small/resolve/main/config.json",
- "apple/mobilevit-x-small": "https://huggingface.co/apple/mobilevit-x-small/resolve/main/config.json",
- "apple/mobilevit-xx-small": "https://huggingface.co/apple/mobilevit-xx-small/resolve/main/config.json",
- "apple/deeplabv3-mobilevit-small": "https://huggingface.co/apple/deeplabv3-mobilevit-small/resolve/main/config.json",
- "apple/deeplabv3-mobilevit-x-small": "https://huggingface.co/apple/deeplabv3-mobilevit-x-small/resolve/main/config.json",
- "apple/deeplabv3-mobilevit-xx-small": "https://huggingface.co/apple/deeplabv3-mobilevit-xx-small/resolve/main/config.json",
- }
-)
-
-MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "apple/mobilevit-small",
- "apple/mobilevit-x-small",
- "apple/mobilevit-xx-small",
- "apple/deeplabv3-mobilevit-small",
- "apple/deeplabv3-mobilevit-x-small",
- "apple/deeplabv3-mobilevit-xx-small",
- ]
-)
-
-TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "apple/mobilevit-small",
- "apple/mobilevit-x-small",
- "apple/mobilevit-xx-small",
- "apple/deeplabv3-mobilevit-small",
- "apple/deeplabv3-mobilevit-x-small",
- "apple/deeplabv3-mobilevit-xx-small",
- ]
-)
-
-MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"apple/mobilevitv2-1.0": "https://huggingface.co/apple/mobilevitv2-1.0/resolve/main/config.json"}
-)
-
-MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["apple/mobilevitv2-1.0-imagenet1k-256"])
-
-MPNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/mpnet-base": "https://huggingface.co/microsoft/mpnet-base/resolve/main/config.json"}
-)
-
-MPNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/mpnet-base"])
-
-TF_MPNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/mpnet-base"])
-
-MPT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"mosaicml/mpt-7b": "https://huggingface.co/mosaicml/mpt-7b/resolve/main/config.json"}
-)
-
-MPT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "mosaicml/mpt-7b",
- "mosaicml/mpt-7b-storywriter",
- "mosaicml/mpt-7b-instruct",
- "mosaicml/mpt-7b-8k",
- "mosaicml/mpt-7b-8k-instruct",
- "mosaicml/mpt-7b-8k-chat",
- "mosaicml/mpt-30b",
- "mosaicml/mpt-30b-instruct",
- "mosaicml/mpt-30b-chat",
- ]
-)
-
-MRA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"uw-madison/mra-base-512-4": "https://huggingface.co/uw-madison/mra-base-512-4/resolve/main/config.json"}
-)
-
-MRA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["uw-madison/mra-base-512-4"])
-
-MUSICGEN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/musicgen-small": "https://huggingface.co/facebook/musicgen-small/resolve/main/config.json"}
-)
-
-MUSICGEN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/musicgen-small"])
-
-MUSICGEN_MELODY_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/musicgen-melody": "https://huggingface.co/facebook/musicgen-melody/resolve/main/config.json"}
-)
-
-MUSICGEN_MELODY_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/musicgen-melody"])
-
-MVP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "RUCAIBox/mvp",
- "RUCAIBox/mvp-data-to-text",
- "RUCAIBox/mvp-open-dialog",
- "RUCAIBox/mvp-question-answering",
- "RUCAIBox/mvp-question-generation",
- "RUCAIBox/mvp-story",
- "RUCAIBox/mvp-summarization",
- "RUCAIBox/mvp-task-dialog",
- "RUCAIBox/mtl-data-to-text",
- "RUCAIBox/mtl-multi-task",
- "RUCAIBox/mtl-open-dialog",
- "RUCAIBox/mtl-question-answering",
- "RUCAIBox/mtl-question-generation",
- "RUCAIBox/mtl-story",
- "RUCAIBox/mtl-summarization",
- ]
-)
-
-NAT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"shi-labs/nat-mini-in1k-224": "https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json"}
-)
-
-NAT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["shi-labs/nat-mini-in1k-224"])
-
-NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"sijunhe/nezha-cn-base": "https://huggingface.co/sijunhe/nezha-cn-base/resolve/main/config.json"}
-)
-
-NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["sijunhe/nezha-cn-base", "sijunhe/nezha-cn-large", "sijunhe/nezha-base-wwm", "sijunhe/nezha-large-wwm"]
-)
-
-NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/nllb-moe-54B": "https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json"}
-)
-
-NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/nllb-moe-54b"])
-
-NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"uw-madison/nystromformer-512": "https://huggingface.co/uw-madison/nystromformer-512/resolve/main/config.json"}
-)
-
-NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["uw-madison/nystromformer-512"])
-
-OLMO_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "allenai/OLMo-1B-hf": "https://huggingface.co/allenai/OLMo-1B-hf/resolve/main/config.json",
- "allenai/OLMo-7B-hf": "https://huggingface.co/allenai/OLMo-7B-hf/resolve/main/config.json",
- "allenai/OLMo-7B-Twin-2T-hf": "https://huggingface.co/allenai/OLMo-7B-Twin-2T-hf/resolve/main/config.json",
- }
-)
-
-ONEFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "shi-labs/oneformer_ade20k_swin_tiny": "https://huggingface.co/shi-labs/oneformer_ade20k_swin_tiny/blob/main/config.json"
- }
-)
-
-ONEFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["shi-labs/oneformer_ade20k_swin_tiny"])
-
-OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"openai-community/openai-gpt": "https://huggingface.co/openai-community/openai-gpt/resolve/main/config.json"}
-)
-
-OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["openai-community/openai-gpt"])
-
-TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["openai-community/openai-gpt"])
-
-OPT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "facebook/opt-125m",
- "facebook/opt-350m",
- "facebook/opt-1.3b",
- "facebook/opt-2.7b",
- "facebook/opt-6.7b",
- "facebook/opt-13b",
- "facebook/opt-30b",
- ]
-)
-
-OWLV2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/owlv2-base-patch16": "https://huggingface.co/google/owlv2-base-patch16/resolve/main/config.json"}
-)
-
-OWLV2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/owlv2-base-patch16-ensemble"])
-
-OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/owlvit-base-patch32": "https://huggingface.co/google/owlvit-base-patch32/resolve/main/config.json",
- "google/owlvit-base-patch16": "https://huggingface.co/google/owlvit-base-patch16/resolve/main/config.json",
- "google/owlvit-large-patch14": "https://huggingface.co/google/owlvit-large-patch14/resolve/main/config.json",
- }
-)
-
-OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["google/owlvit-base-patch32", "google/owlvit-base-patch16", "google/owlvit-large-patch14"]
-)
-
-PATCHTSMIXER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "ibm/patchtsmixer-etth1-pretrain": "https://huggingface.co/ibm/patchtsmixer-etth1-pretrain/resolve/main/config.json"
- }
-)
-
-PATCHTSMIXER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["ibm/patchtsmixer-etth1-pretrain"])
-
-PATCHTST_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"ibm/patchtst-base": "https://huggingface.co/ibm/patchtst-base/resolve/main/config.json"}
-)
-
-PATCHTST_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["ibm/patchtst-etth1-pretrain"])
-
-PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/pegasus-large": "https://huggingface.co/google/pegasus-large/resolve/main/config.json"}
-)
-
-PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/pegasus-x-base": "https://huggingface.co/google/pegasus-x-base/resolve/main/config.json",
- "google/pegasus-x-large": "https://huggingface.co/google/pegasus-x-large/resolve/main/config.json",
- }
-)
-
-PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/pegasus-x-base", "google/pegasus-x-large"])
-
-PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"deepmind/language-perceiver": "https://huggingface.co/deepmind/language-perceiver/resolve/main/config.json"}
-)
-
-PERCEIVER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["deepmind/language-perceiver"])
-
-PERSIMMON_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"adept/persimmon-8b-base": "https://huggingface.co/adept/persimmon-8b-base/resolve/main/config.json"}
-)
-
-PHI_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/phi-1": "https://huggingface.co/microsoft/phi-1/resolve/main/config.json",
- "microsoft/phi-1_5": "https://huggingface.co/microsoft/phi-1_5/resolve/main/config.json",
- "microsoft/phi-2": "https://huggingface.co/microsoft/phi-2/resolve/main/config.json",
- }
-)
-
-PHI_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/phi-1", "microsoft/phi-1_5", "microsoft/phi-2"])
-
-PIX2STRUCT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/pix2struct-textcaps-base": "https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json"
- }
-)
-
-PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/pix2struct-textcaps-base",
- "google/pix2struct-textcaps-large",
- "google/pix2struct-base",
- "google/pix2struct-large",
- "google/pix2struct-ai2d-base",
- "google/pix2struct-ai2d-large",
- "google/pix2struct-widget-captioning-base",
- "google/pix2struct-widget-captioning-large",
- "google/pix2struct-screen2words-base",
- "google/pix2struct-screen2words-large",
- "google/pix2struct-docvqa-base",
- "google/pix2struct-docvqa-large",
- "google/pix2struct-ocrvqa-base",
- "google/pix2struct-ocrvqa-large",
- "google/pix2struct-chartqa-base",
- "google/pix2struct-inforgraphics-vqa-base",
- "google/pix2struct-inforgraphics-vqa-large",
- ]
-)
-
-PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"uclanlp/plbart-base": "https://huggingface.co/uclanlp/plbart-base/resolve/main/config.json"}
-)
-
-PLBART_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["uclanlp/plbart-base", "uclanlp/plbart-cs-java", "uclanlp/plbart-multi_task-all"]
-)
-
-POOLFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"sail/poolformer_s12": "https://huggingface.co/sail/poolformer_s12/resolve/main/config.json"}
-)
-
-POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["sail/poolformer_s12"])
-
-POP2PIANO_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"sweetcocoa/pop2piano": "https://huggingface.co/sweetcocoa/pop2piano/blob/main/config.json"}
-)
-
-POP2PIANO_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["sweetcocoa/pop2piano"])
-
-PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/prophetnet-large-uncased": "https://huggingface.co/microsoft/prophetnet-large-uncased/resolve/main/config.json"
- }
-)
-
-PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/prophetnet-large-uncased"])
-
-PVT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict({"pvt-tiny-224": "https://huggingface.co/Zetatech/pvt-tiny-224"})
-
-PVT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Zetatech/pvt-tiny-224"])
-
-QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google-bert/bert-base-uncased": "https://huggingface.co/google-bert/bert-base-uncased/resolve/main/config.json"}
-)
-
-QDQBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google-bert/bert-base-uncased"])
-
-QWEN2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"Qwen/Qwen2-7B-beta": "https://huggingface.co/Qwen/Qwen2-7B-beta/resolve/main/config.json"}
-)
-
-REALM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/realm-cc-news-pretrained-embedder": "https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json",
- "google/realm-cc-news-pretrained-encoder": "https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json",
- "google/realm-cc-news-pretrained-scorer": "https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json",
- "google/realm-cc-news-pretrained-openqa": "https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json",
- "google/realm-orqa-nq-openqa": "https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json",
- "google/realm-orqa-nq-reader": "https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json",
- "google/realm-orqa-wq-openqa": "https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json",
- "google/realm-orqa-wq-reader": "https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json",
- }
-)
-
-REALM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/realm-cc-news-pretrained-embedder",
- "google/realm-cc-news-pretrained-encoder",
- "google/realm-cc-news-pretrained-scorer",
- "google/realm-cc-news-pretrained-openqa",
- "google/realm-orqa-nq-openqa",
- "google/realm-orqa-nq-reader",
- "google/realm-orqa-wq-openqa",
- "google/realm-orqa-wq-reader",
- ]
-)
-
-REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/reformer-crime-and-punishment": "https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/config.json",
- "google/reformer-enwik8": "https://huggingface.co/google/reformer-enwik8/resolve/main/config.json",
- }
-)
-
-REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["google/reformer-crime-and-punishment", "google/reformer-enwik8"]
-)
-
-REGNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/regnet-y-040": "https://huggingface.co/facebook/regnet-y-040/blob/main/config.json"}
-)
-
-REGNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/regnet-y-040"])
-
-TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/regnet-y-040"])
-
-REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/rembert": "https://huggingface.co/google/rembert/resolve/main/config.json"}
-)
-
-REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/rembert"])
-
-TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/rembert"])
-
-RESNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/resnet-50": "https://huggingface.co/microsoft/resnet-50/blob/main/config.json"}
-)
-
-RESNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/resnet-50"])
-
-TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/resnet-50"])
-
-ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "FacebookAI/roberta-base": "https://huggingface.co/FacebookAI/roberta-base/resolve/main/config.json",
- "FacebookAI/roberta-large": "https://huggingface.co/FacebookAI/roberta-large/resolve/main/config.json",
- "FacebookAI/roberta-large-mnli": "https://huggingface.co/FacebookAI/roberta-large-mnli/resolve/main/config.json",
- "distilbert/distilroberta-base": "https://huggingface.co/distilbert/distilroberta-base/resolve/main/config.json",
- "openai-community/roberta-base-openai-detector": "https://huggingface.co/openai-community/roberta-base-openai-detector/resolve/main/config.json",
- "openai-community/roberta-large-openai-detector": "https://huggingface.co/openai-community/roberta-large-openai-detector/resolve/main/config.json",
- }
-)
-
-ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "FacebookAI/roberta-base",
- "FacebookAI/roberta-large",
- "FacebookAI/roberta-large-mnli",
- "distilbert/distilroberta-base",
- "openai-community/roberta-base-openai-detector",
- "openai-community/roberta-large-openai-detector",
- ]
-)
-
-TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "FacebookAI/roberta-base",
- "FacebookAI/roberta-large",
- "FacebookAI/roberta-large-mnli",
- "distilbert/distilroberta-base",
- ]
-)
-
-ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "andreasmadsen/efficient_mlm_m0.40": "https://huggingface.co/andreasmadsen/efficient_mlm_m0.40/resolve/main/config.json"
- }
-)
-
-ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "andreasmadsen/efficient_mlm_m0.15",
- "andreasmadsen/efficient_mlm_m0.20",
- "andreasmadsen/efficient_mlm_m0.30",
- "andreasmadsen/efficient_mlm_m0.40",
- "andreasmadsen/efficient_mlm_m0.50",
- "andreasmadsen/efficient_mlm_m0.60",
- "andreasmadsen/efficient_mlm_m0.70",
- "andreasmadsen/efficient_mlm_m0.80",
- ]
-)
-
-TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "andreasmadsen/efficient_mlm_m0.15",
- "andreasmadsen/efficient_mlm_m0.20",
- "andreasmadsen/efficient_mlm_m0.30",
- "andreasmadsen/efficient_mlm_m0.40",
- "andreasmadsen/efficient_mlm_m0.50",
- "andreasmadsen/efficient_mlm_m0.60",
- "andreasmadsen/efficient_mlm_m0.70",
- "andreasmadsen/efficient_mlm_m0.80",
- ]
-)
-
-ROC_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"weiweishi/roc-bert-base-zh": "https://huggingface.co/weiweishi/roc-bert-base-zh/resolve/main/config.json"}
-)
-
-ROC_BERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["weiweishi/roc-bert-base-zh"])
-
-ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "junnyu/roformer_chinese_small": "https://huggingface.co/junnyu/roformer_chinese_small/resolve/main/config.json",
- "junnyu/roformer_chinese_base": "https://huggingface.co/junnyu/roformer_chinese_base/resolve/main/config.json",
- "junnyu/roformer_chinese_char_small": "https://huggingface.co/junnyu/roformer_chinese_char_small/resolve/main/config.json",
- "junnyu/roformer_chinese_char_base": "https://huggingface.co/junnyu/roformer_chinese_char_base/resolve/main/config.json",
- "junnyu/roformer_small_discriminator": "https://huggingface.co/junnyu/roformer_small_discriminator/resolve/main/config.json",
- "junnyu/roformer_small_generator": "https://huggingface.co/junnyu/roformer_small_generator/resolve/main/config.json",
- }
-)
-
-ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "junnyu/roformer_chinese_small",
- "junnyu/roformer_chinese_base",
- "junnyu/roformer_chinese_char_small",
- "junnyu/roformer_chinese_char_base",
- "junnyu/roformer_small_discriminator",
- "junnyu/roformer_small_generator",
- ]
-)
-
-TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "junnyu/roformer_chinese_small",
- "junnyu/roformer_chinese_base",
- "junnyu/roformer_chinese_char_small",
- "junnyu/roformer_chinese_char_base",
- "junnyu/roformer_small_discriminator",
- "junnyu/roformer_small_generator",
- ]
-)
-
-RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "RWKV/rwkv-4-169m-pile": "https://huggingface.co/RWKV/rwkv-4-169m-pile/resolve/main/config.json",
- "RWKV/rwkv-4-430m-pile": "https://huggingface.co/RWKV/rwkv-4-430m-pile/resolve/main/config.json",
- "RWKV/rwkv-4-1b5-pile": "https://huggingface.co/RWKV/rwkv-4-1b5-pile/resolve/main/config.json",
- "RWKV/rwkv-4-3b-pile": "https://huggingface.co/RWKV/rwkv-4-3b-pile/resolve/main/config.json",
- "RWKV/rwkv-4-7b-pile": "https://huggingface.co/RWKV/rwkv-4-7b-pile/resolve/main/config.json",
- "RWKV/rwkv-4-14b-pile": "https://huggingface.co/RWKV/rwkv-4-14b-pile/resolve/main/config.json",
- "RWKV/rwkv-raven-1b5": "https://huggingface.co/RWKV/rwkv-raven-1b5/resolve/main/config.json",
- "RWKV/rwkv-raven-3b": "https://huggingface.co/RWKV/rwkv-raven-3b/resolve/main/config.json",
- "RWKV/rwkv-raven-7b": "https://huggingface.co/RWKV/rwkv-raven-7b/resolve/main/config.json",
- "RWKV/rwkv-raven-14b": "https://huggingface.co/RWKV/rwkv-raven-14b/resolve/main/config.json",
- }
-)
-
-RWKV_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "RWKV/rwkv-4-169m-pile",
- "RWKV/rwkv-4-430m-pile",
- "RWKV/rwkv-4-1b5-pile",
- "RWKV/rwkv-4-3b-pile",
- "RWKV/rwkv-4-7b-pile",
- "RWKV/rwkv-4-14b-pile",
- "RWKV/rwkv-raven-1b5",
- "RWKV/rwkv-raven-3b",
- "RWKV/rwkv-raven-7b",
- "RWKV/rwkv-raven-14b",
- ]
-)
-
-SAM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/sam-vit-huge": "https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json",
- "facebook/sam-vit-large": "https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json",
- "facebook/sam-vit-base": "https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json",
- }
-)
-
-SAM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["facebook/sam-vit-huge", "facebook/sam-vit-large", "facebook/sam-vit-base"]
-)
-
-TF_SAM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["facebook/sam-vit-huge", "facebook/sam-vit-large", "facebook/sam-vit-base"]
-)
-
-SEAMLESS_M4T_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/hf-seamless-m4t-medium": "https://huggingface.co/facebook/hf-seamless-m4t-medium/resolve/main/config.json"
- }
-)
-
-SEAMLESS_M4T_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/hf-seamless-m4t-medium"])
-
-SEAMLESS_M4T_V2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"": "https://huggingface.co//resolve/main/config.json"}
-)
-
-SEAMLESS_M4T_V2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/seamless-m4t-v2-large"])
-
-SEGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "nvidia/segformer-b0-finetuned-ade-512-512": "https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512/resolve/main/config.json"
- }
-)
-
-SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["nvidia/segformer-b0-finetuned-ade-512-512"])
-
-TF_SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["nvidia/segformer-b0-finetuned-ade-512-512"])
-
-SEGGPT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"BAAI/seggpt-vit-large": "https://huggingface.co/BAAI/seggpt-vit-large/resolve/main/config.json"}
-)
-
-SEGGPT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["BAAI/seggpt-vit-large"])
-
-SEW_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"asapp/sew-tiny-100k": "https://huggingface.co/asapp/sew-tiny-100k/resolve/main/config.json"}
-)
-
-SEW_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["asapp/sew-tiny-100k", "asapp/sew-small-100k", "asapp/sew-mid-100k"]
-)
-
-SEW_D_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"asapp/sew-d-tiny-100k": "https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json"}
-)
-
-SEW_D_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "asapp/sew-d-tiny-100k",
- "asapp/sew-d-small-100k",
- "asapp/sew-d-mid-100k",
- "asapp/sew-d-mid-k127-100k",
- "asapp/sew-d-base-100k",
- "asapp/sew-d-base-plus-100k",
- "asapp/sew-d-mid-400k",
- "asapp/sew-d-mid-k127-400k",
- "asapp/sew-d-base-plus-400k",
- ]
-)
-
-SIGLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/siglip-base-patch16-224": "https://huggingface.co/google/siglip-base-patch16-224/resolve/main/config.json"
- }
-)
-
-SIGLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/siglip-base-patch16-224"])
-
-SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/s2t-small-librispeech-asr": "https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/config.json"
- }
-)
-
-SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/s2t-small-librispeech-asr"])
-
-TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/s2t-small-librispeech-asr"])
-
-SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/s2t-wav2vec2-large-en-de": "https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/config.json"
- }
-)
-
-SPEECHT5_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/speecht5_asr": "https://huggingface.co/microsoft/speecht5_asr/resolve/main/config.json",
- "microsoft/speecht5_tts": "https://huggingface.co/microsoft/speecht5_tts/resolve/main/config.json",
- "microsoft/speecht5_vc": "https://huggingface.co/microsoft/speecht5_vc/resolve/main/config.json",
- }
-)
-
-SPEECHT5_PRETRAINED_HIFIGAN_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/speecht5_hifigan": "https://huggingface.co/microsoft/speecht5_hifigan/resolve/main/config.json"}
-)
-
-SPEECHT5_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["microsoft/speecht5_asr", "microsoft/speecht5_tts", "microsoft/speecht5_vc"]
-)
-
-SPLINTER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "tau/splinter-base": "https://huggingface.co/tau/splinter-base/resolve/main/config.json",
- "tau/splinter-base-qass": "https://huggingface.co/tau/splinter-base-qass/resolve/main/config.json",
- "tau/splinter-large": "https://huggingface.co/tau/splinter-large/resolve/main/config.json",
- "tau/splinter-large-qass": "https://huggingface.co/tau/splinter-large-qass/resolve/main/config.json",
- }
-)
-
-SPLINTER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["tau/splinter-base", "tau/splinter-base-qass", "tau/splinter-large", "tau/splinter-large-qass"]
-)
-
-SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "squeezebert/squeezebert-uncased": "https://huggingface.co/squeezebert/squeezebert-uncased/resolve/main/config.json",
- "squeezebert/squeezebert-mnli": "https://huggingface.co/squeezebert/squeezebert-mnli/resolve/main/config.json",
- "squeezebert/squeezebert-mnli-headless": "https://huggingface.co/squeezebert/squeezebert-mnli-headless/resolve/main/config.json",
- }
-)
-
-SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["squeezebert/squeezebert-uncased", "squeezebert/squeezebert-mnli", "squeezebert/squeezebert-mnli-headless"]
-)
-
-STABLELM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"stabilityai/stablelm-3b-4e1t": "https://huggingface.co/stabilityai/stablelm-3b-4e1t/resolve/main/config.json"}
-)
-
-STARCODER2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict({})
-
-SWIFTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"MBZUAI/swiftformer-xs": "https://huggingface.co/MBZUAI/swiftformer-xs/resolve/main/config.json"}
-)
-
-SWIFTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["MBZUAI/swiftformer-xs"])
-
-SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/swin-tiny-patch4-window7-224": "https://huggingface.co/microsoft/swin-tiny-patch4-window7-224/resolve/main/config.json"
- }
-)
-
-SWIN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/swin-tiny-patch4-window7-224"])
-
-TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/swin-tiny-patch4-window7-224"])
-
-SWIN2SR_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "caidas/swin2sr-classicalsr-x2-64": "https://huggingface.co/caidas/swin2sr-classicalsr-x2-64/resolve/main/config.json"
- }
-)
-
-SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["caidas/swin2SR-classical-sr-x2-64"])
-
-SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/swinv2-tiny-patch4-window8-256": "https://huggingface.co/microsoft/swinv2-tiny-patch4-window8-256/resolve/main/config.json"
- }
-)
-
-SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/swinv2-tiny-patch4-window8-256"])
-
-SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/switch-base-8": "https://huggingface.co/google/switch-base-8/blob/main/config.json"}
-)
-
-SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/switch-base-8",
- "google/switch-base-16",
- "google/switch-base-32",
- "google/switch-base-64",
- "google/switch-base-128",
- "google/switch-base-256",
- "google/switch-large-128",
- "google/switch-xxl-128",
- "google/switch-c-2048",
- ]
-)
-
-T5_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google-t5/t5-small": "https://huggingface.co/google-t5/t5-small/resolve/main/config.json",
- "google-t5/t5-base": "https://huggingface.co/google-t5/t5-base/resolve/main/config.json",
- "google-t5/t5-large": "https://huggingface.co/google-t5/t5-large/resolve/main/config.json",
- "google-t5/t5-3b": "https://huggingface.co/google-t5/t5-3b/resolve/main/config.json",
- "google-t5/t5-11b": "https://huggingface.co/google-t5/t5-11b/resolve/main/config.json",
- }
-)
-
-T5_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["google-t5/t5-small", "google-t5/t5-base", "google-t5/t5-large", "google-t5/t5-3b", "google-t5/t5-11b"]
-)
-
-TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["google-t5/t5-small", "google-t5/t5-base", "google-t5/t5-large", "google-t5/t5-3b", "google-t5/t5-11b"]
-)
-
-TABLE_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/table-transformer-detection": "https://huggingface.co/microsoft/table-transformer-detection/resolve/main/config.json"
- }
-)
-
-TABLE_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/table-transformer-detection"])
-
-TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/tapas-base-finetuned-sqa": "https://huggingface.co/google/tapas-base-finetuned-sqa/resolve/main/config.json",
- "google/tapas-base-finetuned-wtq": "https://huggingface.co/google/tapas-base-finetuned-wtq/resolve/main/config.json",
- "google/tapas-base-finetuned-wikisql-supervised": "https://huggingface.co/google/tapas-base-finetuned-wikisql-supervised/resolve/main/config.json",
- "google/tapas-base-finetuned-tabfact": "https://huggingface.co/google/tapas-base-finetuned-tabfact/resolve/main/config.json",
- }
-)
-
-TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/tapas-large",
- "google/tapas-large-finetuned-sqa",
- "google/tapas-large-finetuned-wtq",
- "google/tapas-large-finetuned-wikisql-supervised",
- "google/tapas-large-finetuned-tabfact",
- "google/tapas-base",
- "google/tapas-base-finetuned-sqa",
- "google/tapas-base-finetuned-wtq",
- "google/tapas-base-finetuned-wikisql-supervised",
- "google/tapas-base-finetuned-tabfact",
- "google/tapas-small",
- "google/tapas-small-finetuned-sqa",
- "google/tapas-small-finetuned-wtq",
- "google/tapas-small-finetuned-wikisql-supervised",
- "google/tapas-small-finetuned-tabfact",
- "google/tapas-mini",
- "google/tapas-mini-finetuned-sqa",
- "google/tapas-mini-finetuned-wtq",
- "google/tapas-mini-finetuned-wikisql-supervised",
- "google/tapas-mini-finetuned-tabfact",
- "google/tapas-tiny",
- "google/tapas-tiny-finetuned-sqa",
- "google/tapas-tiny-finetuned-wtq",
- "google/tapas-tiny-finetuned-wikisql-supervised",
- "google/tapas-tiny-finetuned-tabfact",
- ]
-)
-
-TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "google/tapas-large",
- "google/tapas-large-finetuned-sqa",
- "google/tapas-large-finetuned-wtq",
- "google/tapas-large-finetuned-wikisql-supervised",
- "google/tapas-large-finetuned-tabfact",
- "google/tapas-base",
- "google/tapas-base-finetuned-sqa",
- "google/tapas-base-finetuned-wtq",
- "google/tapas-base-finetuned-wikisql-supervised",
- "google/tapas-base-finetuned-tabfact",
- "google/tapas-small",
- "google/tapas-small-finetuned-sqa",
- "google/tapas-small-finetuned-wtq",
- "google/tapas-small-finetuned-wikisql-supervised",
- "google/tapas-small-finetuned-tabfact",
- "google/tapas-mini",
- "google/tapas-mini-finetuned-sqa",
- "google/tapas-mini-finetuned-wtq",
- "google/tapas-mini-finetuned-wikisql-supervised",
- "google/tapas-mini-finetuned-tabfact",
- "google/tapas-tiny",
- "google/tapas-tiny-finetuned-sqa",
- "google/tapas-tiny-finetuned-wtq",
- "google/tapas-tiny-finetuned-wikisql-supervised",
- "google/tapas-tiny-finetuned-tabfact",
- ]
-)
-
-TIME_SERIES_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "huggingface/time-series-transformer-tourism-monthly": "https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"
- }
-)
-
-TIME_SERIES_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["huggingface/time-series-transformer-tourism-monthly"]
-)
-
-TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/timesformer": "https://huggingface.co/facebook/timesformer/resolve/main/config.json"}
-)
-
-TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/timesformer-base-finetuned-k400"])
-
-TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/trocr-base-handwritten": "https://huggingface.co/microsoft/trocr-base-handwritten/resolve/main/config.json"
- }
-)
-
-TROCR_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/trocr-base-handwritten"])
-
-TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"ZinengTang/tvlt-base": "https://huggingface.co/ZinengTang/tvlt-base/blob/main/config.json"}
-)
-
-TVLT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["ZinengTang/tvlt-base"])
-
-TVP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"Intel/tvp-base": "https://huggingface.co/Intel/tvp-base/resolve/main/config.json"}
-)
-
-TVP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["Intel/tvp-base", "Intel/tvp-base-ANet"])
-
-UDOP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/udop-large": "https://huggingface.co/microsoft/udop-large/resolve/main/config.json"}
-)
-
-UDOP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/udop-large"])
-
-UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/unispeech-large-1500h-cv": "https://huggingface.co/microsoft/unispeech-large-1500h-cv/resolve/main/config.json"
- }
-)
-
-UNISPEECH_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["microsoft/unispeech-large-1500h-cv", "microsoft/unispeech-large-multi-lingual-1500h-cv"]
-)
-
-UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/unispeech-sat-base-100h-libri-ft": "https://huggingface.co/microsoft/unispeech-sat-base-100h-libri-ft/resolve/main/config.json"
- }
-)
-
-UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList([])
-
-UNIVNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"dg845/univnet-dev": "https://huggingface.co/dg845/univnet-dev/resolve/main/config.json"}
-)
-
-UNIVNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["dg845/univnet-dev"])
-
-VIDEOMAE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"MCG-NJU/videomae-base": "https://huggingface.co/MCG-NJU/videomae-base/resolve/main/config.json"}
-)
-
-VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["MCG-NJU/videomae-base"])
-
-VILT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"dandelin/vilt-b32-mlm": "https://huggingface.co/dandelin/vilt-b32-mlm/blob/main/config.json"}
-)
-
-VILT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["dandelin/vilt-b32-mlm"])
-
-VIPLLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"ybelkada/vip-llava-7b-hf": "https://huggingface.co/llava-hf/vip-llava-7b-hf/resolve/main/config.json"}
-)
-
-VIPLLAVA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["llava-hf/vip-llava-7b-hf"])
-
-VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "uclanlp/visualbert-vqa": "https://huggingface.co/uclanlp/visualbert-vqa/resolve/main/config.json",
- "uclanlp/visualbert-vqa-pre": "https://huggingface.co/uclanlp/visualbert-vqa-pre/resolve/main/config.json",
- "uclanlp/visualbert-vqa-coco-pre": "https://huggingface.co/uclanlp/visualbert-vqa-coco-pre/resolve/main/config.json",
- "uclanlp/visualbert-vcr": "https://huggingface.co/uclanlp/visualbert-vcr/resolve/main/config.json",
- "uclanlp/visualbert-vcr-pre": "https://huggingface.co/uclanlp/visualbert-vcr-pre/resolve/main/config.json",
- "uclanlp/visualbert-vcr-coco-pre": "https://huggingface.co/uclanlp/visualbert-vcr-coco-pre/resolve/main/config.json",
- "uclanlp/visualbert-nlvr2": "https://huggingface.co/uclanlp/visualbert-nlvr2/resolve/main/config.json",
- "uclanlp/visualbert-nlvr2-pre": "https://huggingface.co/uclanlp/visualbert-nlvr2-pre/resolve/main/config.json",
- "uclanlp/visualbert-nlvr2-coco-pre": "https://huggingface.co/uclanlp/visualbert-nlvr2-coco-pre/resolve/main/config.json",
- }
-)
-
-VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "uclanlp/visualbert-vqa",
- "uclanlp/visualbert-vqa-pre",
- "uclanlp/visualbert-vqa-coco-pre",
- "uclanlp/visualbert-vcr",
- "uclanlp/visualbert-vcr-pre",
- "uclanlp/visualbert-vcr-coco-pre",
- "uclanlp/visualbert-nlvr2",
- "uclanlp/visualbert-nlvr2-pre",
- "uclanlp/visualbert-nlvr2-coco-pre",
- ]
-)
-
-VIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/vit-base-patch16-224": "https://huggingface.co/vit-base-patch16-224/resolve/main/config.json"}
-)
-
-VIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/vit-base-patch16-224"])
-
-VIT_HYBRID_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"google/vit-hybrid-base-bit-384": "https://huggingface.co/vit-hybrid-base-bit-384/resolve/main/config.json"}
-)
-
-VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/vit-hybrid-base-bit-384"])
-
-VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/vit-mae-base": "https://huggingface.co/facebook/vit-mae-base/resolve/main/config.json"}
-)
-
-VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/vit-mae-base"])
-
-VIT_MSN_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"sayakpaul/vit-msn-base": "https://huggingface.co/sayakpaul/vit-msn-base/resolve/main/config.json"}
-)
-
-VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/vit-msn-small"])
-
-VITDET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/vit-det-base": "https://huggingface.co/facebook/vit-det-base/resolve/main/config.json"}
-)
-
-VITDET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/vit-det-base"])
-
-VITMATTE_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "hustvl/vitmatte-small-composition-1k": "https://huggingface.co/hustvl/vitmatte-small-composition-1k/resolve/main/config.json"
- }
-)
-
-VITMATTE_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["hustvl/vitmatte-small-composition-1k"])
-
-VITS_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/mms-tts-eng": "https://huggingface.co/facebook/mms-tts-eng/resolve/main/config.json"}
-)
-
-VITS_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/mms-tts-eng"])
-
-VIVIT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "google/vivit-b-16x2-kinetics400": "https://huggingface.co/google/vivit-b-16x2-kinetics400/resolve/main/config.json"
- }
-)
-
-VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["google/vivit-b-16x2-kinetics400"])
-
-WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/wav2vec2-base-960h": "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/config.json"}
-)
-
-WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "facebook/wav2vec2-base-960h",
- "facebook/wav2vec2-large-960h",
- "facebook/wav2vec2-large-960h-lv60",
- "facebook/wav2vec2-large-960h-lv60-self",
- ]
-)
-
-TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "facebook/wav2vec2-base-960h",
- "facebook/wav2vec2-large-960h",
- "facebook/wav2vec2-large-960h-lv60",
- "facebook/wav2vec2-large-960h-lv60-self",
- ]
-)
-
-WAV2VEC2_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/w2v-bert-2.0": "https://huggingface.co/facebook/w2v-bert-2.0/resolve/main/config.json"}
-)
-
-WAV2VEC2_BERT_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/w2v-bert-2.0"])
-
-WAV2VEC2_CONFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/wav2vec2-conformer-rel-pos-large": "https://huggingface.co/facebook/wav2vec2-conformer-rel-pos-large/resolve/main/config.json"
- }
-)
-
-WAV2VEC2_CONFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/wav2vec2-conformer-rel-pos-large"])
-
-WAVLM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/wavlm-base": "https://huggingface.co/microsoft/wavlm-base/resolve/main/config.json"}
-)
-
-WAVLM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["microsoft/wavlm-base", "microsoft/wavlm-base-plus", "microsoft/wavlm-large"]
-)
-
-WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"openai/whisper-base": "https://huggingface.co/openai/whisper-base/resolve/main/config.json"}
-)
-
-WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["openai/whisper-base"])
-
-TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["openai/whisper-base"])
-
-XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"microsoft/xclip-base-patch32": "https://huggingface.co/microsoft/xclip-base-patch32/resolve/main/config.json"}
-)
-
-XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/xclip-base-patch32"])
-
-XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"facebook/xglm-564M": "https://huggingface.co/facebook/xglm-564M/resolve/main/config.json"}
-)
-
-XGLM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/xglm-564M"])
-
-TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/xglm-564M"])
-
-XLM_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "FacebookAI/xlm-mlm-en-2048": "https://huggingface.co/FacebookAI/xlm-mlm-en-2048/resolve/main/config.json",
- "FacebookAI/xlm-mlm-ende-1024": "https://huggingface.co/FacebookAI/xlm-mlm-ende-1024/resolve/main/config.json",
- "FacebookAI/xlm-mlm-enfr-1024": "https://huggingface.co/FacebookAI/xlm-mlm-enfr-1024/resolve/main/config.json",
- "FacebookAI/xlm-mlm-enro-1024": "https://huggingface.co/FacebookAI/xlm-mlm-enro-1024/resolve/main/config.json",
- "FacebookAI/xlm-mlm-tlm-xnli15-1024": "https://huggingface.co/FacebookAI/xlm-mlm-tlm-xnli15-1024/resolve/main/config.json",
- "FacebookAI/xlm-mlm-xnli15-1024": "https://huggingface.co/FacebookAI/xlm-mlm-xnli15-1024/resolve/main/config.json",
- "FacebookAI/xlm-clm-enfr-1024": "https://huggingface.co/FacebookAI/xlm-clm-enfr-1024/resolve/main/config.json",
- "FacebookAI/xlm-clm-ende-1024": "https://huggingface.co/FacebookAI/xlm-clm-ende-1024/resolve/main/config.json",
- "FacebookAI/xlm-mlm-17-1280": "https://huggingface.co/FacebookAI/xlm-mlm-17-1280/resolve/main/config.json",
- "FacebookAI/xlm-mlm-100-1280": "https://huggingface.co/FacebookAI/xlm-mlm-100-1280/resolve/main/config.json",
- }
-)
-
-XLM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "FacebookAI/xlm-mlm-en-2048",
- "FacebookAI/xlm-mlm-ende-1024",
- "FacebookAI/xlm-mlm-enfr-1024",
- "FacebookAI/xlm-mlm-enro-1024",
- "FacebookAI/xlm-mlm-tlm-xnli15-1024",
- "FacebookAI/xlm-mlm-xnli15-1024",
- "FacebookAI/xlm-clm-enfr-1024",
- "FacebookAI/xlm-clm-ende-1024",
- "FacebookAI/xlm-mlm-17-1280",
- "FacebookAI/xlm-mlm-100-1280",
- ]
-)
-
-TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "FacebookAI/xlm-mlm-en-2048",
- "FacebookAI/xlm-mlm-ende-1024",
- "FacebookAI/xlm-mlm-enfr-1024",
- "FacebookAI/xlm-mlm-enro-1024",
- "FacebookAI/xlm-mlm-tlm-xnli15-1024",
- "FacebookAI/xlm-mlm-xnli15-1024",
- "FacebookAI/xlm-clm-enfr-1024",
- "FacebookAI/xlm-clm-ende-1024",
- "FacebookAI/xlm-mlm-17-1280",
- "FacebookAI/xlm-mlm-100-1280",
- ]
-)
-
-XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "microsoft/xprophetnet-large-wiki100-cased": "https://huggingface.co/microsoft/xprophetnet-large-wiki100-cased/resolve/main/config.json"
- }
-)
-
-XLM_PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["microsoft/xprophetnet-large-wiki100-cased"])
-
-XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "FacebookAI/xlm-roberta-base": "https://huggingface.co/FacebookAI/xlm-roberta-base/resolve/main/config.json",
- "FacebookAI/xlm-roberta-large": "https://huggingface.co/FacebookAI/xlm-roberta-large/resolve/main/config.json",
- "FacebookAI/xlm-roberta-large-finetuned-conll02-dutch": "https://huggingface.co/FacebookAI/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/config.json",
- "FacebookAI/xlm-roberta-large-finetuned-conll02-spanish": "https://huggingface.co/FacebookAI/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/config.json",
- "FacebookAI/xlm-roberta-large-finetuned-conll03-english": "https://huggingface.co/FacebookAI/xlm-roberta-large-finetuned-conll03-english/resolve/main/config.json",
- "FacebookAI/xlm-roberta-large-finetuned-conll03-german": "https://huggingface.co/FacebookAI/xlm-roberta-large-finetuned-conll03-german/resolve/main/config.json",
- }
-)
-
-XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "FacebookAI/xlm-roberta-base",
- "FacebookAI/xlm-roberta-large",
- "FacebookAI/xlm-roberta-large-finetuned-conll02-dutch",
- "FacebookAI/xlm-roberta-large-finetuned-conll02-spanish",
- "FacebookAI/xlm-roberta-large-finetuned-conll03-english",
- "FacebookAI/xlm-roberta-large-finetuned-conll03-german",
- ]
-)
-
-TF_XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "FacebookAI/xlm-roberta-base",
- "FacebookAI/xlm-roberta-large",
- "joeddav/xlm-roberta-large-xnli",
- "cardiffnlp/twitter-xlm-roberta-base-sentiment",
- ]
-)
-
-FLAX_XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- ["FacebookAI/xlm-roberta-base", "FacebookAI/xlm-roberta-large"]
-)
-
-XLM_ROBERTA_XL_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/xlm-roberta-xl": "https://huggingface.co/facebook/xlm-roberta-xl/resolve/main/config.json",
- "facebook/xlm-roberta-xxl": "https://huggingface.co/facebook/xlm-roberta-xxl/resolve/main/config.json",
- }
-)
-
-XLM_ROBERTA_XL_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["facebook/xlm-roberta-xl", "facebook/xlm-roberta-xxl"])
-
-XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "xlnet/xlnet-base-cased": "https://huggingface.co/xlnet/xlnet-base-cased/resolve/main/config.json",
- "xlnet/xlnet-large-cased": "https://huggingface.co/xlnet/xlnet-large-cased/resolve/main/config.json",
- }
-)
-
-XLNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["xlnet/xlnet-base-cased", "xlnet/xlnet-large-cased"])
-
-TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["xlnet/xlnet-base-cased", "xlnet/xlnet-large-cased"])
-
-XMOD_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {
- "facebook/xmod-base": "https://huggingface.co/facebook/xmod-base/resolve/main/config.json",
- "facebook/xmod-large-prenorm": "https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json",
- "facebook/xmod-base-13-125k": "https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json",
- "facebook/xmod-base-30-125k": "https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json",
- "facebook/xmod-base-30-195k": "https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json",
- "facebook/xmod-base-60-125k": "https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json",
- "facebook/xmod-base-60-265k": "https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json",
- "facebook/xmod-base-75-125k": "https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json",
- "facebook/xmod-base-75-269k": "https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json",
- }
-)
-
-XMOD_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(
- [
- "facebook/xmod-base",
- "facebook/xmod-large-prenorm",
- "facebook/xmod-base-13-125k",
- "facebook/xmod-base-30-125k",
- "facebook/xmod-base-30-195k",
- "facebook/xmod-base-60-125k",
- "facebook/xmod-base-60-265k",
- "facebook/xmod-base-75-125k",
- "facebook/xmod-base-75-269k",
- ]
-)
-
-YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"hustvl/yolos-small": "https://huggingface.co/hustvl/yolos-small/resolve/main/config.json"}
-)
-
-YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["hustvl/yolos-small"])
-
-YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP = DeprecatedDict(
- {"uw-madison/yoso-4096": "https://huggingface.co/uw-madison/yoso-4096/resolve/main/config.json"}
-)
-
-YOSO_PRETRAINED_MODEL_ARCHIVE_LIST = DeprecatedList(["uw-madison/yoso-4096"])
-
-
-CONFIG_ARCHIVE_MAP_MAPPING_NAMES = OrderedDict(
- [
- # Add archive maps here)
- ("albert", "ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("align", "ALIGN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("altclip", "ALTCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("audio-spectrogram-transformer", "AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("autoformer", "AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("bark", "BARK_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("bart", "BART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("beit", "BEIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("bert", "BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("big_bird", "BIG_BIRD_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("bigbird_pegasus", "BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("biogpt", "BIOGPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("bit", "BIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("blenderbot", "BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("blenderbot-small", "BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("blip", "BLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("blip-2", "BLIP_2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("bloom", "BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("bridgetower", "BRIDGETOWER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("bros", "BROS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("camembert", "CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("canine", "CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("chinese_clip", "CHINESE_CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("clap", "CLAP_PRETRAINED_MODEL_ARCHIVE_LIST"),
- ("clip", "CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("clipseg", "CLIPSEG_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("clvp", "CLVP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("codegen", "CODEGEN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("conditional_detr", "CONDITIONAL_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("convbert", "CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("convnext", "CONVNEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("convnextv2", "CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("cpmant", "CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("ctrl", "CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("cvt", "CVT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("data2vec-audio", "DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("data2vec-text", "DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("data2vec-vision", "DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("deberta", "DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("deberta-v2", "DEBERTA_V2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("deformable_detr", "DEFORMABLE_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("deit", "DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("depth_anything", "DEPTH_ANYTHING_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("deta", "DETA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("detr", "DETR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("dinat", "DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("dinov2", "DINOV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("distilbert", "DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("donut-swin", "DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("dpr", "DPR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("dpt", "DPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("efficientformer", "EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("efficientnet", "EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("electra", "ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("encodec", "ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("ernie", "ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("ernie_m", "ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("esm", "ESM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("falcon", "FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("fastspeech2_conformer", "FASTSPEECH2_CONFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("flaubert", "FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("flava", "FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("fnet", "FNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("focalnet", "FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("fsmt", "FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("funnel", "FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("fuyu", "FUYU_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("gemma", "GEMMA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("git", "GIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("glpn", "GLPN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("gpt2", "GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("gpt_bigcode", "GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("gpt_neo", "GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("gpt_neox", "GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("gpt_neox_japanese", "GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("gptj", "GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("gptsan-japanese", "GPTSAN_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("graphormer", "GRAPHORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("groupvit", "GROUPVIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("hubert", "HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("ibert", "IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("idefics", "IDEFICS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("imagegpt", "IMAGEGPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("informer", "INFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("instructblip", "INSTRUCTBLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("jukebox", "JUKEBOX_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("kosmos-2", "KOSMOS2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("layoutlm", "LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("layoutlmv2", "LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("layoutlmv3", "LAYOUTLMV3_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("led", "LED_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("levit", "LEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("lilt", "LILT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("llama", "LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("llava", "LLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("longformer", "LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("longt5", "LONGT5_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("luke", "LUKE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("lxmert", "LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("m2m_100", "M2M_100_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mamba", "MAMBA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("markuplm", "MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mask2former", "MASK2FORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("maskformer", "MASKFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mbart", "MBART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mctct", "MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mega", "MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("megatron-bert", "MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mgp-str", "MGP_STR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mistral", "MISTRAL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mixtral", "MIXTRAL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mobilenet_v1", "MOBILENET_V1_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mobilenet_v2", "MOBILENET_V2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mobilevit", "MOBILEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mobilevitv2", "MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mpnet", "MPNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mpt", "MPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mra", "MRA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("musicgen", "MUSICGEN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("mvp", "MVP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("nat", "NAT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("nezha", "NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("nllb-moe", "NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("nystromformer", "NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("oneformer", "ONEFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("olmo", "OLMO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("open-llama", "OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("openai-gpt", "OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("opt", "OPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("owlv2", "OWLV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("owlvit", "OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("patchtsmixer", "PATCHTSMIXER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("patchtst", "PATCHTST_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("pegasus", "PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("pegasus_x", "PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("perceiver", "PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("persimmon", "PERSIMMON_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("phi", "PHI_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("pix2struct", "PIX2STRUCT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("plbart", "PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("poolformer", "POOLFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("pop2piano", "POP2PIANO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("prophetnet", "PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("pvt", "PVT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("qdqbert", "QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("qwen2", "QWEN2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("realm", "REALM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("regnet", "REGNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("rembert", "REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("resnet", "RESNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("retribert", "RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("roberta", "ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("roberta-prelayernorm", "ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("roc_bert", "ROC_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("roformer", "ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("rwkv", "RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("sam", "SAM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("seamless_m4t", "SEAMLESS_M4T_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("seamless_m4t_v2", "SEAMLESS_M4T_V2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("segformer", "SEGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("seggpt", "SEGGPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("sew", "SEW_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("sew-d", "SEW_D_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("siglip", "SIGLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("speech_to_text", "SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("speech_to_text_2", "SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("speecht5", "SPEECHT5_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("splinter", "SPLINTER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("squeezebert", "SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("stablelm", "STABLELM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("starcoder2", "STARCODER2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("swiftformer", "SWIFTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("swin", "SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("swin2sr", "SWIN2SR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("swinv2", "SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("switch_transformers", "SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("t5", "T5_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("table-transformer", "TABLE_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("tapas", "TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("time_series_transformer", "TIME_SERIES_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("timesformer", "TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("transfo-xl", "TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("tvlt", "TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("tvp", "TVP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("udop", "UDOP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("unispeech", "UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("unispeech-sat", "UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("univnet", "UNIVNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("van", "VAN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("videomae", "VIDEOMAE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vilt", "VILT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vipllava", "VIPLLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("visual_bert", "VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vit", "VIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vit_hybrid", "VIT_HYBRID_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vit_mae", "VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vit_msn", "VIT_MSN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vitdet", "VITDET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vitmatte", "VITMATTE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vits", "VITS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("vivit", "VIVIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("wav2vec2", "WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("wav2vec2-bert", "WAV2VEC2_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("wav2vec2-conformer", "WAV2VEC2_CONFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("whisper", "WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("xclip", "XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("xglm", "XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("xlm", "XLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("xlm-prophetnet", "XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("xlm-roberta", "XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("xlnet", "XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("xmod", "XMOD_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("yolos", "YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ("yoso", "YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
- ]
-)
diff --git a/transformers/models/deprecated/bort/__init__.py b/transformers/models/deprecated/bort/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/transformers/models/deprecated/bort/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/bort/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 8d58dd4836d8d36696b1961114f2409d9429490b..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/bort/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/bort/__pycache__/convert_bort_original_gluonnlp_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/deprecated/bort/__pycache__/convert_bort_original_gluonnlp_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 439eeb068829c2583b93b76785dc1e0ab0779be1..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/bort/__pycache__/convert_bort_original_gluonnlp_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/bort/convert_bort_original_gluonnlp_checkpoint_to_pytorch.py b/transformers/models/deprecated/bort/convert_bort_original_gluonnlp_checkpoint_to_pytorch.py
deleted file mode 100644
index 5dc9a244c43c78c58e5b1076cccb82847193301b..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/bort/convert_bort_original_gluonnlp_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,319 +0,0 @@
-# coding=utf-8
-# Copyright 2020, The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Bort checkpoint."""
-
-
-import argparse
-import os
-
-import gluonnlp as nlp
-import mxnet as mx
-import numpy as np
-import torch
-from gluonnlp.base import get_home_dir
-from gluonnlp.model.bert import BERTEncoder
-from gluonnlp.model.utils import _load_vocab
-from gluonnlp.vocab import Vocab
-from packaging import version
-from torch import nn
-
-from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
-from transformers.models.bert.modeling_bert import (
- BertIntermediate,
- BertLayer,
- BertOutput,
- BertSelfAttention,
- BertSelfOutput,
-)
-from transformers.utils import logging
-
-
-if version.parse(nlp.__version__) != version.parse("0.8.3"):
- raise Exception("requires gluonnlp == 0.8.3")
-
-if version.parse(mx.__version__) != version.parse("1.5.0"):
- raise Exception("requires mxnet == 1.5.0")
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-SAMPLE_TEXT = "The Nymphenburg Palace is a beautiful palace in Munich!"
-
-
-def convert_bort_checkpoint_to_pytorch(bort_checkpoint_path: str, pytorch_dump_folder_path: str):
- """
- Convert the original Bort checkpoint (based on MXNET and Gluonnlp) to our BERT structure-
- """
-
- # Original Bort configuration
- bort_4_8_768_1024_hparams = {
- "attention_cell": "multi_head",
- "num_layers": 4,
- "units": 1024,
- "hidden_size": 768,
- "max_length": 512,
- "num_heads": 8,
- "scaled": True,
- "dropout": 0.1,
- "use_residual": True,
- "embed_size": 1024,
- "embed_dropout": 0.1,
- "word_embed": None,
- "layer_norm_eps": 1e-5,
- "token_type_vocab_size": 2,
- }
-
- predefined_args = bort_4_8_768_1024_hparams
-
- # Let's construct the original Bort model here
- # Taken from official BERT implementation, see:
- # https://github.com/alexa/bort/blob/master/bort/bort.py
- encoder = BERTEncoder(
- attention_cell=predefined_args["attention_cell"],
- num_layers=predefined_args["num_layers"],
- units=predefined_args["units"],
- hidden_size=predefined_args["hidden_size"],
- max_length=predefined_args["max_length"],
- num_heads=predefined_args["num_heads"],
- scaled=predefined_args["scaled"],
- dropout=predefined_args["dropout"],
- output_attention=False,
- output_all_encodings=False,
- use_residual=predefined_args["use_residual"],
- activation=predefined_args.get("activation", "gelu"),
- layer_norm_eps=predefined_args.get("layer_norm_eps", None),
- )
-
- # Vocab information needs to be fetched first
- # It's the same as RoBERTa, so RobertaTokenizer can be used later
- vocab_name = "openwebtext_ccnews_stories_books_cased"
-
- # Specify download folder to Gluonnlp's vocab
- gluon_cache_dir = os.path.join(get_home_dir(), "models")
- bort_vocab = _load_vocab(vocab_name, None, gluon_cache_dir, cls=Vocab)
-
- original_bort = nlp.model.BERTModel(
- encoder,
- len(bort_vocab),
- units=predefined_args["units"],
- embed_size=predefined_args["embed_size"],
- embed_dropout=predefined_args["embed_dropout"],
- word_embed=predefined_args["word_embed"],
- use_pooler=False,
- use_token_type_embed=False,
- token_type_vocab_size=predefined_args["token_type_vocab_size"],
- use_classifier=False,
- use_decoder=False,
- )
-
- original_bort.load_parameters(bort_checkpoint_path, cast_dtype=True, ignore_extra=True)
- params = original_bort._collect_params_with_prefix()
-
- # Build our config 🤗
- hf_bort_config_json = {
- "architectures": ["BertForMaskedLM"],
- "attention_probs_dropout_prob": predefined_args["dropout"],
- "hidden_act": "gelu",
- "hidden_dropout_prob": predefined_args["dropout"],
- "hidden_size": predefined_args["embed_size"],
- "initializer_range": 0.02,
- "intermediate_size": predefined_args["hidden_size"],
- "layer_norm_eps": predefined_args["layer_norm_eps"],
- "max_position_embeddings": predefined_args["max_length"],
- "model_type": "bort",
- "num_attention_heads": predefined_args["num_heads"],
- "num_hidden_layers": predefined_args["num_layers"],
- "pad_token_id": 1, # 2 = BERT, 1 = RoBERTa
- "type_vocab_size": 1, # 2 = BERT, 1 = RoBERTa
- "vocab_size": len(bort_vocab),
- }
-
- hf_bort_config = BertConfig.from_dict(hf_bort_config_json)
- hf_bort_model = BertForMaskedLM(hf_bort_config)
- hf_bort_model.eval()
-
- # Parameter mapping table (Gluonnlp to Transformers)
- # * denotes layer index
- #
- # | Gluon Parameter | Transformers Parameter
- # | -------------------------------------------------------------- | ----------------------
- # | `encoder.layer_norm.beta` | `bert.embeddings.LayerNorm.bias`
- # | `encoder.layer_norm.gamma` | `bert.embeddings.LayerNorm.weight`
- # | `encoder.position_weight` | `bert.embeddings.position_embeddings.weight`
- # | `word_embed.0.weight` | `bert.embeddings.word_embeddings.weight`
- # | `encoder.transformer_cells.*.attention_cell.proj_key.bias` | `bert.encoder.layer.*.attention.self.key.bias`
- # | `encoder.transformer_cells.*.attention_cell.proj_key.weight` | `bert.encoder.layer.*.attention.self.key.weight`
- # | `encoder.transformer_cells.*.attention_cell.proj_query.bias` | `bert.encoder.layer.*.attention.self.query.bias`
- # | `encoder.transformer_cells.*.attention_cell.proj_query.weight` | `bert.encoder.layer.*.attention.self.query.weight`
- # | `encoder.transformer_cells.*.attention_cell.proj_value.bias` | `bert.encoder.layer.*.attention.self.value.bias`
- # | `encoder.transformer_cells.*.attention_cell.proj_value.weight` | `bert.encoder.layer.*.attention.self.value.weight`
- # | `encoder.transformer_cells.*.ffn.ffn_2.bias` | `bert.encoder.layer.*.attention.output.dense.bias`
- # | `encoder.transformer_cells.*.ffn.ffn_2.weight` | `bert.encoder.layer.*.attention.output.dense.weight`
- # | `encoder.transformer_cells.*.layer_norm.beta` | `bert.encoder.layer.*.attention.output.LayerNorm.bias`
- # | `encoder.transformer_cells.*.layer_norm.gamma` | `bert.encoder.layer.*.attention.output.LayerNorm.weight`
- # | `encoder.transformer_cells.*.ffn.ffn_1.bias` | `bert.encoder.layer.*.intermediate.dense.bias`
- # | `encoder.transformer_cells.*.ffn.ffn_1.weight` | `bert.encoder.layer.*.intermediate.dense.weight`
- # | `encoder.transformer_cells.*.ffn.layer_norm.beta` | `bert.encoder.layer.*.output.LayerNorm.bias`
- # | `encoder.transformer_cells.*.ffn.layer_norm.gamma` | `bert.encoder.layer.*.output.LayerNorm.weight`
- # | `encoder.transformer_cells.*.proj.bias` | `bert.encoder.layer.*.output.dense.bias`
- # | `encoder.transformer_cells.*.proj.weight` | `bert.encoder.layer.*.output.dense.weight`
-
- # Helper function to convert MXNET Arrays to PyTorch
- def to_torch(mx_array) -> nn.Parameter:
- return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()))
-
- # Check param shapes and map new HF param back
- def check_and_map_params(hf_param, gluon_param):
- shape_hf = hf_param.shape
-
- gluon_param = to_torch(params[gluon_param])
- shape_gluon = gluon_param.shape
-
- assert (
- shape_hf == shape_gluon
- ), f"The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers"
-
- return gluon_param
-
- hf_bort_model.bert.embeddings.word_embeddings.weight = check_and_map_params(
- hf_bort_model.bert.embeddings.word_embeddings.weight, "word_embed.0.weight"
- )
- hf_bort_model.bert.embeddings.position_embeddings.weight = check_and_map_params(
- hf_bort_model.bert.embeddings.position_embeddings.weight, "encoder.position_weight"
- )
- hf_bort_model.bert.embeddings.LayerNorm.bias = check_and_map_params(
- hf_bort_model.bert.embeddings.LayerNorm.bias, "encoder.layer_norm.beta"
- )
- hf_bort_model.bert.embeddings.LayerNorm.weight = check_and_map_params(
- hf_bort_model.bert.embeddings.LayerNorm.weight, "encoder.layer_norm.gamma"
- )
-
- # Inspired by RoBERTa conversion script, we just zero them out (Bort does not use them)
- hf_bort_model.bert.embeddings.token_type_embeddings.weight.data = torch.zeros_like(
- hf_bort_model.bert.embeddings.token_type_embeddings.weight.data
- )
-
- for i in range(hf_bort_config.num_hidden_layers):
- layer: BertLayer = hf_bort_model.bert.encoder.layer[i]
-
- # self attention
- self_attn: BertSelfAttention = layer.attention.self
-
- self_attn.key.bias.data = check_and_map_params(
- self_attn.key.bias.data, f"encoder.transformer_cells.{i}.attention_cell.proj_key.bias"
- )
-
- self_attn.key.weight.data = check_and_map_params(
- self_attn.key.weight.data, f"encoder.transformer_cells.{i}.attention_cell.proj_key.weight"
- )
- self_attn.query.bias.data = check_and_map_params(
- self_attn.query.bias.data, f"encoder.transformer_cells.{i}.attention_cell.proj_query.bias"
- )
- self_attn.query.weight.data = check_and_map_params(
- self_attn.query.weight.data, f"encoder.transformer_cells.{i}.attention_cell.proj_query.weight"
- )
- self_attn.value.bias.data = check_and_map_params(
- self_attn.value.bias.data, f"encoder.transformer_cells.{i}.attention_cell.proj_value.bias"
- )
- self_attn.value.weight.data = check_and_map_params(
- self_attn.value.weight.data, f"encoder.transformer_cells.{i}.attention_cell.proj_value.weight"
- )
-
- # self attention output
- self_output: BertSelfOutput = layer.attention.output
-
- self_output.dense.bias = check_and_map_params(
- self_output.dense.bias, f"encoder.transformer_cells.{i}.proj.bias"
- )
- self_output.dense.weight = check_and_map_params(
- self_output.dense.weight, f"encoder.transformer_cells.{i}.proj.weight"
- )
- self_output.LayerNorm.bias = check_and_map_params(
- self_output.LayerNorm.bias, f"encoder.transformer_cells.{i}.layer_norm.beta"
- )
- self_output.LayerNorm.weight = check_and_map_params(
- self_output.LayerNorm.weight, f"encoder.transformer_cells.{i}.layer_norm.gamma"
- )
-
- # intermediate
- intermediate: BertIntermediate = layer.intermediate
-
- intermediate.dense.bias = check_and_map_params(
- intermediate.dense.bias, f"encoder.transformer_cells.{i}.ffn.ffn_1.bias"
- )
- intermediate.dense.weight = check_and_map_params(
- intermediate.dense.weight, f"encoder.transformer_cells.{i}.ffn.ffn_1.weight"
- )
-
- # output
- bert_output: BertOutput = layer.output
-
- bert_output.dense.bias = check_and_map_params(
- bert_output.dense.bias, f"encoder.transformer_cells.{i}.ffn.ffn_2.bias"
- )
- bert_output.dense.weight = check_and_map_params(
- bert_output.dense.weight, f"encoder.transformer_cells.{i}.ffn.ffn_2.weight"
- )
- bert_output.LayerNorm.bias = check_and_map_params(
- bert_output.LayerNorm.bias, f"encoder.transformer_cells.{i}.ffn.layer_norm.beta"
- )
- bert_output.LayerNorm.weight = check_and_map_params(
- bert_output.LayerNorm.weight, f"encoder.transformer_cells.{i}.ffn.layer_norm.gamma"
- )
-
- # Save space and energy 🎄
- hf_bort_model.half()
-
- # Compare output of both models
- tokenizer = RobertaTokenizer.from_pretrained("FacebookAI/roberta-base")
-
- input_ids = tokenizer.encode_plus(SAMPLE_TEXT)["input_ids"]
-
- # Get gluon output
- gluon_input_ids = mx.nd.array([input_ids])
- output_gluon = original_bort(inputs=gluon_input_ids, token_types=[])
-
- # Get Transformer output (save and reload model again)
- hf_bort_model.save_pretrained(pytorch_dump_folder_path)
- hf_bort_model = BertModel.from_pretrained(pytorch_dump_folder_path)
- hf_bort_model.eval()
-
- input_ids = tokenizer.encode_plus(SAMPLE_TEXT, return_tensors="pt")
- output_hf = hf_bort_model(**input_ids)[0]
-
- gluon_layer = output_gluon[0].asnumpy()
- hf_layer = output_hf[0].detach().numpy()
-
- max_absolute_diff = np.max(np.abs(hf_layer - gluon_layer)).item()
- success = np.allclose(gluon_layer, hf_layer, atol=1e-3)
-
- if success:
- print("✔️ Both model do output the same tensors")
- else:
- print("❌ Both model do **NOT** output the same tensors")
- print("Absolute difference is:", max_absolute_diff)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--bort_checkpoint_path", default=None, type=str, required=True, help="Path the official Bort params file."
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
- args = parser.parse_args()
- convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)
diff --git a/transformers/models/deprecated/mctct/__init__.py b/transformers/models/deprecated/mctct/__init__.py
deleted file mode 100644
index 567be97b7cd8631e71367e713dc2f0ef23bd76f5..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/mctct/__init__.py
+++ /dev/null
@@ -1,56 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {
- "configuration_mctct": ["MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP", "MCTCTConfig"],
- "feature_extraction_mctct": ["MCTCTFeatureExtractor"],
- "processing_mctct": ["MCTCTProcessor"],
-}
-
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_mctct"] = [
- "MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "MCTCTForCTC",
- "MCTCTModel",
- "MCTCTPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_mctct import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP, MCTCTConfig
- from .feature_extraction_mctct import MCTCTFeatureExtractor
- from .processing_mctct import MCTCTProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_mctct import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST, MCTCTForCTC, MCTCTModel, MCTCTPreTrainedModel
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deprecated/mctct/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/mctct/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 282910296e2dc5cebd26ccc1b336dadae443dbbb..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/mctct/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/mctct/__pycache__/configuration_mctct.cpython-310.pyc b/transformers/models/deprecated/mctct/__pycache__/configuration_mctct.cpython-310.pyc
deleted file mode 100644
index 070efbaa319e77523de73abc12c969e6be6bd9b1..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/mctct/__pycache__/configuration_mctct.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/mctct/__pycache__/feature_extraction_mctct.cpython-310.pyc b/transformers/models/deprecated/mctct/__pycache__/feature_extraction_mctct.cpython-310.pyc
deleted file mode 100644
index 0322536368955c0c20c7c574f7896c2880d27943..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/mctct/__pycache__/feature_extraction_mctct.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/mctct/__pycache__/modeling_mctct.cpython-310.pyc b/transformers/models/deprecated/mctct/__pycache__/modeling_mctct.cpython-310.pyc
deleted file mode 100644
index eafda60201d786efb961a00826435adaf2af1efc..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/mctct/__pycache__/modeling_mctct.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/mctct/__pycache__/processing_mctct.cpython-310.pyc b/transformers/models/deprecated/mctct/__pycache__/processing_mctct.cpython-310.pyc
deleted file mode 100644
index b11428e3453e670521288034a4da4daacfab2a20..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/mctct/__pycache__/processing_mctct.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/mctct/configuration_mctct.py b/transformers/models/deprecated/mctct/configuration_mctct.py
deleted file mode 100644
index 6546b18eab0522ce8c618da7bf3a8baab005dd0d..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/mctct/configuration_mctct.py
+++ /dev/null
@@ -1,184 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""M-CTC-T model configuration"""
-
-from ....configuration_utils import PretrainedConfig
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from .._archive_maps import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class MCTCTConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`MCTCTModel`]. It is used to instantiate an
- M-CTC-T model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the M-CTC-T
- [speechbrain/m-ctc-t-large](https://huggingface.co/speechbrain/m-ctc-t-large) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 8065):
- Vocabulary size of the M-CTC-T model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`MCTCTModel`].
- hidden_size (`int`, *optional*, defaults to 1536):
- Dimension of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 36):
- Number of hidden layers in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 6144):
- Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 4):
- Number of attention heads for each attention layer in the Transformer encoder.
- attention_head_dim (`int`, *optional*, defaults to 384):
- Dimensions of each attention head for each attention layer in the Transformer encoder.
- max_position_embeddings (`int`, *optional*, defaults to 920):
- The maximum sequence length that this model might ever be used with (after log-mel spectrogram extraction).
- layer_norm_eps (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the layer normalization layers.
- layerdrop (`float`, *optional*, defaults to 0.3):
- The probability of dropping an encoder layer during training. The default 0.3 value is used in the original
- implementation.
- hidden_act (`str` or `function`, *optional*, defaults to `"relu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.3):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.3):
- The dropout ratio for the attention probabilities.
- pad_token_id (`int`, *optional*, defaults to 1):
- The tokenizer index of the pad token.
- bos_token_id (`int`, *optional*, defaults to 0):
- The tokenizer index of the bos token.
- eos_token_id (`int`, *optional*, defaults to 2):
- The tokenizer index of the eos token.
- conv_glu_dim (`int`, *optional*, defaults to 1):
- The dimension of the output of the `Conv1dSubsampler` layer in which GLU is applied on. Though the original
- Flashlight code uses the value of 2, here it's adapted to 1 due to transposition differences.
- conv_dropout (`int`, *optional*, defaults to 0.3):
- The probability of randomly dropping the `Conv1dSubsampler` layer during training.
- num_conv_layers (`int`, *optional*, defaults to 1):
- Number of convolution layers before applying transformer encoder layers.
- conv_kernel (`Sequence[int]`, *optional*, defaults to `(7,)`):
- The kernel size of the 1D convolution applied before transformer layers. `len(conv_kernel)` must be equal
- to `num_conv_layers`.
- conv_stride (`Sequence[int]`, *optional*, defaults to `(3,)`):
- The stride length of the 1D convolution applied before transformer layers. `len(conv_stride)` must be equal
- to `num_conv_layers`.
- input_feat_per_channel (`int`, *optional*, defaults to 80):
- Feature dimensions of the channels of the input to the Conv1D layer.
- input_channels (`int`, *optional*, defaults to 1):
- Number of input channels of the input to the Conv1D layer.
- conv_channels (`List[int]`, *optional*):
- Channel sizes of intermediate Conv1D layers.
- ctc_loss_reduction (`str`, *optional*, defaults to `"sum"`):
- Specifies the reduction to apply to the output of `torch.nn.CTCLoss`. Only relevant when training an
- instance of [`MCTCTForCTC`].
- ctc_zero_infinity (`bool`, *optional*, defaults to `False`):
- Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly
- occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance
- of [`MCTCTForCTC`].
-
- Example:
-
- ```python
- >>> from transformers import MCTCTConfig, MCTCTModel
-
- >>> # Initializing a M-CTC-T mctct-large style configuration
- >>> configuration = MCTCTConfig()
-
- >>> # Initializing a model (with random weights) from the mctct-large style configuration
- >>> model = MCTCTModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "mctct"
-
- def __init__(
- self,
- vocab_size=8065,
- hidden_size=1536,
- num_hidden_layers=36,
- intermediate_size=6144,
- num_attention_heads=4,
- attention_head_dim=384,
- max_position_embeddings=920,
- layer_norm_eps=1e-5,
- layerdrop=0.3,
- hidden_act="relu",
- initializer_range=0.02,
- hidden_dropout_prob=0.3,
- attention_probs_dropout_prob=0.3,
- pad_token_id=1,
- bos_token_id=0,
- eos_token_id=2,
- conv_glu_dim=1,
- conv_dropout=0.3,
- num_conv_layers=1,
- conv_kernel=(7,),
- conv_stride=(3,),
- input_feat_per_channel=80,
- input_channels=1,
- conv_channels=None,
- ctc_loss_reduction="sum",
- ctc_zero_infinity=False,
- **kwargs,
- ):
- super().__init__(**kwargs, pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id)
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.intermediate_size = intermediate_size
- self.num_attention_heads = num_attention_heads
- self.attention_head_dim = attention_head_dim
- self.max_position_embeddings = max_position_embeddings
- self.layer_norm_eps = layer_norm_eps
- self.layerdrop = layerdrop
- self.hidden_act = hidden_act
- self.initializer_range = initializer_range
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.pad_token_id = pad_token_id
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
- self.conv_glu_dim = conv_glu_dim
- self.conv_dropout = conv_dropout
- self.num_conv_layers = num_conv_layers
- self.input_feat_per_channel = input_feat_per_channel
- self.input_channels = input_channels
- self.conv_channels = conv_channels
- self.ctc_loss_reduction = ctc_loss_reduction
- self.ctc_zero_infinity = ctc_zero_infinity
-
- # prevents config testing fail with exporting to json
- self.conv_kernel = list(conv_kernel)
- self.conv_stride = list(conv_stride)
-
- if len(self.conv_kernel) != self.num_conv_layers:
- raise ValueError(
- "Configuration for convolutional module is incorrect. "
- "It is required that `len(config.conv_kernel)` == `config.num_conv_layers` "
- f"but is `len(config.conv_kernel) = {len(self.conv_kernel)}`, "
- f"`config.num_conv_layers = {self.num_conv_layers}`."
- )
diff --git a/transformers/models/deprecated/mctct/feature_extraction_mctct.py b/transformers/models/deprecated/mctct/feature_extraction_mctct.py
deleted file mode 100644
index e1e17c4b12f91dc25284e30a70388137e52ab82b..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/mctct/feature_extraction_mctct.py
+++ /dev/null
@@ -1,288 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Feature extractor class for M-CTC-T
-"""
-
-from typing import List, Optional, Union
-
-import numpy as np
-
-from ....audio_utils import mel_filter_bank, optimal_fft_length, spectrogram, window_function
-from ....feature_extraction_sequence_utils import SequenceFeatureExtractor
-from ....feature_extraction_utils import BatchFeature
-from ....file_utils import PaddingStrategy, TensorType
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class MCTCTFeatureExtractor(SequenceFeatureExtractor):
- r"""
- Constructs a M-CTC-T feature extractor.
-
- This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains
- most of the main methods. Users should refer to this superclass for more information regarding those methods. This
- code has been adapted from Flashlight's C++ code. For more information about the implementation, one can refer to
- this [notebook](https://colab.research.google.com/drive/1GLtINkkhzms-IsdcGy_-tVCkv0qNF-Gt#scrollTo=pMCRGMmUC_an)
- that takes the user step-by-step in the implementation.
-
- Args:
- feature_size (`int`, defaults to 80):
- The feature dimension of the extracted features. This is the number of mel_frequency
- sampling_rate (`int`, defaults to 16000):
- The sampling rate at which the audio files should be digitalized expressed in hertz (Hz).
- padding_value (`float`, defaults to 0.0):
- The value that is used to fill the padding values.
- hop_length (`int`, defaults to 10):
- Number of audio samples between windows. Otherwise referred to as "shift" in many papers.
- win_length (`int`, defaults to 25):
- Number of ms per window
- win_function (`str`, defaults to `"hamming_window"`):
- Name for the window function used for windowing, must be accessible via `torch.{win_function}`
- frame_signal_scale (`float`, defaults to 32768.0):
- Constant multiplied in creating the frames before applying DFT.
- preemphasis_coeff (`float`, defaults to 0.97):
- Constant multiplied in applying Pre-emphasis before DFT.
- mel_floor (`float` defaults to 1.0):
- Minimum value of mel frequency banks.
- normalize_means (`bool`, *optional*, defaults to `True`):
- Whether or not to zero-mean normalize the extracted features.
- normalize_vars (`bool`, *optional*, defaults to `True`):
- Whether or not to unit-variance normalize the extracted features.
- """
-
- model_input_names = ["input_features", "attention_mask"]
-
- def __init__(
- self,
- feature_size=80,
- sampling_rate=16000,
- padding_value=0.0,
- hop_length=10,
- win_length=25,
- win_function="hamming_window",
- frame_signal_scale=32768.0,
- preemphasis_coeff=0.97,
- mel_floor=1.0,
- normalize_means=True,
- normalize_vars=True,
- return_attention_mask=False,
- **kwargs,
- ):
- super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
-
- self.feature_size = feature_size
- self.sampling_rate = sampling_rate
- self.padding_value = padding_value
- self.hop_length = hop_length
- self.win_length = win_length
- self.frame_signal_scale = frame_signal_scale
- self.preemphasis_coeff = preemphasis_coeff
- self.mel_floor = mel_floor
- self.normalize_means = normalize_means
- self.normalize_vars = normalize_vars
- self.win_function = win_function
- self.return_attention_mask = return_attention_mask
-
- self.sample_size = win_length * sampling_rate // 1000
- self.sample_stride = hop_length * sampling_rate // 1000
-
- self.n_fft = optimal_fft_length(self.sample_size)
- self.n_freqs = (self.n_fft // 2) + 1
-
- def _extract_mfsc_features(self, one_waveform: np.array) -> np.ndarray:
- """
- Extracts MFSC Features for one waveform vector (unbatched). Adapted from Flashlight's C++ MFSC code.
- """
- if self.win_function == "hamming_window":
- window = window_function(window_length=self.sample_size, name=self.win_function, periodic=False)
- else:
- window = window_function(window_length=self.sample_size, name=self.win_function)
-
- fbanks = mel_filter_bank(
- num_frequency_bins=self.n_freqs,
- num_mel_filters=self.feature_size,
- min_frequency=0.0,
- max_frequency=self.sampling_rate / 2.0,
- sampling_rate=self.sampling_rate,
- )
-
- msfc_features = spectrogram(
- one_waveform * self.frame_signal_scale,
- window=window,
- frame_length=self.sample_size,
- hop_length=self.sample_stride,
- fft_length=self.n_fft,
- center=False,
- preemphasis=self.preemphasis_coeff,
- mel_filters=fbanks,
- mel_floor=self.mel_floor,
- log_mel="log",
- )
- return msfc_features.T
-
- def _normalize_one(self, x, input_length, padding_value):
- # make sure we normalize float32 arrays
- if self.normalize_means:
- mean = x[:input_length].mean(axis=0)
- x = np.subtract(x, mean)
- if self.normalize_vars:
- std = x[:input_length].std(axis=0)
- x = np.divide(x, std)
-
- if input_length < x.shape[0]:
- x[input_length:] = padding_value
-
- # make sure array is in float32
- x = x.astype(np.float32)
-
- return x
-
- def normalize(
- self, input_features: List[np.ndarray], attention_mask: Optional[np.ndarray] = None
- ) -> List[np.ndarray]:
- lengths = attention_mask.sum(-1) if attention_mask is not None else [x.shape[0] for x in input_features]
- return [self._normalize_one(x, n, self.padding_value) for x, n in zip(input_features, lengths)]
-
- def __call__(
- self,
- raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],
- padding: Union[bool, str, PaddingStrategy] = False,
- max_length: Optional[int] = None,
- truncation: bool = False,
- pad_to_multiple_of: Optional[int] = None,
- return_attention_mask: Optional[bool] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- sampling_rate: Optional[int] = None,
- **kwargs,
- ) -> BatchFeature:
- """
- Main method to featurize and prepare for the model one or several sequence(s). sequences. It returns the
- log-mel spectrogram of the input audio, as implemented in the original Flashlight MFSC feature extraction code.
-
- Args:
- raw_speech (`torch.Tensor`, `np.ndarray`, `List[float]`, `List[torch.Tensor]`, `List[np.ndarray]`, `List[List[float]]`):
- The sequence or batch of sequences to be padded. Each sequence can be a tensor, a numpy array, a list
- of float values, a list of tensors, a list of numpy arrays or a list of list of float values. Must be
- mono channel audio, not stereo, i.e. single float per timestep.
- padding (`bool`, `str` or [`~file_utils.PaddingStrategy`], *optional*, defaults to `False`):
- Select a strategy to pad the returned sequences (according to the model's padding side and padding
- index) among:
-
- - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
- sequence if provided).
- - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided.
- - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
- lengths).
- max_length (`int`, *optional*):
- Maximum length of the returned list and optionally padding length (see above).
- truncation (`bool`):
- Activates truncation to cut input sequences longer than *max_length* to *max_length*.
- pad_to_multiple_of (`int`, *optional*):
- If set will pad the sequence to a multiple of the provided value.
-
- This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
- `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128.
- return_attention_mask (`bool`, *optional*):
- Whether to return the attention mask. If left to the default, will return the attention mask according
- to the specific feature_extractor's default.
-
- [What are attention masks?](../glossary#attention-mask)
-
- return_tensors (`str` or [`~file_utils.TensorType`], *optional*):
- If set, will return tensors instead of list of python integers. Acceptable values are:
-
- - `'tf'`: Return TensorFlow `tf.constant` objects.
- - `'pt'`: Return PyTorch `torch.Tensor` objects.
- - `'np'`: Return Numpy `np.ndarray` objects.
- sampling_rate (`int`, *optional*):
- The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass
- `sampling_rate` at the forward call to prevent silent errors.
- padding_value (`float`, defaults to 0.0):
- """
-
- if sampling_rate is not None:
- if sampling_rate != self.sampling_rate:
- raise ValueError(
- f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of"
- f" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with"
- f" {self.sampling_rate} and not {sampling_rate}."
- )
- else:
- logger.warning(
- "It is strongly recommended to pass the ``sampling_rate`` argument to this function. "
- "Failing to do so can result in silent errors that might be hard to debug."
- )
-
- is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
- if is_batched_numpy and len(raw_speech.shape) > 2:
- raise ValueError(f"Only mono-channel audio is supported for input to {self}")
- is_batched = is_batched_numpy or (
- isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
- )
-
- if is_batched:
- raw_speech = [np.asarray(speech, dtype=np.float32) for speech in raw_speech]
- elif not is_batched and not isinstance(raw_speech, np.ndarray):
- raw_speech = np.asarray(raw_speech, dtype=np.float32)
- elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):
- raw_speech = raw_speech.astype(np.float32)
-
- # always return batch
- if not is_batched:
- raw_speech = [raw_speech]
-
- # extract fbank features
- features = [self._extract_mfsc_features(one_waveform) for one_waveform in raw_speech]
-
- # convert into correct format for padding
- encoded_inputs = BatchFeature({"input_features": features})
-
- padded_inputs = self.pad(
- encoded_inputs,
- padding=padding,
- max_length=max_length,
- truncation=truncation,
- pad_to_multiple_of=pad_to_multiple_of,
- return_attention_mask=True,
- **kwargs,
- )
- # make sure list is in array format
- input_features = padded_inputs.get("input_features")
- if isinstance(input_features[0], list):
- padded_inputs["input_features"] = [np.asarray(feature, dtype=np.float32) for feature in input_features]
-
- attention_mask = padded_inputs.get("attention_mask")
- if attention_mask is not None:
- padded_inputs["attention_mask"] = [np.asarray(array, dtype=np.int32) for array in attention_mask]
-
- if self.normalize_means or self.normalize_vars:
- attention_mask = (
- np.array(attention_mask, dtype=np.int32)
- if self._get_padding_strategies(padding, max_length=max_length) is not PaddingStrategy.DO_NOT_PAD
- and padding
- else None
- )
- padded_inputs["input_features"] = self.normalize(
- padded_inputs["input_features"], attention_mask=attention_mask
- )
-
- if return_tensors is not None:
- padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
-
- return padded_inputs
diff --git a/transformers/models/deprecated/mctct/modeling_mctct.py b/transformers/models/deprecated/mctct/modeling_mctct.py
deleted file mode 100644
index 2d9ef6cf724c28a708de4be12161caabb663da12..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/mctct/modeling_mctct.py
+++ /dev/null
@@ -1,792 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch M-CTC-T model."""
-
-
-import math
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from ....activations import ACT2FN
-from ....file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
-from ....integrations.deepspeed import is_deepspeed_zero3_enabled
-from ....modeling_attn_mask_utils import _prepare_4d_attention_mask
-from ....modeling_outputs import BaseModelOutput, CausalLMOutput
-from ....modeling_utils import (
- PreTrainedModel,
- apply_chunking_to_forward,
- find_pruneable_heads_and_indices,
- prune_linear_layer,
-)
-from ....utils import logging
-from .configuration_mctct import MCTCTConfig
-
-
-logger = logging.get_logger(__name__)
-
-_HIDDEN_STATES_START_POSITION = 1
-
-_CONFIG_FOR_DOC = "MCTCTConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "speechbrain/m-ctc-t-large"
-_EXPECTED_OUTPUT_SHAPE = [1, 195, 1536]
-
-# CTC docstring
-_CTC_EXPECTED_OUTPUT = '"Mr. Quilter is the apostle of the middle classes, and we\'re glad to welcome his gospel."'
-_CTC_EXPECTED_LOSS = 1885.65
-
-
-from .._archive_maps import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class MCTCTConv1dSubsampler(nn.Module):
- """
- Convolutional subsampler: a stack of 1D convolution (along temporal dimension) followed by non-linear activation
- via gated linear units (https://arxiv.org/abs/1911.08460)
- """
-
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.glu_dim = config.conv_glu_dim
-
- self.dropout = nn.Dropout(config.conv_dropout)
-
- self.num_layers = config.num_conv_layers
- self.in_channels = config.input_feat_per_channel * config.input_channels
-
- if self.num_layers > 1:
- if config.conv_channels is None:
- raise ValueError(
- "Need to specify `conv_channels` configuration in `MCTCTConfig` to use multiple convolution"
- " layers."
- )
-
- self.mid_channels = config.conv_channels
- else:
- self.mid_channels = None
-
- self.out_channels = config.hidden_size * 2 # considering GLU halving
- self.kernel_size = config.conv_kernel
- self.stride = config.conv_stride
-
- # NOTE: MCTCT by construction only uses one convolution kernel. I've made this flexible to allow for
- # multiple layers of convolutions, but not sure if this model definition should just restrict it
- # to one layer. This becomes especially relevant when considering the padding like line 1 of forward().
- self.conv_layers = nn.ModuleList(
- nn.Conv1d(
- self.in_channels if i == 0 else self.mid_channels[i],
- self.mid_channels[i] if i < self.num_layers - 1 else self.out_channels,
- kernel_size=k,
- stride=self.stride[i],
- padding="valid",
- )
- for i, k in enumerate(self.kernel_size)
- )
-
- def forward(self, input_features):
- # NOTE: in reference to the NOTE in __init__, right now it just calculates padding as if
- # there will be just one conv layer.
- padding = sum([size // 2 for size in self.kernel_size]) # (7, 7) -> (3, 3)
-
- input_features = torch.nn.functional.pad(input_features, (0, 0, padding, padding), "constant", 0)
- hidden_states = input_features.transpose(1, 2).contiguous() # -> Batch x Frame x Time
- for conv in self.conv_layers:
- hidden_states = conv(hidden_states)
- hidden_states = nn.functional.glu(hidden_states, dim=self.glu_dim)
- hidden_states = self.dropout(hidden_states)
-
- hidden_states = hidden_states.transpose(1, 2).contiguous() # -> Batch x Time x Frame
- return hidden_states
-
-
-class MCTCTEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
-
- # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
- # any TensorFlow checkpoint file
- # self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.LayerNorm = MCTCTLayerNorm()
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
- self.register_buffer(
- "token_type_ids",
- torch.zeros(self.position_ids.size(), dtype=torch.long, device=self.position_ids.device),
- persistent=False,
- )
-
- def forward(
- self, input_features=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
- ):
- input_shape = input_features.size() if input_features is not None else inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
-
- # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
- # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
- # issue #5664
- if token_type_ids is None:
- if hasattr(self, "token_type_ids"):
- buffered_token_type_ids = self.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_features)
-
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
- embeddings = inputs_embeds + token_type_embeddings
-
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-class MCTCTSelfAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = config.attention_head_dim
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=False)
- self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=False)
- self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=False)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- self.max_position_embeddings = config.max_position_embeddings
- self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
-
- self.is_decoder = config.is_decoder
-
- def transpose_for_scores(self, x):
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(*new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def reshape_fortran(self, x, shape):
- if len(x.shape) > 0:
- x = x.permute(*reversed(range(len(x.shape))))
- return x.reshape(*reversed(shape)).permute(*reversed(range(len(shape))))
-
- def relative_position_embedding_rotate(self, scores):
- # NOTE: should re-evaluate whether this re-implementation was truly necessary
- # or the reason why my complete re-haul worked was due to some other part
- # of the code. Adding this and the reshape fortrain code seems very undesirable.
- scores = scores.permute(0, 2, 3, 1) # e.g. [10, 1839, 14, 4]
-
- batch, hidden_state, seq_len, heads = scores.shape
-
- # e.g. [10, 1853, 14, 4]
- scores = torch.cat((scores, torch.zeros((batch, seq_len, seq_len, heads), device=scores.device)), dim=1)
-
- # e.g. [10, 25942, 1, 4]
- scores = self.reshape_fortran(scores, [batch, (hidden_state + seq_len) * seq_len, 1, heads])
-
- # e.g. [10, 25928, 1, 4]
- scores = scores[:, : (seq_len + hidden_state - 1) * seq_len]
-
- # e.g. [10, 1852, 14, 4]
- scores = self.reshape_fortran(scores, [batch, hidden_state + seq_len - 1, seq_len, heads])
-
- halfpoint = hidden_state // 2
- scores = scores[:, halfpoint : halfpoint + seq_len].transpose(1, 2) # e.g. [10, 14, 14, 4]
-
- return scores.permute(0, 3, 1, 2)
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- output_attentions=False,
- ):
- mixed_query_layer = self.query(hidden_states)
- mixed_query_layer = mixed_query_layer / math.sqrt(self.attention_head_size)
-
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- # relative key position embeddings
- positional_embedding = self.distance_embedding.weight
- relative_position_scores = torch.einsum("lh, bche -> bcle", positional_embedding, query_layer.transpose(2, 3))
-
- relative_position_scores = self.relative_position_embedding_rotate(relative_position_scores)
- attention_scores = attention_scores + relative_position_scores
-
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in MCTCTModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).flatten(start_dim=-2)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-class MCTCTLayerNorm(nn.Module):
- def __init__(self):
- super().__init__()
- self.singleton_weight = nn.Parameter(torch.ones(1))
- self.singleton_bias = nn.Parameter(torch.zeros(1))
-
- def forward(self, hidden_states):
- return (hidden_states * self.singleton_weight) + self.singleton_bias
-
-
-class MCTCTSelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class MCTCTAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.self = MCTCTSelfAttention(config)
- self.output = MCTCTSelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- output_attentions=False,
- ):
- self_outputs = self.self(
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions,
- )
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
-
- return outputs
-
-
-class MCTCTIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-class MCTCTOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class MCTCTLayer(nn.Module):
- def __init__(self, config: MCTCTConfig):
- super().__init__()
-
- self.seq_len_dim = 1
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
-
- self.intermediate = MCTCTIntermediate(config)
- self.attention = MCTCTAttention(config)
- self.is_decoder = config.is_decoder
- self.output = MCTCTOutput(config)
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- output_attentions=False,
- ):
- self_attention_outputs = self.attention(
- hidden_states, attention_mask, head_mask, output_attentions=output_attentions
- )
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- layer_output = apply_chunking_to_forward(
- self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
- )
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
- def feed_forward_chunk(self, attention_output):
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.output(intermediate_output, attention_output)
- return layer_output
-
-
-class MCTCTPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = MCTCTConfig
- base_model_prefix = "mctct"
- main_input_name = "input_features"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- std = self.config.initializer_range
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
- elif isinstance(module, MCTCTLayerNorm):
- module.singleton_weight.data.fill_(1.0)
- module.singleton_bias.data.zero_()
- if isinstance(module, (nn.Linear, nn.Conv1d)):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
-
- def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
- """
- Computes the output length of the convolutional layers
- """
- dilation = 1
- for _, kernel_sz, stride in zip(
- range(self.config.num_conv_layers), self.config.conv_kernel, self.config.conv_stride
- ):
- padding = kernel_sz // 2
- input_lengths = input_lengths + 2 * padding - dilation * (kernel_sz - 1) - 1
- input_lengths = torch.div(input_lengths, stride, rounding_mode="trunc") + 1
-
- return input_lengths
-
- def _get_feature_vector_attention_mask(self, feature_vector_length, attention_mask):
- # generate creates 3D attention mask, because of the shape of input_features
- # convert it to 2D if thats the case
- if len(attention_mask.shape) > 2:
- attention_mask = attention_mask[:, :, -1]
-
- # subsampled_lengths = attention_mask.sum(-1)
- subsampled_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1))
- bsz = attention_mask.size()[0]
- attention_mask = torch.zeros(
- (bsz, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device
- )
-
- # these two operations makes sure that all values
- # before the output lengths indices are attended to
- attention_mask[(torch.arange(bsz, device=attention_mask.device), subsampled_lengths - 1)] = 1
- attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).long()
- return attention_mask
-
-
-MCTCT_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`MCTCTConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-MCTCT_INPUTS_DOCSTRING = r"""
- Args:
- input_features (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`Wav2Vec2CTCTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class MCTCTEncoder(MCTCTPreTrainedModel):
- def __init__(self, config: MCTCTConfig):
- super().__init__(config)
- self.hidden_dropout_prob = config.hidden_dropout_prob
-
- self.layer_norm = MCTCTLayerNorm()
- self.conv = MCTCTConv1dSubsampler(config)
- self.layers = nn.ModuleList([MCTCTLayer(config) for _ in range(config.num_hidden_layers)])
-
- self.gradient_checkpointing = False
-
- def forward(
- self,
- input_features: torch.Tensor,
- attention_mask: torch.Tensor,
- head_mask: torch.Tensor,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[Tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- input_features = self.layer_norm(input_features)
-
- inputs_embeds = self.conv(input_features)
-
- # subsample attention mask if necessary
- if attention_mask is not None:
- attention_mask = self._get_feature_vector_attention_mask(inputs_embeds.shape[1], attention_mask)
-
- hidden_states = nn.functional.dropout(inputs_embeds, p=self.hidden_dropout_prob, training=self.training)
-
- # expand attention_mask
- if attention_mask is not None:
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
- attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
-
- encoder_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- # check if head_mask has a correct number of layers specified if desired
- if head_mask is not None:
- if head_mask.size()[0] != len(self.layers):
- raise ValueError(
- f"The head_mask should be specified for {len(self.layers)} layers, "
- f"but it is for {head_mask.size()[0]}."
- )
-
- deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled()
- for idx, encoder_layer in enumerate(self.layers):
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
-
- # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
- dropout_probability = torch.rand([])
-
- skip_the_layer = True if self.training and (dropout_probability < self.config.layerdrop) else False
- if not skip_the_layer or deepspeed_zero3_is_enabled:
- # under deepspeed zero3 all gpus must run in sync
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- encoder_layer.__call__,
- hidden_states,
- attention_mask,
- (head_mask[idx] if head_mask is not None else None),
- output_attentions,
- )
- else:
- layer_outputs = encoder_layer(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- if skip_the_layer:
- layer_outputs = (None, None)
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
- )
-
-
-@add_start_docstrings(
- "The bare M-CTC-T Model transformer outputting raw hidden-states without any specific head on top.",
- MCTCT_START_DOCSTRING,
-)
-class MCTCTModel(MCTCTPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.config = config
-
- self.encoder = MCTCTEncoder(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(MCTCT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="audio",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- input_features: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_features is None:
- raise ValueError("You have to specify input_features.")
-
- encoder_outputs = self.encoder(
- input_features,
- attention_mask=attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[1:]
-
- return BaseModelOutput(
- last_hidden_state=sequence_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """MCTCT Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).""",
- MCTCT_START_DOCSTRING,
-)
-class MCTCTForCTC(MCTCTPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.mctct = MCTCTModel(config)
-
- if config.vocab_size is None:
- raise ValueError(
- f"You are trying to instantiate {self.__class__} with a configuration that "
- "does not define the vocabulary size of the language model head. Please "
- "instantiate the model as follows: `MCTCTForCTC.from_pretrained(..., vocab_size=vocab_size)`. "
- "or define `vocab_size` of your model's configuration."
- )
- output_hidden_size = config.hidden_size
-
- self.ctc_head = nn.Linear(output_hidden_size, config.vocab_size)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(MCTCT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_CTC_EXPECTED_OUTPUT,
- expected_loss=_CTC_EXPECTED_LOSS,
- )
- def forward(
- self,
- input_features: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: Optional[torch.LongTensor] = None,
- ) -> Union[Tuple, CausalLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, target_length)`, *optional*):
- Labels for connectionist temporal classification. Note that `target_length` has to be smaller or equal to
- the sequence length of the output logits. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`.
- All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ...,
- config.vocab_size - 1]`.
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- outputs = self.mctct(
- input_features,
- attention_mask=attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
-
- logits = self.ctc_head(hidden_states)
-
- loss = None
- if labels is not None:
- if labels.max() >= self.config.vocab_size:
- raise ValueError(f"Label values must be <= vocab_size: {self.config.vocab_size}")
-
- # retrieve loss input_lengths from attention_mask
- attention_mask = (
- attention_mask
- if attention_mask is not None
- else torch.ones(input_features.shape[:-1], dtype=torch.long)
- )
- input_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long)
- # assuming that padded tokens are filled with -100
- # when not being attended to
- labels_mask = labels >= 0
- target_lengths = labels_mask.sum(-1)
- flattened_targets = labels.masked_select(labels_mask)
-
- # ctc_loss doesn't support fp16
- log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)
-
- with torch.backends.cudnn.flags(enabled=False):
- loss = nn.functional.ctc_loss(
- log_probs,
- flattened_targets,
- input_lengths,
- target_lengths,
- blank=self.config.pad_token_id,
- reduction=self.config.ctc_loss_reduction,
- zero_infinity=self.config.ctc_zero_infinity,
- )
-
- if not return_dict:
- output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutput(
- loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
diff --git a/transformers/models/deprecated/mctct/processing_mctct.py b/transformers/models/deprecated/mctct/processing_mctct.py
deleted file mode 100644
index 4e0cbe27dd9be0244d63a23256808cc421fa1fa5..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/mctct/processing_mctct.py
+++ /dev/null
@@ -1,142 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Speech processor class for M-CTC-T
-"""
-import warnings
-from contextlib import contextmanager
-
-from ....processing_utils import ProcessorMixin
-
-
-class MCTCTProcessor(ProcessorMixin):
- r"""
- Constructs a MCTCT processor which wraps a MCTCT feature extractor and a MCTCT tokenizer into a single processor.
-
- [`MCTCTProcessor`] offers all the functionalities of [`MCTCTFeatureExtractor`] and [`AutoTokenizer`]. See the
- [`~MCTCTProcessor.__call__`] and [`~MCTCTProcessor.decode`] for more information.
-
- Args:
- feature_extractor (`MCTCTFeatureExtractor`):
- An instance of [`MCTCTFeatureExtractor`]. The feature extractor is a required input.
- tokenizer (`AutoTokenizer`):
- An instance of [`AutoTokenizer`]. The tokenizer is a required input.
- """
-
- feature_extractor_class = "MCTCTFeatureExtractor"
- tokenizer_class = "AutoTokenizer"
-
- def __init__(self, feature_extractor, tokenizer):
- super().__init__(feature_extractor, tokenizer)
- self.current_processor = self.feature_extractor
- self._in_target_context_manager = False
-
- def __call__(self, *args, **kwargs):
- """
- When used in normal mode, this method forwards all its arguments to MCTCTFeatureExtractor's
- [`~MCTCTFeatureExtractor.__call__`] and returns its output. If used in the context
- [`~MCTCTProcessor.as_target_processor`] this method forwards all its arguments to AutoTokenizer's
- [`~AutoTokenizer.__call__`]. Please refer to the doctsring of the above two methods for more information.
- """
- # For backward compatibility
- if self._in_target_context_manager:
- return self.current_processor(*args, **kwargs)
-
- if "raw_speech" in kwargs:
- warnings.warn("Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.")
- audio = kwargs.pop("raw_speech")
- else:
- audio = kwargs.pop("audio", None)
- sampling_rate = kwargs.pop("sampling_rate", None)
- text = kwargs.pop("text", None)
- if len(args) > 0:
- audio = args[0]
- args = args[1:]
-
- if audio is None and text is None:
- raise ValueError("You need to specify either an `audio` or `text` input to process.")
-
- if audio is not None:
- inputs = self.feature_extractor(audio, *args, sampling_rate=sampling_rate, **kwargs)
- if text is not None:
- encodings = self.tokenizer(text, **kwargs)
-
- if text is None:
- return inputs
- elif audio is None:
- return encodings
- else:
- inputs["labels"] = encodings["input_ids"]
- return inputs
-
- def batch_decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to AutoTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please refer
- to the docstring of this method for more information.
- """
- return self.tokenizer.batch_decode(*args, **kwargs)
-
- def pad(self, *args, **kwargs):
- """
- When used in normal mode, this method forwards all its arguments to MCTCTFeatureExtractor's
- [`~MCTCTFeatureExtractor.pad`] and returns its output. If used in the context
- [`~MCTCTProcessor.as_target_processor`] this method forwards all its arguments to PreTrainedTokenizer's
- [`~PreTrainedTokenizer.pad`]. Please refer to the docstring of the above two methods for more information.
- """
- # For backward compatibility
- if self._in_target_context_manager:
- return self.current_processor.pad(*args, **kwargs)
-
- input_features = kwargs.pop("input_features", None)
- labels = kwargs.pop("labels", None)
- if len(args) > 0:
- input_features = args[0]
- args = args[1:]
-
- if input_features is not None:
- input_features = self.feature_extractor.pad(input_features, *args, **kwargs)
- if labels is not None:
- labels = self.tokenizer.pad(labels, **kwargs)
-
- if labels is None:
- return input_features
- elif input_features is None:
- return labels
- else:
- input_features["labels"] = labels["input_ids"]
- return input_features
-
- def decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to AutoTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to the
- docstring of this method for more information.
- """
- return self.tokenizer.decode(*args, **kwargs)
-
- @contextmanager
- def as_target_processor(self):
- """
- Temporarily sets the tokenizer for processing the input. Useful for encoding the labels when fine-tuning MCTCT.
- """
- warnings.warn(
- "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
- "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
- "your audio inputs, or in a separate call."
- )
- self._in_target_context_manager = True
- self.current_processor = self.tokenizer
- yield
- self.current_processor = self.feature_extractor
- self._in_target_context_manager = False
diff --git a/transformers/models/deprecated/mmbt/__init__.py b/transformers/models/deprecated/mmbt/__init__.py
deleted file mode 100644
index e467090cb4fbfa55ec51ec8232a54180c532ad6c..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/mmbt/__init__.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {"configuration_mmbt": ["MMBTConfig"]}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_mmbt"] = ["MMBTForClassification", "MMBTModel", "ModalEmbeddings"]
-
-
-if TYPE_CHECKING:
- from .configuration_mmbt import MMBTConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_mmbt import MMBTForClassification, MMBTModel, ModalEmbeddings
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deprecated/mmbt/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/mmbt/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 5f7198e1459eb1f54d1d379e935dc2cd8c04a155..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/mmbt/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/mmbt/__pycache__/configuration_mmbt.cpython-310.pyc b/transformers/models/deprecated/mmbt/__pycache__/configuration_mmbt.cpython-310.pyc
deleted file mode 100644
index 81165cf772a8615fe780e2bfa338e44726d41805..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/mmbt/__pycache__/configuration_mmbt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/mmbt/__pycache__/modeling_mmbt.cpython-310.pyc b/transformers/models/deprecated/mmbt/__pycache__/modeling_mmbt.cpython-310.pyc
deleted file mode 100644
index 255399ec4f961e2dee8a6329cf0daca33af49393..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/mmbt/__pycache__/modeling_mmbt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/mmbt/configuration_mmbt.py b/transformers/models/deprecated/mmbt/configuration_mmbt.py
deleted file mode 100644
index df5161b0927ad26279a273216d1d9ab6d465063a..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/mmbt/configuration_mmbt.py
+++ /dev/null
@@ -1,42 +0,0 @@
-# coding=utf-8
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Copyright (c) HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" MMBT configuration"""
-
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class MMBTConfig(object):
- """
- This is the configuration class to store the configuration of a [`MMBTModel`]. It is used to instantiate a MMBT
- model according to the specified arguments, defining the model architecture.
-
- Args:
- config ([`PreTrainedConfig`]):
- Config of the underlying Transformer models. Its values are copied over to use a single config.
- num_labels (`int`, *optional*):
- Size of final Linear layer for classification.
- modal_hidden_size (`int`, *optional*, defaults to 2048):
- Embedding dimension of the non-text modality encoder.
- """
-
- def __init__(self, config, num_labels=None, modal_hidden_size=2048):
- self.__dict__ = config.__dict__
- self.modal_hidden_size = modal_hidden_size
- if num_labels:
- self.num_labels = num_labels
diff --git a/transformers/models/deprecated/mmbt/modeling_mmbt.py b/transformers/models/deprecated/mmbt/modeling_mmbt.py
deleted file mode 100644
index 8dc450ce8f6c13346f30e7da045a927a1186e089..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/mmbt/modeling_mmbt.py
+++ /dev/null
@@ -1,408 +0,0 @@
-# coding=utf-8
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Copyright (c) HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch MMBT model."""
-
-
-import torch
-from torch import nn
-from torch.nn import CrossEntropyLoss, MSELoss
-
-from ....modeling_outputs import BaseModelOutputWithPooling, SequenceClassifierOutput
-from ....modeling_utils import ModuleUtilsMixin
-from ....utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "MMBTConfig"
-
-
-class ModalEmbeddings(nn.Module):
- """Generic Modal Embeddings which takes in an encoder, and a transformer embedding."""
-
- def __init__(self, config, encoder, embeddings):
- super().__init__()
- self.config = config
- self.encoder = encoder
- self.proj_embeddings = nn.Linear(config.modal_hidden_size, config.hidden_size)
- self.position_embeddings = embeddings.position_embeddings
- self.token_type_embeddings = embeddings.token_type_embeddings
- self.word_embeddings = embeddings.word_embeddings
- self.LayerNorm = embeddings.LayerNorm
- self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
-
- def forward(self, input_modal, start_token=None, end_token=None, position_ids=None, token_type_ids=None):
- token_embeddings = self.proj_embeddings(self.encoder(input_modal))
- seq_length = token_embeddings.size(1)
-
- if start_token is not None:
- start_token_embeds = self.word_embeddings(start_token)
- seq_length += 1
- token_embeddings = torch.cat([start_token_embeds.unsqueeze(1), token_embeddings], dim=1)
-
- if end_token is not None:
- end_token_embeds = self.word_embeddings(end_token)
- seq_length += 1
- token_embeddings = torch.cat([token_embeddings, end_token_embeds.unsqueeze(1)], dim=1)
-
- if position_ids is None:
- position_ids = torch.arange(seq_length, dtype=torch.long, device=input_modal.device)
- position_ids = position_ids.unsqueeze(0).expand(input_modal.size(0), seq_length)
-
- if token_type_ids is None:
- token_type_ids = torch.zeros(
- (input_modal.size(0), seq_length), dtype=torch.long, device=input_modal.device
- )
-
- position_embeddings = self.position_embeddings(position_ids)
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
- embeddings = token_embeddings + position_embeddings + token_type_embeddings
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-MMBT_START_DOCSTRING = r"""
- MMBT model was proposed in [Supervised Multimodal Bitransformers for Classifying Images and
- Text](https://github.com/facebookresearch/mmbt) by Douwe Kiela, Suvrat Bhooshan, Hamed Firooz, Davide Testuggine.
- It's a supervised multimodal bitransformer model that fuses information from text and other image encoders, and
- obtain state-of-the-art performance on various multimodal classification benchmark tasks.
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`MMBTConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration.
- transformer (`nn.Module`): A text transformer that is used by MMBT.
- It should have embeddings, encoder, and pooler attributes.
- encoder (`nn.Module`): Encoder for the second modality.
- It should take in a batch of modal inputs and return k, n dimension embeddings.
-"""
-
-MMBT_INPUTS_DOCSTRING = r"""
- Args:
- input_modal (`torch.FloatTensor` of shape `(batch_size, ***)`):
- The other modality data. It will be the shape that the encoder for that type expects. e.g. With an Image
- Encoder, the shape would be (batch_size, channels, height, width)
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. It does not expect [CLS] token to be added as it's
- appended to the end of other modality embeddings. Indices can be obtained using [`AutoTokenizer`]. See
- [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- modal_start_tokens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Optional start token to be added to Other Modality Embedding. [CLS] Most commonly used for classification
- tasks.
- modal_end_tokens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Optional end token to be added to Other Modality Embedding. [SEP] Most commonly used.
- attention_mask (*optional*) `torch.FloatTensor` of shape `(batch_size, sequence_length)`:
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (*optional*) `torch.LongTensor` of shape `(batch_size, sequence_length)`:
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- modal_token_type_ids (*optional*) `torch.LongTensor` of shape `(batch_size, modal_sequence_length)`:
- Segment token indices to indicate different portions of the non-text modality. The embeddings from these
- tokens will be summed with the respective token embeddings for the non-text modality.
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- modal_position_ids (`torch.LongTensor` of shape `(batch_size, modal_sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings for the non-text modality.
- Selected in the range `[0, config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, embedding_dim)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare MMBT Model outputting raw hidden-states without any specific head on top.",
- MMBT_START_DOCSTRING,
-)
-class MMBTModel(nn.Module, ModuleUtilsMixin):
- def __init__(self, config, transformer, encoder):
- super().__init__()
- self.config = config
- self.transformer = transformer
- self.modal_encoder = ModalEmbeddings(config, encoder, transformer.embeddings)
-
- @add_start_docstrings_to_model_forward(MMBT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_modal,
- input_ids=None,
- modal_start_tokens=None,
- modal_end_tokens=None,
- attention_mask=None,
- token_type_ids=None,
- modal_token_type_ids=None,
- position_ids=None,
- modal_position_ids=None,
- head_mask=None,
- inputs_embeds=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- ):
- r"""
- Returns:
-
- Examples:
-
- ```python
- # For example purposes. Not runnable.
- transformer = BertModel.from_pretrained("google-bert/bert-base-uncased")
- encoder = ImageEncoder(args)
- mmbt = MMBTModel(config, transformer, encoder)
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_txt_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_txt_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- modal_embeddings = self.modal_encoder(
- input_modal,
- start_token=modal_start_tokens,
- end_token=modal_end_tokens,
- position_ids=modal_position_ids,
- token_type_ids=modal_token_type_ids,
- )
-
- input_modal_shape = modal_embeddings.size()[:-1]
-
- if token_type_ids is None:
- token_type_ids = torch.ones(input_txt_shape, dtype=torch.long, device=device)
-
- txt_embeddings = self.transformer.embeddings(
- input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
- )
-
- embedding_output = torch.cat([modal_embeddings, txt_embeddings], 1)
-
- input_shape = embedding_output.size()[:-1]
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device)
- else:
- attention_mask = torch.cat(
- [torch.ones(input_modal_shape, device=device, dtype=torch.long), attention_mask], dim=1
- )
- if encoder_attention_mask is None:
- encoder_attention_mask = torch.ones(input_shape, device=device)
- else:
- encoder_attention_mask = torch.cat(
- [torch.ones(input_modal_shape, device=device), encoder_attention_mask], dim=1
- )
-
- extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)
- encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- encoder_outputs = self.transformer.encoder(
- embedding_output,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_extended_attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = encoder_outputs[0]
- pooled_output = self.transformer.pooler(sequence_output)
-
- if not return_dict:
- return (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
-
-@add_start_docstrings(
- """
- MMBT Model with a sequence classification/regression head on top (a linear layer on top of the pooled output)
- """,
- MMBT_START_DOCSTRING,
- MMBT_INPUTS_DOCSTRING,
-)
-class MMBTForClassification(nn.Module):
- r"""
- **labels**: (*optional*) `torch.LongTensor` of shape `(batch_size,)`:
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-
- Returns: *Tuple* comprising various elements depending on the configuration (config) and inputs: **loss**:
- (*optional*, returned when `labels` is provided) `torch.FloatTensor` of shape `(1,)`: Classification (or
- regression if config.num_labels==1) loss. **logits**:
- `torch.FloatTensor` of shape `(batch_size, config.num_labels)` Classification (or regression if
- config.num_labels==1) scores (before SoftMax).
- **hidden_states**: (*optional*, returned when `output_hidden_states=True`) list of `torch.FloatTensor` (one for
- the output of each layer + the output of the embeddings) of shape `(batch_size, sequence_length, hidden_size)`:
- Hidden-states of the model at the output of each layer plus the initial embedding outputs. **attentions**:
- (*optional*, returned when `output_attentions=True`) list of `torch.FloatTensor` (one for each layer) of shape
- `(batch_size, num_heads, sequence_length, sequence_length)`: Attentions weights after the attention softmax, used
- to compute the weighted average in the self-attention heads.
-
- Examples:
-
- ```python
- # For example purposes. Not runnable.
- transformer = BertModel.from_pretrained("google-bert/bert-base-uncased")
- encoder = ImageEncoder(args)
- model = MMBTForClassification(config, transformer, encoder)
- outputs = model(input_modal, input_ids, labels=labels)
- loss, logits = outputs[:2]
- ```"""
-
- def __init__(self, config, transformer, encoder):
- super().__init__()
- self.num_labels = config.num_labels
-
- self.mmbt = MMBTModel(config, transformer, encoder)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- def forward(
- self,
- input_modal,
- input_ids=None,
- modal_start_tokens=None,
- modal_end_tokens=None,
- attention_mask=None,
- token_type_ids=None,
- modal_token_type_ids=None,
- position_ids=None,
- modal_position_ids=None,
- head_mask=None,
- inputs_embeds=None,
- labels=None,
- return_dict=None,
- ):
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.mmbt(
- input_modal=input_modal,
- input_ids=input_ids,
- modal_start_tokens=modal_start_tokens,
- modal_end_tokens=modal_end_tokens,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- modal_token_type_ids=modal_token_type_ids,
- position_ids=position_ids,
- modal_position_ids=modal_position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.num_labels == 1:
- # We are doing regression
- loss_fct = MSELoss()
- loss = loss_fct(logits.view(-1), labels.view(-1))
- else:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/deprecated/open_llama/__init__.py b/transformers/models/deprecated/open_llama/__init__.py
deleted file mode 100644
index 446c9f076d31347c496300f432908d56895f7e67..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/open_llama/__init__.py
+++ /dev/null
@@ -1,95 +0,0 @@
-# Copyright 2023 EleutherAI and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ....utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_sentencepiece_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_open_llama": ["OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP", "OpenLlamaConfig"],
-}
-
-try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_open_llama"] = ["LlamaTokenizer"]
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_open_llama_fast"] = ["LlamaTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_open_llama"] = [
- "OpenLlamaForCausalLM",
- "OpenLlamaModel",
- "OpenLlamaPreTrainedModel",
- "OpenLlamaForSequenceClassification",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_open_llama import OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenLlamaConfig
-
- try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from transformers import LlamaTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from transformers import LlamaTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_open_llama import (
- OpenLlamaForCausalLM,
- OpenLlamaForSequenceClassification,
- OpenLlamaModel,
- OpenLlamaPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deprecated/open_llama/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/open_llama/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 3aab5df782466c83e43a8f4ee22f170b9f1541f5..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/open_llama/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/open_llama/__pycache__/configuration_open_llama.cpython-310.pyc b/transformers/models/deprecated/open_llama/__pycache__/configuration_open_llama.cpython-310.pyc
deleted file mode 100644
index b0788c1b26ff00dfedb8a945f953d7329bc79195..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/open_llama/__pycache__/configuration_open_llama.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/open_llama/__pycache__/modeling_open_llama.cpython-310.pyc b/transformers/models/deprecated/open_llama/__pycache__/modeling_open_llama.cpython-310.pyc
deleted file mode 100644
index d4667d4333c9069f8636a7cb1609dc640d39715a..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/open_llama/__pycache__/modeling_open_llama.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/open_llama/configuration_open_llama.py b/transformers/models/deprecated/open_llama/configuration_open_llama.py
deleted file mode 100644
index 0111e031251a2c7108419827f7b8c35eab0edcc3..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/open_llama/configuration_open_llama.py
+++ /dev/null
@@ -1,170 +0,0 @@
-# coding=utf-8
-# Copyright 2023 EleutherAI and the HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Open-Llama model configuration"""
-
-from ....configuration_utils import PretrainedConfig
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from .._archive_maps import OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class OpenLlamaConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`OpenLlamaModel`]. It is used to instantiate an
- Open-Llama model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the
- [s-JoL/Open-Llama-V1](https://huggingface.co/s-JoL/Open-Llama-V1).
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 32000):
- Vocabulary size of the Open-Llama model. Defines the number of different tokens that can be represented by
- the `inputs_ids` passed when calling [`OpenLlamaModel`]
- hidden_size (`int`, *optional*, defaults to 4096):
- Dimension of the hidden representations.
- intermediate_size (`int`, *optional*, defaults to 11008):
- Dimension of the MLP representations.
- num_hidden_layers (`int`, *optional*, defaults to 32):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 32):
- Number of attention heads for each attention layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
- The non-linear activation function (function or string) in the decoder.
- max_position_embeddings (`int`, *optional*, defaults to 2048):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- rms_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the rms normalization layers.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- tie_word_embeddings(`bool`, *optional*, defaults to `False`):
- Whether to tie weight embeddings
- rope_theta (`float`, *optional*, defaults to 10000.0):
- The base period of the RoPE embeddings.
- rope_scaling (`Dict`, *optional*):
- Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
- strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
- `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
- `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
- these scaling strategies behave:
- https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
- experimental feature, subject to breaking API changes in future versions.
-
- Example:
-
- ```python
- >>> from transformers import OpenLlamaModel, OpenLlamaConfig
-
- >>> # Initializing a Open-Llama open_llama-7b style configuration
- >>> configuration = OpenLlamaConfig()
-
- >>> # Initializing a model from the open_llama-7b style configuration
- >>> model = OpenLlamaModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "open-llama"
-
- def __init__(
- self,
- vocab_size=100000,
- hidden_size=4096,
- intermediate_size=11008,
- num_hidden_layers=32,
- num_attention_heads=32,
- hidden_act="silu",
- max_position_embeddings=2048,
- initializer_range=0.02,
- rms_norm_eps=1e-6,
- use_cache=True,
- pad_token_id=0,
- bos_token_id=1,
- eos_token_id=2,
- tie_word_embeddings=False,
- use_memory_efficient_attention=True,
- hidden_dropout_prob=0.1,
- attention_dropout_prob=0.1,
- use_stable_embedding=True,
- shared_input_output_embedding=True,
- rope_theta=10000.0,
- rope_scaling=None,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.hidden_size = hidden_size
- self.intermediate_size = intermediate_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.hidden_act = hidden_act
- self.initializer_range = initializer_range
- self.rms_norm_eps = rms_norm_eps
- self.use_cache = use_cache
- self.use_memory_efficient_attention = kwargs.pop(
- "use_memorry_efficient_attention", use_memory_efficient_attention
- )
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_dropout_prob = attention_dropout_prob
- self.use_stable_embedding = use_stable_embedding
- self.shared_input_output_embedding = shared_input_output_embedding
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
- self._rope_scaling_validation()
-
- super().__init__(
- pad_token_id=pad_token_id,
- bos_token_id=bos_token_id,
- eos_token_id=eos_token_id,
- tie_word_embeddings=tie_word_embeddings,
- **kwargs,
- )
-
- # Copied from transformers.models.llama.configuration_llama.LlamaConfig._rope_scaling_validation
- def _rope_scaling_validation(self):
- """
- Validate the `rope_scaling` configuration.
- """
- if self.rope_scaling is None:
- return
-
- if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
- raise ValueError(
- "`rope_scaling` must be a dictionary with two fields, `type` and `factor`, " f"got {self.rope_scaling}"
- )
- rope_scaling_type = self.rope_scaling.get("type", None)
- rope_scaling_factor = self.rope_scaling.get("factor", None)
- if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
- raise ValueError(
- f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
- )
- if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
- raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
diff --git a/transformers/models/deprecated/open_llama/modeling_open_llama.py b/transformers/models/deprecated/open_llama/modeling_open_llama.py
deleted file mode 100644
index 098f8c7da50d5e3b737ad6a5e12ccb456e71356d..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/open_llama/modeling_open_llama.py
+++ /dev/null
@@ -1,969 +0,0 @@
-# coding=utf-8
-# Copyright 2023 EleutherAI and the HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Open-Llama model."""
-import math
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ....activations import ACT2FN
-from ....modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
-from ....modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
-from ....modeling_utils import PreTrainedModel
-from ....utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from .configuration_open_llama import OpenLlamaConfig
-
-
-logger = logging.get_logger(__name__)
-
-try:
- from xformers import ops as xops
-except ImportError:
- xops = None
-
-
-_CONFIG_FOR_DOC = "OpenLlamaConfig"
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->OpenLlama
-class OpenLlamaRMSNorm(nn.Module):
- def __init__(self, hidden_size, eps=1e-6):
- """
- OpenLlamaRMSNorm is equivalent to T5LayerNorm
- """
- super().__init__()
- self.weight = nn.Parameter(torch.ones(hidden_size))
- self.variance_epsilon = eps
-
- def forward(self, hidden_states):
- input_dtype = hidden_states.dtype
- hidden_states = hidden_states.to(torch.float32)
- variance = hidden_states.pow(2).mean(-1, keepdim=True)
- hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
- return self.weight * hidden_states.to(input_dtype)
-
-
-# Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->OpenLlama
-class OpenLlamaRotaryEmbedding(nn.Module):
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
- super().__init__()
-
- self.dim = dim
- self.max_position_embeddings = max_position_embeddings
- self.base = base
- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
-
- # Build here to make `torch.jit.trace` work.
- self._set_cos_sin_cache(
- seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
- )
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
-
- def forward(self, x, seq_len=None):
- # x: [bs, num_attention_heads, seq_len, head_size]
- if seq_len > self.max_seq_len_cached:
- self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-
- return (
- self.cos_cached[:seq_len].to(dtype=x.dtype),
- self.sin_cached[:seq_len].to(dtype=x.dtype),
- )
-
-
-# Copied from transformers.models.falcon.modeling_falcon.FalconLinearScalingRotaryEmbedding with Falcon->OpenLlama
-class OpenLlamaLinearScalingRotaryEmbedding(OpenLlamaRotaryEmbedding):
- """OpenLlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
-
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
- self.scaling_factor = scaling_factor
- super().__init__(dim, max_position_embeddings, base, device)
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
- t = t / self.scaling_factor
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
-
-
-# Copied from transformers.models.falcon.modeling_falcon.FalconDynamicNTKScalingRotaryEmbedding with Falcon->OpenLlama
-class OpenLlamaDynamicNTKScalingRotaryEmbedding(OpenLlamaRotaryEmbedding):
- """OpenLlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
-
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
- self.scaling_factor = scaling_factor
- super().__init__(dim, max_position_embeddings, base, device)
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
-
- if seq_len > self.max_position_embeddings:
- base = self.base * (
- (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
- ) ** (self.dim / (self.dim - 2))
- inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
-
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
-
-
-def rotate_half(x):
- """Rotates half the hidden dims of the input."""
- x1 = x[..., : x.shape[-1] // 2]
- x2 = x[..., x.shape[-1] // 2 :]
- return torch.cat((-x2, x1), dim=-1)
-
-
-# Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
- """Applies Rotary Position Embedding to the query and key tensors.
-
- Args:
- q (`torch.Tensor`): The query tensor.
- k (`torch.Tensor`): The key tensor.
- cos (`torch.Tensor`): The cosine part of the rotary embedding.
- sin (`torch.Tensor`): The sine part of the rotary embedding.
- position_ids (`torch.Tensor`):
- The position indices of the tokens corresponding to the query and key tensors. For example, this can be
- used to pass offsetted position ids when working with a KV-cache.
- unsqueeze_dim (`int`, *optional*, defaults to 1):
- The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
- sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
- that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
- k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
- cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
- the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
- Returns:
- `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
- """
- cos = cos[position_ids].unsqueeze(unsqueeze_dim)
- sin = sin[position_ids].unsqueeze(unsqueeze_dim)
- q_embed = (q * cos) + (rotate_half(q) * sin)
- k_embed = (k * cos) + (rotate_half(k) * sin)
- return q_embed, k_embed
-
-
-class OpenLlamaMLP(nn.Module):
- def __init__(
- self,
- hidden_size: int,
- intermediate_size: int,
- hidden_act: str,
- dropout_prob: float,
- ):
- super().__init__()
- self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
- self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
- self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
- self.act_fn = ACT2FN[hidden_act]
- self.dropout = nn.Dropout(dropout_prob)
-
- def forward(self, x):
- out = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
- return self.dropout(out)
-
-
-class OpenLlamaAttention(nn.Module):
- """Multi-headed attention from 'Attention Is All You Need' paper"""
-
- def __init__(self, config: OpenLlamaConfig):
- super().__init__()
- self.config = config
- self.hidden_size = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.hidden_size // self.num_heads
- self.max_position_embeddings = config.max_position_embeddings
- self.dropout_prob = config.attention_dropout_prob
- self.rope_theta = config.rope_theta
-
- if (self.head_dim * self.num_heads) != self.hidden_size:
- raise ValueError(
- f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
- f" and `num_heads`: {self.num_heads})."
- )
- self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
- self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
- self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
- self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
- self._init_rope()
-
- # Copied from transformers.models.llama.modeling_llama.LlamaAttention._init_rope with Llama->OpenLlama
- def _init_rope(self):
- if self.config.rope_scaling is None:
- self.rotary_emb = OpenLlamaRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- base=self.rope_theta,
- )
- else:
- scaling_type = self.config.rope_scaling["type"]
- scaling_factor = self.config.rope_scaling["factor"]
- if scaling_type == "linear":
- self.rotary_emb = OpenLlamaLinearScalingRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- scaling_factor=scaling_factor,
- base=self.rope_theta,
- )
- elif scaling_type == "dynamic":
- self.rotary_emb = OpenLlamaDynamicNTKScalingRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- scaling_factor=scaling_factor,
- base=self.rope_theta,
- )
- else:
- raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
-
- def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
- return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- bsz, q_len, _ = hidden_states.size()
-
- query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
- key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
- value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-
- kv_seq_len = key_states.shape[-2]
- if past_key_value is not None:
- kv_seq_len += past_key_value[0].shape[-2]
- cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
- # [bsz, nh, t, hd]
-
- if past_key_value is not None:
- # reuse k, v, self_attention
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
-
- past_key_value = (key_states, value_states) if use_cache else None
-
- if self.config.use_memory_efficient_attention and xops is not None and self.training:
- attn_weights = None
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.transpose(1, 2)
- attn_output = xops.memory_efficient_attention(
- query_states, key_states, value_states, attn_bias=xops.LowerTriangularMask(), p=self.dropout_prob
- )
- else:
- attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-
- if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
- raise ValueError(
- f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- if attention_mask is not None:
- if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
- raise ValueError(
- f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
- )
- attn_weights = attn_weights + attention_mask
- attn_weights = torch.max(
- attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min, device=attn_weights.device)
- )
-
- # upcast attention to fp32
- attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
- attn_output = torch.matmul(attn_weights, value_states)
-
- if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.transpose(1, 2)
- attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-
- attn_output = self.o_proj(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights, past_key_value
-
-
-class OpenLlamaDecoderLayer(nn.Module):
- def __init__(self, config: OpenLlamaConfig):
- super().__init__()
- self.hidden_size = config.hidden_size
- self.self_attn = OpenLlamaAttention(config=config)
- self.mlp = OpenLlamaMLP(
- hidden_size=self.hidden_size,
- intermediate_size=config.intermediate_size,
- hidden_act=config.hidden_act,
- dropout_prob=config.hidden_dropout_prob,
- )
- self.input_layernorm = OpenLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
- self.post_attention_layernorm = OpenLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: Optional[bool] = False,
- use_cache: Optional[bool] = False,
- ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
- `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
- """
-
- residual = hidden_states
-
- hidden_states = self.input_layernorm(hidden_states)
-
- # Self Attention
- hidden_states, self_attn_weights, present_key_value = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- )
- hidden_states = residual + hidden_states
-
- # Fully Connected
- residual = hidden_states
- hidden_states = self.post_attention_layernorm(hidden_states)
- hidden_states = self.mlp(hidden_states)
- hidden_states = residual + hidden_states
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights,)
-
- if use_cache:
- outputs += (present_key_value,)
-
- return outputs
-
-
-OPEN_LLAMA_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`OpenLlamaConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-@add_start_docstrings(
- "The bare Open-Llama Model outputting raw hidden-states without any specific head on top.",
- OPEN_LLAMA_START_DOCSTRING,
-)
-class OpenLlamaPreTrainedModel(PreTrainedModel):
- config_class = OpenLlamaConfig
- base_model_prefix = "model"
- supports_gradient_checkpointing = True
- _no_split_modules = ["OpenLlamaDecoderLayer"]
-
- def _init_weights(self, module):
- std = self.config.initializer_range
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- if self.config.use_stable_embedding:
- torch.nn.init.xavier_normal_(module.weight.data)
- else:
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
-
-OPEN_LLAMA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
- it.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
- `past_key_values`).
-
- If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
- and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
- information on the default strategy.
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
- `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
- `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
-
- Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
- blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Open-Llama Model outputting raw hidden-states without any specific head on top.",
- OPEN_LLAMA_START_DOCSTRING,
-)
-class OpenLlamaModel(OpenLlamaPreTrainedModel):
- """
- Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`OpenLlamaDecoderLayer`]
-
- Args:
- config: OpenLlamaConfig
- """
-
- def __init__(self, config: OpenLlamaConfig):
- super().__init__(config)
- self.padding_idx = config.pad_token_id
- self.vocab_size = config.vocab_size
-
- self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
- if config.use_stable_embedding:
- self.embed_layer_norm = nn.LayerNorm(config.hidden_size)
- else:
- self.embed_layer_norm = None
- self.layers = nn.ModuleList([OpenLlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
- self.norm = OpenLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
- self.gradient_checkpointing = False
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embed_tokens
-
- def set_input_embeddings(self, value):
- self.embed_tokens = value
-
- @add_start_docstrings_to_model_forward(OPEN_LLAMA_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPast]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # retrieve input_ids and inputs_embeds
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
- elif input_ids is not None:
- batch_size, seq_length = input_ids.shape
- elif inputs_embeds is not None:
- batch_size, seq_length, _ = inputs_embeds.shape
- else:
- raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
-
- seq_length_with_past = seq_length
- past_key_values_length = 0
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- if past_key_values is not None:
- past_key_values_length = past_key_values[0][0].shape[2]
- seq_length_with_past = seq_length_with_past + past_key_values_length
-
- if position_ids is None:
- device = input_ids.device if input_ids is not None else inputs_embeds.device
- position_ids = torch.arange(
- past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
- )
- position_ids = position_ids.unsqueeze(0)
-
- if inputs_embeds is None:
- inputs_embeds = self.embed_tokens(input_ids)
- if self.embed_layer_norm:
- inputs_embeds = self.embed_layer_norm(inputs_embeds)
- # embed positions
- if self.config.use_memory_efficient_attention and self.training:
- attention_mask = None
- elif attention_mask is None:
- attention_mask = torch.ones(
- (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
- )
-
- input_shape = (batch_size, seq_length)
- attention_mask = _prepare_4d_causal_attention_mask(
- attention_mask, input_shape, inputs_embeds, past_key_values_length
- )
-
- hidden_states = inputs_embeds
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- next_decoder_cache = () if use_cache else None
-
- for idx, decoder_layer in enumerate(self.layers):
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- past_key_value = past_key_values[idx] if past_key_values is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- decoder_layer.__call__,
- hidden_states,
- attention_mask,
- position_ids,
- None,
- output_attentions,
- None,
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- )
-
- hidden_states = layer_outputs[0]
-
- if use_cache:
- next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- hidden_states = self.norm(hidden_states)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- next_cache = next_decoder_cache if use_cache else None
- if not return_dict:
- return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=next_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- )
-
-
-class OpenLlamaForCausalLM(OpenLlamaPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.model = OpenLlamaModel(config)
- if config.shared_input_output_embedding:
- self.lm_head = None
- else:
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.model.embed_tokens
-
- def set_input_embeddings(self, value):
- self.model.embed_tokens = value
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def set_decoder(self, decoder):
- self.model = decoder
-
- def get_decoder(self):
- return self.model
-
- @add_start_docstrings_to_model_forward(OPEN_LLAMA_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- Args:
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, OpenLlamaForCausalLM
-
- >>> model = OpenLlamaForCausalLM.from_pretrained("openlm-research/open_llama_7b")
- >>> tokenizer = AutoTokenizer.from_pretrained("openlm-research/open_llama_7b")
-
- >>> prompt = "Hey, are you conscious? Can you talk to me?"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
- ```"""
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- if self.config.shared_input_output_embedding:
- logits = torch.einsum(
- "blh,vh->blv", hidden_states.to(self.model.embed_tokens.weight.device), self.model.embed_tokens.weight
- )
- else:
- logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(logits.device)
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- loss = loss_fct(shift_logits, shift_labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def prepare_inputs_for_generation(
- self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
- ):
- if past_key_values is not None:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
-
- position_ids = kwargs.get("position_ids", None)
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- model_inputs = {"input_ids": input_ids}
-
- model_inputs.update(
- {
- "position_ids": position_ids,
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "attention_mask": attention_mask,
- }
- )
- return model_inputs
-
- @staticmethod
- def _reorder_cache(past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
- )
- return reordered_past
-
-
-@add_start_docstrings(
- """
- The LLaMa Model transformer with a sequence classification head on top (linear layer).
-
- [`OpenLlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal
- models (e.g. GPT-2) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- OPEN_LLAMA_START_DOCSTRING,
-)
-class OpenLlamaForSequenceClassification(OpenLlamaPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.model = OpenLlamaModel(config)
- self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.model.embed_tokens
-
- def set_input_embeddings(self, value):
- self.model.embed_tokens = value
-
- @add_start_docstrings_to_model_forward(OPEN_LLAMA_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.model(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size = input_ids.shape[0]
- else:
- batch_size = inputs_embeds.shape[0]
-
- if self.config.pad_token_id is None and batch_size != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
-
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
diff --git a/transformers/models/deprecated/retribert/__init__.py b/transformers/models/deprecated/retribert/__init__.py
deleted file mode 100644
index dba5e14594e16c19fc1a269a92e968fec35afc26..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/retribert/__init__.py
+++ /dev/null
@@ -1,73 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available
-
-
-_import_structure = {
- "configuration_retribert": ["RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "RetriBertConfig"],
- "tokenization_retribert": ["RetriBertTokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_retribert_fast"] = ["RetriBertTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_retribert"] = [
- "RETRIBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "RetriBertModel",
- "RetriBertPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_retribert import RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RetriBertConfig
- from .tokenization_retribert import RetriBertTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_retribert_fast import RetriBertTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_retribert import (
- RETRIBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
- RetriBertModel,
- RetriBertPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deprecated/retribert/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/retribert/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index eeb7c9d545eb2d050a52bc737e944597928b01be..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/retribert/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/retribert/__pycache__/configuration_retribert.cpython-310.pyc b/transformers/models/deprecated/retribert/__pycache__/configuration_retribert.cpython-310.pyc
deleted file mode 100644
index e2fd60717b1d89eb472e0f3eb0117640704697a1..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/retribert/__pycache__/configuration_retribert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/retribert/__pycache__/modeling_retribert.cpython-310.pyc b/transformers/models/deprecated/retribert/__pycache__/modeling_retribert.cpython-310.pyc
deleted file mode 100644
index 178c0328d1237c05dfba6b8cd0eda50fd7cc7863..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/retribert/__pycache__/modeling_retribert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/retribert/__pycache__/tokenization_retribert.cpython-310.pyc b/transformers/models/deprecated/retribert/__pycache__/tokenization_retribert.cpython-310.pyc
deleted file mode 100644
index 0f7c01f0d0e3ef59f61018ba37ccf5c11c1818db..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/retribert/__pycache__/tokenization_retribert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/retribert/__pycache__/tokenization_retribert_fast.cpython-310.pyc b/transformers/models/deprecated/retribert/__pycache__/tokenization_retribert_fast.cpython-310.pyc
deleted file mode 100644
index 31776530d14b99a3c225a9bd5a6bacca44447d55..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/retribert/__pycache__/tokenization_retribert_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/retribert/configuration_retribert.py b/transformers/models/deprecated/retribert/configuration_retribert.py
deleted file mode 100644
index c188c7347a8fb85c14748095b762e0a2583cee00..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/retribert/configuration_retribert.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" RetriBERT model configuration"""
-
-from ....configuration_utils import PretrainedConfig
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-from .._archive_maps import RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class RetriBertConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`RetriBertModel`]. It is used to instantiate a
- RetriBertModel model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the RetriBERT
- [yjernite/retribert-base-uncased](https://huggingface.co/yjernite/retribert-base-uncased) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the RetriBERT model. Defines the number of different tokens that can be represented by
- the `inputs_ids` passed when calling [`RetriBertModel`]
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 2):
- The vocabulary size of the *token_type_ids* passed into [`BertModel`].
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- share_encoders (`bool`, *optional*, defaults to `True`):
- Whether or not to use the same Bert-type encoder for the queries and document
- projection_dim (`int`, *optional*, defaults to 128):
- Final dimension of the query and document representation after projection
- """
-
- model_type = "retribert"
-
- def __init__(
- self,
- vocab_size=30522,
- hidden_size=768,
- num_hidden_layers=8,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=2,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- share_encoders=True,
- projection_dim=128,
- pad_token_id=0,
- **kwargs,
- ):
- super().__init__(pad_token_id=pad_token_id, **kwargs)
-
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.hidden_act = hidden_act
- self.intermediate_size = intermediate_size
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.share_encoders = share_encoders
- self.projection_dim = projection_dim
diff --git a/transformers/models/deprecated/retribert/modeling_retribert.py b/transformers/models/deprecated/retribert/modeling_retribert.py
deleted file mode 100644
index 7dba8a276eeb56c379f782ea2bd5e1be7119cfe7..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/retribert/modeling_retribert.py
+++ /dev/null
@@ -1,218 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-RetriBERT model
-"""
-
-
-import math
-from typing import Optional
-
-import torch
-import torch.utils.checkpoint as checkpoint
-from torch import nn
-
-from ....modeling_utils import PreTrainedModel
-from ....utils import add_start_docstrings, logging
-from ...bert.modeling_bert import BertModel
-from .configuration_retribert import RetriBertConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-from .._archive_maps import RETRIBERT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# INTERFACE FOR ENCODER AND TASK SPECIFIC MODEL #
-class RetriBertPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = RetriBertConfig
- load_tf_weights = None
- base_model_prefix = "retribert"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-RETRIBERT_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-@add_start_docstrings(
- """Bert Based model to embed queries or document for document retrieval.""",
- RETRIBERT_START_DOCSTRING,
-)
-class RetriBertModel(RetriBertPreTrainedModel):
- def __init__(self, config: RetriBertConfig) -> None:
- super().__init__(config)
- self.projection_dim = config.projection_dim
-
- self.bert_query = BertModel(config)
- self.bert_doc = None if config.share_encoders else BertModel(config)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.project_query = nn.Linear(config.hidden_size, config.projection_dim, bias=False)
- self.project_doc = nn.Linear(config.hidden_size, config.projection_dim, bias=False)
-
- self.ce_loss = nn.CrossEntropyLoss(reduction="mean")
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def embed_sentences_checkpointed(
- self,
- input_ids,
- attention_mask,
- sent_encoder,
- checkpoint_batch_size=-1,
- ):
- # reproduces BERT forward pass with checkpointing
- if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
- return sent_encoder(input_ids, attention_mask=attention_mask)[1]
- else:
- # prepare implicit variables
- device = input_ids.device
- input_shape = input_ids.size()
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
- head_mask = [None] * sent_encoder.config.num_hidden_layers
- extended_attention_mask: torch.Tensor = sent_encoder.get_extended_attention_mask(
- attention_mask, input_shape
- )
-
- # define function for checkpointing
- def partial_encode(*inputs):
- encoder_outputs = sent_encoder.encoder(
- inputs[0],
- attention_mask=inputs[1],
- head_mask=head_mask,
- )
- sequence_output = encoder_outputs[0]
- pooled_output = sent_encoder.pooler(sequence_output)
- return pooled_output
-
- # run embedding layer on everything at once
- embedding_output = sent_encoder.embeddings(
- input_ids=input_ids, position_ids=None, token_type_ids=token_type_ids, inputs_embeds=None
- )
- # run encoding and pooling on one mini-batch at a time
- pooled_output_list = []
- for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size)):
- b_embedding_output = embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
- b_attention_mask = extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
- pooled_output = checkpoint.checkpoint(partial_encode, b_embedding_output, b_attention_mask)
- pooled_output_list.append(pooled_output)
- return torch.cat(pooled_output_list, dim=0)
-
- def embed_questions(
- self,
- input_ids,
- attention_mask=None,
- checkpoint_batch_size=-1,
- ):
- q_reps = self.embed_sentences_checkpointed(
- input_ids,
- attention_mask,
- self.bert_query,
- checkpoint_batch_size,
- )
- return self.project_query(q_reps)
-
- def embed_answers(
- self,
- input_ids,
- attention_mask=None,
- checkpoint_batch_size=-1,
- ):
- a_reps = self.embed_sentences_checkpointed(
- input_ids,
- attention_mask,
- self.bert_query if self.bert_doc is None else self.bert_doc,
- checkpoint_batch_size,
- )
- return self.project_doc(a_reps)
-
- def forward(
- self,
- input_ids_query: torch.LongTensor,
- attention_mask_query: Optional[torch.FloatTensor],
- input_ids_doc: torch.LongTensor,
- attention_mask_doc: Optional[torch.FloatTensor],
- checkpoint_batch_size: int = -1,
- ) -> torch.FloatTensor:
- r"""
- Args:
- input_ids_query (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary for the queries in a batch.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask_query (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- input_ids_doc (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary for the documents in a batch.
- attention_mask_doc (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on documents padding token indices.
- checkpoint_batch_size (`int`, *optional*, defaults to `-1`):
- If greater than 0, uses gradient checkpointing to only compute sequence representation on
- `checkpoint_batch_size` examples at a time on the GPU. All query representations are still compared to
- all document representations in the batch.
-
- Return:
- `torch.FloatTensor``: The bidirectional cross-entropy loss obtained while trying to match each query to its
- corresponding document and each document to its corresponding query in the batch
- """
- device = input_ids_query.device
- q_reps = self.embed_questions(input_ids_query, attention_mask_query, checkpoint_batch_size)
- a_reps = self.embed_answers(input_ids_doc, attention_mask_doc, checkpoint_batch_size)
- compare_scores = torch.mm(q_reps, a_reps.t())
- loss_qa = self.ce_loss(compare_scores, torch.arange(compare_scores.shape[1]).to(device))
- loss_aq = self.ce_loss(compare_scores.t(), torch.arange(compare_scores.shape[0]).to(device))
- loss = (loss_qa + loss_aq) / 2
- return loss
diff --git a/transformers/models/deprecated/retribert/tokenization_retribert.py b/transformers/models/deprecated/retribert/tokenization_retribert.py
deleted file mode 100644
index c991f3972230bd9027f43137b2a51550d4655f4d..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/retribert/tokenization_retribert.py
+++ /dev/null
@@ -1,517 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for RetriBERT."""
-
-import collections
-import os
-import unicodedata
-from typing import List, Optional, Tuple
-
-from ....tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
-
-
-# Copied from transformers.models.bert.tokenization_bert.load_vocab
-def load_vocab(vocab_file):
- """Loads a vocabulary file into a dictionary."""
- vocab = collections.OrderedDict()
- with open(vocab_file, "r", encoding="utf-8") as reader:
- tokens = reader.readlines()
- for index, token in enumerate(tokens):
- token = token.rstrip("\n")
- vocab[token] = index
- return vocab
-
-
-# Copied from transformers.models.bert.tokenization_bert.whitespace_tokenize
-def whitespace_tokenize(text):
- """Runs basic whitespace cleaning and splitting on a piece of text."""
- text = text.strip()
- if not text:
- return []
- tokens = text.split()
- return tokens
-
-
-class RetriBertTokenizer(PreTrainedTokenizer):
- r"""
- Constructs a RetriBERT tokenizer.
-
- [`RetriBertTokenizer`] is identical to [`BertTokenizer`] and runs end-to-end tokenization: punctuation splitting
- and wordpiece.
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer
- to: this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- do_basic_tokenize (`bool`, *optional*, defaults to `True`):
- Whether or not to do basic tokenization before WordPiece.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.__init__
- def __init__(
- self,
- vocab_file,
- do_lower_case=True,
- do_basic_tokenize=True,
- never_split=None,
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- if not os.path.isfile(vocab_file):
- raise ValueError(
- f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
- " model use `tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- self.vocab = load_vocab(vocab_file)
- self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
- self.do_basic_tokenize = do_basic_tokenize
- if do_basic_tokenize:
- self.basic_tokenizer = BasicTokenizer(
- do_lower_case=do_lower_case,
- never_split=never_split,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- )
-
- self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token))
-
- super().__init__(
- do_lower_case=do_lower_case,
- do_basic_tokenize=do_basic_tokenize,
- never_split=never_split,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- @property
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.do_lower_case
- def do_lower_case(self):
- return self.basic_tokenizer.do_lower_case
-
- @property
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.vocab_size
- def vocab_size(self):
- return len(self.vocab)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_vocab
- def get_vocab(self):
- return dict(self.vocab, **self.added_tokens_encoder)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._tokenize
- def _tokenize(self, text, split_special_tokens=False):
- split_tokens = []
- if self.do_basic_tokenize:
- for token in self.basic_tokenizer.tokenize(
- text, never_split=self.all_special_tokens if not split_special_tokens else None
- ):
- # If the token is part of the never_split set
- if token in self.basic_tokenizer.never_split:
- split_tokens.append(token)
- else:
- split_tokens += self.wordpiece_tokenizer.tokenize(token)
- else:
- split_tokens = self.wordpiece_tokenizer.tokenize(text)
- return split_tokens
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.vocab.get(token, self.vocab.get(self.unk_token))
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.ids_to_tokens.get(index, self.unk_token)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.convert_tokens_to_string
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- out_string = " ".join(tokens).replace(" ##", "").strip()
- return out_string
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A BERT sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_special_tokens_mask
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- index = 0
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- else:
- vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(token + "\n")
- index += 1
- return (vocab_file,)
-
-
-# Copied from transformers.models.bert.tokenization_bert.BasicTokenizer
-class BasicTokenizer(object):
- """
- Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
-
- Args:
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- do_split_on_punc (`bool`, *optional*, defaults to `True`):
- In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
- the full context of the words, such as contractions.
- """
-
- def __init__(
- self,
- do_lower_case=True,
- never_split=None,
- tokenize_chinese_chars=True,
- strip_accents=None,
- do_split_on_punc=True,
- ):
- if never_split is None:
- never_split = []
- self.do_lower_case = do_lower_case
- self.never_split = set(never_split)
- self.tokenize_chinese_chars = tokenize_chinese_chars
- self.strip_accents = strip_accents
- self.do_split_on_punc = do_split_on_punc
-
- def tokenize(self, text, never_split=None):
- """
- Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
-
- Args:
- never_split (`List[str]`, *optional*)
- Kept for backward compatibility purposes. Now implemented directly at the base class level (see
- [`PreTrainedTokenizer.tokenize`]) List of token not to split.
- """
- # union() returns a new set by concatenating the two sets.
- never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
- text = self._clean_text(text)
-
- # This was added on November 1st, 2018 for the multilingual and Chinese
- # models. This is also applied to the English models now, but it doesn't
- # matter since the English models were not trained on any Chinese data
- # and generally don't have any Chinese data in them (there are Chinese
- # characters in the vocabulary because Wikipedia does have some Chinese
- # words in the English Wikipedia.).
- if self.tokenize_chinese_chars:
- text = self._tokenize_chinese_chars(text)
- # prevents treating the same character with different unicode codepoints as different characters
- unicode_normalized_text = unicodedata.normalize("NFC", text)
- orig_tokens = whitespace_tokenize(unicode_normalized_text)
- split_tokens = []
- for token in orig_tokens:
- if token not in never_split:
- if self.do_lower_case:
- token = token.lower()
- if self.strip_accents is not False:
- token = self._run_strip_accents(token)
- elif self.strip_accents:
- token = self._run_strip_accents(token)
- split_tokens.extend(self._run_split_on_punc(token, never_split))
-
- output_tokens = whitespace_tokenize(" ".join(split_tokens))
- return output_tokens
-
- def _run_strip_accents(self, text):
- """Strips accents from a piece of text."""
- text = unicodedata.normalize("NFD", text)
- output = []
- for char in text:
- cat = unicodedata.category(char)
- if cat == "Mn":
- continue
- output.append(char)
- return "".join(output)
-
- def _run_split_on_punc(self, text, never_split=None):
- """Splits punctuation on a piece of text."""
- if not self.do_split_on_punc or (never_split is not None and text in never_split):
- return [text]
- chars = list(text)
- i = 0
- start_new_word = True
- output = []
- while i < len(chars):
- char = chars[i]
- if _is_punctuation(char):
- output.append([char])
- start_new_word = True
- else:
- if start_new_word:
- output.append([])
- start_new_word = False
- output[-1].append(char)
- i += 1
-
- return ["".join(x) for x in output]
-
- def _tokenize_chinese_chars(self, text):
- """Adds whitespace around any CJK character."""
- output = []
- for char in text:
- cp = ord(char)
- if self._is_chinese_char(cp):
- output.append(" ")
- output.append(char)
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
- def _is_chinese_char(self, cp):
- """Checks whether CP is the codepoint of a CJK character."""
- # This defines a "chinese character" as anything in the CJK Unicode block:
- # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
- #
- # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
- # despite its name. The modern Korean Hangul alphabet is a different block,
- # as is Japanese Hiragana and Katakana. Those alphabets are used to write
- # space-separated words, so they are not treated specially and handled
- # like the all of the other languages.
- if (
- (cp >= 0x4E00 and cp <= 0x9FFF)
- or (cp >= 0x3400 and cp <= 0x4DBF) #
- or (cp >= 0x20000 and cp <= 0x2A6DF) #
- or (cp >= 0x2A700 and cp <= 0x2B73F) #
- or (cp >= 0x2B740 and cp <= 0x2B81F) #
- or (cp >= 0x2B820 and cp <= 0x2CEAF) #
- or (cp >= 0xF900 and cp <= 0xFAFF)
- or (cp >= 0x2F800 and cp <= 0x2FA1F) #
- ): #
- return True
-
- return False
-
- def _clean_text(self, text):
- """Performs invalid character removal and whitespace cleanup on text."""
- output = []
- for char in text:
- cp = ord(char)
- if cp == 0 or cp == 0xFFFD or _is_control(char):
- continue
- if _is_whitespace(char):
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
-
-# Copied from transformers.models.bert.tokenization_bert.WordpieceTokenizer
-class WordpieceTokenizer(object):
- """Runs WordPiece tokenization."""
-
- def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
- self.vocab = vocab
- self.unk_token = unk_token
- self.max_input_chars_per_word = max_input_chars_per_word
-
- def tokenize(self, text):
- """
- Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
- tokenization using the given vocabulary.
-
- For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
-
- Args:
- text: A single token or whitespace separated tokens. This should have
- already been passed through *BasicTokenizer*.
-
- Returns:
- A list of wordpiece tokens.
- """
-
- output_tokens = []
- for token in whitespace_tokenize(text):
- chars = list(token)
- if len(chars) > self.max_input_chars_per_word:
- output_tokens.append(self.unk_token)
- continue
-
- is_bad = False
- start = 0
- sub_tokens = []
- while start < len(chars):
- end = len(chars)
- cur_substr = None
- while start < end:
- substr = "".join(chars[start:end])
- if start > 0:
- substr = "##" + substr
- if substr in self.vocab:
- cur_substr = substr
- break
- end -= 1
- if cur_substr is None:
- is_bad = True
- break
- sub_tokens.append(cur_substr)
- start = end
-
- if is_bad:
- output_tokens.append(self.unk_token)
- else:
- output_tokens.extend(sub_tokens)
- return output_tokens
diff --git a/transformers/models/deprecated/retribert/tokenization_retribert_fast.py b/transformers/models/deprecated/retribert/tokenization_retribert_fast.py
deleted file mode 100644
index 97fbfc07d30ca65cd4d31dc2e221d0ef22073175..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/retribert/tokenization_retribert_fast.py
+++ /dev/null
@@ -1,180 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for RetriBERT."""
-
-import json
-from typing import List, Optional, Tuple
-
-from tokenizers import normalizers
-
-from ....tokenization_utils_fast import PreTrainedTokenizerFast
-from ....utils import logging
-from .tokenization_retribert import RetriBertTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
-
-
-class RetriBertTokenizerFast(PreTrainedTokenizerFast):
- r"""
- Construct a "fast" RetriBERT tokenizer (backed by HuggingFace's *tokenizers* library).
-
- [`RetriBertTokenizerFast`] is identical to [`BertTokenizerFast`] and runs end-to-end tokenization: punctuation
- splitting and wordpiece.
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- clean_text (`bool`, *optional*, defaults to `True`):
- Whether or not to clean the text before tokenization by removing any control characters and replacing all
- whitespaces by the classic one.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see [this
- issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- wordpieces_prefix (`str`, *optional*, defaults to `"##"`):
- The prefix for subwords.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = RetriBertTokenizer
- model_input_names = ["input_ids", "attention_mask"]
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.__init__
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- do_lower_case=True,
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- super().__init__(
- vocab_file,
- tokenizer_file=tokenizer_file,
- do_lower_case=do_lower_case,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- normalizer_state = json.loads(self.backend_tokenizer.normalizer.__getstate__())
- if (
- normalizer_state.get("lowercase", do_lower_case) != do_lower_case
- or normalizer_state.get("strip_accents", strip_accents) != strip_accents
- or normalizer_state.get("handle_chinese_chars", tokenize_chinese_chars) != tokenize_chinese_chars
- ):
- normalizer_class = getattr(normalizers, normalizer_state.pop("type"))
- normalizer_state["lowercase"] = do_lower_case
- normalizer_state["strip_accents"] = strip_accents
- normalizer_state["handle_chinese_chars"] = tokenize_chinese_chars
- self.backend_tokenizer.normalizer = normalizer_class(**normalizer_state)
-
- self.do_lower_case = do_lower_case
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A BERT sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- output = [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
-
- if token_ids_1 is not None:
- output += token_ids_1 + [self.sep_token_id]
-
- return output
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
diff --git a/transformers/models/deprecated/tapex/__init__.py b/transformers/models/deprecated/tapex/__init__.py
deleted file mode 100644
index 82bbacd15b0d00509972e16ac406005ee97370f7..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/tapex/__init__.py
+++ /dev/null
@@ -1,29 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ....utils import _LazyModule
-
-
-_import_structure = {"tokenization_tapex": ["TapexTokenizer"]}
-
-
-if TYPE_CHECKING:
- from .tokenization_tapex import TapexTokenizer
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
diff --git a/transformers/models/deprecated/tapex/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/tapex/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index d9c17b9c8070c366faa7fd35ff02e184db1bdd3c..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/tapex/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/tapex/__pycache__/tokenization_tapex.cpython-310.pyc b/transformers/models/deprecated/tapex/__pycache__/tokenization_tapex.cpython-310.pyc
deleted file mode 100644
index 389ee933edbe1b36f49b16c8fb3350610b648de1..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/tapex/__pycache__/tokenization_tapex.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/tapex/tokenization_tapex.py b/transformers/models/deprecated/tapex/tokenization_tapex.py
deleted file mode 100644
index cd3d353b526c4a8d4ba033ce8c0ed47137852b30..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/tapex/tokenization_tapex.py
+++ /dev/null
@@ -1,1467 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Microsoft Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for TAPEX."""
-
-import json
-import os
-import random
-from functools import lru_cache
-from typing import Dict, List, Optional, Tuple, Union
-
-import regex as re
-
-from ....file_utils import ExplicitEnum, PaddingStrategy, TensorType, add_end_docstrings, is_pandas_available
-from ....tokenization_utils import AddedToken, PreTrainedTokenizer
-from ....tokenization_utils_base import ENCODE_KWARGS_DOCSTRING, BatchEncoding, TextInput, TruncationStrategy
-from ....utils import logging
-
-
-if is_pandas_available():
- import pandas as pd
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt"}
-
-
-class TapexTruncationStrategy(ExplicitEnum):
- """
- Possible values for the `truncation` argument in [`~TapasTokenizer.__call__`]. Useful for tab-completion in an IDE.
- """
-
- DROP_ROWS_TO_FIT = "drop_rows_to_fit"
-
-
-TAPEX_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING = r"""
- add_special_tokens (`bool`, *optional*, defaults to `True`):
- Whether or not to encode the sequences with the special tokens relative to their model.
- padding (`bool`, `str` or [`~file_utils.PaddingStrategy`], *optional*, defaults to `False`):
- Activates and controls padding. Accepts the following values:
-
- - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
- sequence if provided).
- - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided.
- - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
- lengths).
- truncation (`bool`, `str`, [`TapexTruncationStrategy`] or [`~tokenization_utils_base.TruncationStrategy`],
- *optional*, defaults to `False`):
-
- Activates and controls truncation. Accepts the following values:
-
- - `'drop_rows_to_fit'`: Truncate to a maximum length specified with the argument `max_length` or to the
- maximum acceptable input length for the model if that argument is not provided. This will truncate
- row by row, removing rows from the table.
- - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or
- to the maximum acceptable input length for the model if that argument is not provided. This will
- truncate token by token, removing a token from the longest sequence in the pair if a pair of
- sequences (or a batch of pairs) is provided.
- - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the
- maximum acceptable input length for the model if that argument is not provided. This will only
- truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
- - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the
- maximum acceptable input length for the model if that argument is not provided. This will only
- truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
- - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths
- greater than the model maximum admissible input size).
- max_length (`int`, *optional*):
- Controls the maximum length to use by one of the truncation/padding parameters. If left unset or set to
- `None`, this will use the predefined model maximum length if a maximum length is required by one of the
- truncation/padding parameters. If the model has no specific maximum input length (like XLNet)
- truncation/padding to a maximum length will be deactivated.
- stride (`int`, *optional*, defaults to 0):
- If set to a number along with `max_length`, the overflowing tokens returned when
- `return_overflowing_tokens=True` will contain some tokens from the end of the truncated sequence
- returned to provide some overlap between truncated and overflowing sequences. The value of this
- argument defines the number of overlapping tokens.
- pad_to_multiple_of (`int`, *optional*):
- If set will pad the sequence to a multiple of the provided value. This is especially useful to enable
- the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta).
- return_tensors (`str` or [`~file_utils.TensorType`], *optional*):
- If set, will return tensors instead of list of python integers. Acceptable values are:
-
- - `'tf'`: Return TensorFlow `tf.constant` objects.
- - `'pt'`: Return PyTorch `torch.Tensor` objects.
- - `'np'`: Return Numpy `np.ndarray` objects.
-"""
-
-
-@lru_cache()
-def bytes_to_unicode():
- """
- Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
- characters the bpe code barfs on. The reversible bpe codes work on unicode strings. This means you need a large #
- of unicode characters in your vocab if you want to avoid UNKs. When you're at something like a 10B token dataset
- you end up needing around 5K for decent coverage. This is a significant percentage of your normal, say, 32K bpe
- vocab. To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
- """
- bs = (
- list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
- )
- cs = bs[:]
- n = 0
- for b in range(2**8):
- if b not in bs:
- bs.append(b)
- cs.append(2**8 + n)
- n += 1
- cs = [chr(n) for n in cs]
- return dict(zip(bs, cs))
-
-
-def get_pairs(word):
- """
- Return set of symbol pairs in a word. Word is represented as tuple of symbols (symbols being variable-length
- strings).
- """
- pairs = set()
- prev_char = word[0]
- for char in word[1:]:
- pairs.add((prev_char, char))
- prev_char = char
- return pairs
-
-
-class IndexedRowTableLinearize:
- """
- FORMAT: col: col1 | col2 | col 3 row 1 : val1 | val2 | val3 row 2 : ...
- """
-
- def process_table(self, table_content: Dict):
- """
- Given a table, TableLinearize aims at converting it into a flatten sequence with special symbols.
- """
- assert "header" in table_content and "rows" in table_content, self.PROMPT_MESSAGE
- # process header
- table_str = self.process_header(table_content["header"]) + " "
- # process rows
- for i, row_example in enumerate(table_content["rows"]):
- # NOTE: the row should start from row 1 instead of 0
- table_str += self.process_row(row_example, row_index=i + 1) + " "
- return table_str.strip()
-
- def process_header(self, headers: List):
- """
- Given a list of headers, TableLinearize aims at converting it into a flatten sequence with special symbols.
- """
- return "col : " + " | ".join(headers)
-
- def process_row(self, row: List, row_index: int):
- """
- Given a row, TableLinearize aims at converting it into a flatten sequence with special symbols.
- """
- row_str = ""
- row_cell_values = []
- for cell_value in row:
- if isinstance(cell_value, int):
- row_cell_values.append(str(cell_value))
- else:
- row_cell_values.append(cell_value)
- row_str += " | ".join(row_cell_values)
- return "row " + str(row_index) + " : " + row_str
-
-
-class TapexTokenizer(PreTrainedTokenizer):
- r"""
- Construct a TAPEX tokenizer. Based on byte-level Byte-Pair-Encoding (BPE).
-
- This tokenizer can be used to flatten one or more table(s) and concatenate them with one or more related sentences
- to be used by TAPEX models. The format that the TAPEX tokenizer creates is the following:
-
- sentence col: col1 | col2 | col 3 row 1 : val1 | val2 | val3 row 2 : ...
-
- The tokenizer supports a single table + single query, a single table and multiple queries (in which case the table
- will be duplicated for every query), a single query and multiple tables (in which case the query will be duplicated
- for every table), and multiple tables and queries. In other words, you can provide a batch of tables + questions to
- the tokenizer for instance to prepare them for the model.
-
- Tokenization itself is based on the BPE algorithm. It is identical to the one used by BART, RoBERTa and GPT-2.
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- merges_file (`str`):
- Path to the merges file.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- errors (`str`, *optional*, defaults to `"replace"`):
- Paradigm to follow when decoding bytes to UTF-8. See
- [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the beginning of
- sequence. The token used is the `cls_token`.
-
-
-
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sequence token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the end of sequence.
- The token used is the `sep_token`.
-
-
-
- sep_token (`str`, *optional*, defaults to `""`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- cls_token (`str`, *optional*, defaults to `""`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- mask_token (`str`, *optional*, defaults to `""`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial space to the input. This allows to treat the leading word just as any
- other word. (BART tokenizer detect beginning of words by the preceding space).
- max_cell_length (`int`, *optional*, defaults to 15):
- Maximum number of characters per cell when linearizing a table. If this number is exceeded, truncation
- takes place.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- merges_file,
- do_lower_case=True,
- errors="replace",
- bos_token="",
- eos_token="",
- sep_token="",
- cls_token="",
- unk_token="",
- pad_token="",
- mask_token="",
- add_prefix_space=False,
- max_cell_length=15,
- **kwargs,
- ):
- bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
- eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
- sep_token = AddedToken(sep_token, lstrip=False, rstrip=False) if isinstance(sep_token, str) else sep_token
- cls_token = AddedToken(cls_token, lstrip=False, rstrip=False) if isinstance(cls_token, str) else cls_token
- unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
- pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
-
- # Mask token behave like a normal word, i.e. include the space before it
- mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token
-
- with open(vocab_file, encoding="utf-8") as vocab_handle:
- self.encoder = json.load(vocab_handle)
- self.decoder = {v: k for k, v in self.encoder.items()}
- self.errors = errors # how to handle errors in decoding
- self.byte_encoder = bytes_to_unicode()
- self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
- with open(merges_file, encoding="utf-8") as merges_handle:
- bpe_merges = merges_handle.read().split("\n")[1:-1]
- bpe_merges = [tuple(merge.split()) for merge in bpe_merges]
- self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
- self.cache = {}
- self.add_prefix_space = add_prefix_space
- self.do_lower_case = do_lower_case
-
- # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
- self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
-
- # additional properties
-
- super().__init__(
- vocab_file=vocab_file,
- merges_file=merges_file,
- do_lower_case=do_lower_case,
- errors=errors,
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- sep_token=sep_token,
- cls_token=cls_token,
- pad_token=pad_token,
- mask_token=mask_token,
- add_prefix_space=add_prefix_space,
- max_cell_length=max_cell_length,
- **kwargs,
- )
-
- self.max_cell_length = max_cell_length
- self.table_linearize = IndexedRowTableLinearize()
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A TAPEX sequence has the following format:
- - single sequence: ` X `
- - pair of sequences: ` A B `
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + sep + token_ids_1 + sep
-
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Args:
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is None:
- return [1] + ([0] * len(token_ids_0)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1]
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Args:
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. TAPEX does not:
- make use of token type ids, therefore a list of zeros is returned.
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- Returns:
- `List[int]`: List of zeros.
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
-
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep + sep + token_ids_1 + sep) * [0]
-
- def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
- add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space)
- if (is_split_into_words or add_prefix_space) and (len(text) > 0 and not text[0].isspace()):
- text = " " + text
- return (text, kwargs)
-
- @property
- def vocab_size(self):
- return len(self.encoder)
-
- def get_vocab(self):
- return dict(self.encoder, **self.added_tokens_encoder)
-
- def bpe(self, token):
- if token in self.cache:
- return self.cache[token]
- word = tuple(token)
- pairs = get_pairs(word)
-
- if not pairs:
- return token
-
- while True:
- bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
- if bigram not in self.bpe_ranks:
- break
- first, second = bigram
- new_word = []
- i = 0
- while i < len(word):
- try:
- j = word.index(first, i)
- except ValueError:
- new_word.extend(word[i:])
- break
- else:
- new_word.extend(word[i:j])
- i = j
-
- if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
- new_word.append(first + second)
- i += 2
- else:
- new_word.append(word[i])
- i += 1
- new_word = tuple(new_word)
- word = new_word
- if len(word) == 1:
- break
- else:
- pairs = get_pairs(word)
- word = " ".join(word)
- self.cache[token] = word
- return word
-
- def _tokenize(self, text):
- """Tokenize a string."""
- bpe_tokens = []
- for token in re.findall(self.pat, text):
- token = "".join(
- self.byte_encoder[b] for b in token.encode("utf-8")
- ) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
- bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
- return bpe_tokens
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- text = "".join(tokens)
- text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
- return text
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- merge_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
- )
-
- with open(vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
- index = 0
- with open(merge_file, "w", encoding="utf-8") as writer:
- writer.write("#version: 0.2\n")
- for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
- " Please check that the tokenizer is not corrupted!"
- )
- index = token_index
- writer.write(" ".join(bpe_tokens) + "\n")
- index += 1
-
- return vocab_file, merge_file
-
- @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, TAPEX_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING)
- def __call__(
- self,
- table: Union["pd.DataFrame", List["pd.DataFrame"]] = None,
- query: Optional[Union[TextInput, List[TextInput]]] = None,
- answer: Union[str, List[str]] = None,
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str, TruncationStrategy] = None,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- """
- Main method to tokenize and prepare for the model one or several table-sequence pair(s).
-
- Args:
- table (`pd.DataFrame`, `List[pd.DataFrame]`):
- Table(s) containing tabular data.
- query (`str` or `List[str]`, *optional*):
- Sentence or batch of sentences related to one or more table(s) to be encoded. Note that the number of
- sentences must match the number of tables.
- answer (`str` or `List[str]`, *optional*):
- Optionally, the corresponding answer to the questions as supervision.
- """
-
- if table is not None:
- return self.source_call_func(
- table=table,
- query=query,
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
- elif answer is not None:
- return self.target_call_func(
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
- else:
- raise ValueError("You need to provide either a `table` or an `answer`.")
-
- def source_call_func(
- self,
- table: Union["pd.DataFrame", List["pd.DataFrame"]],
- query: Optional[Union[TextInput, List[TextInput]]] = None,
- answer: Union[str, List[str]] = None,
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str, TruncationStrategy] = None,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- # Input type checking for clearer error
- valid_table = False
- valid_query = False
-
- # Check that table have a valid type
- if isinstance(table, pd.DataFrame):
- valid_table = True
- elif isinstance(table, (list, tuple)) and isinstance(table[0], pd.DataFrame):
- valid_table = True
-
- # Check that query have a valid type
- if query is None or isinstance(query, str):
- valid_query = True
- elif isinstance(query, (list, tuple)):
- if len(query) == 0 or isinstance(query[0], str):
- valid_query = True
-
- if not valid_table:
- raise ValueError(
- "table input must of type `pd.DataFrame` (single example), `List[pd.DataFrame]` (batch of examples). "
- )
- if not valid_query:
- raise ValueError("query input must of type `str` (single example), `List[str]` (batch of examples). ")
- is_batched = isinstance(table, (list, tuple)) or isinstance(query, (list, tuple))
-
- if is_batched:
- return self.batch_encode_plus(
- table=table,
- query=query,
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
- else:
- return self.encode_plus(
- table=table,
- query=query,
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
-
- @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, TAPEX_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING)
- def batch_encode_plus(
- self,
- table: Union["pd.DataFrame", List["pd.DataFrame"]],
- query: Optional[List[TextInput]] = None,
- answer: List[str] = None,
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str] = None,
- max_length: Optional[int] = None,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- """
-
-
- This method is deprecated, `__call__` should be used instead.
-
-
- """
- # Backward compatibility for 'truncation_strategy', 'pad_to_max_length'
- padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- verbose=verbose,
- **kwargs,
- )
-
- return self._batch_encode_plus(
- table=table,
- query=query,
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding_strategy=padding_strategy,
- truncation_strategy=truncation_strategy,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
-
- def _batch_encode_plus(
- self,
- table: Union["pd.DataFrame", List["pd.DataFrame"]],
- query: Optional[List[TextInput]] = None,
- answer: Optional[List[str]] = None,
- add_special_tokens: bool = True,
- padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
- truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- if return_offsets_mapping:
- raise NotImplementedError(
- "return_offset_mapping is not available when using Python tokenizers. "
- "To use this feature, change your tokenizer to one deriving from "
- "transformers.PreTrainedTokenizerFast."
- )
-
- if isinstance(table, pd.DataFrame) and isinstance(query, (list, tuple)):
- # single table, many queries case
- # duplicate table for every query
- table = [table] * len(query)
- if isinstance(table, (list, tuple)) and isinstance(query, str):
- # many tables, single query case
- # duplicate query for every table
- query = [query] * len(table)
-
- batch_outputs = self._batch_prepare_for_model(
- table=table,
- query=query,
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding_strategy=padding_strategy,
- truncation_strategy=truncation_strategy,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=pad_to_multiple_of,
- return_attention_mask=return_attention_mask,
- return_token_type_ids=return_token_type_ids,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_length=return_length,
- return_tensors=return_tensors,
- verbose=verbose,
- )
-
- return BatchEncoding(batch_outputs)
-
- @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, TAPEX_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING)
- def _batch_prepare_for_model(
- self,
- table: Union["pd.DataFrame", List["pd.DataFrame"]],
- query: Optional[Union[TextInput, List[TextInput]]] = None,
- answer: Optional[Union[str, List[str]]] = None,
- add_special_tokens: bool = True,
- padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
- truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[str] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- ) -> BatchEncoding:
- """
- This method adds special tokens, truncates sequences if overflowing while taking into account the special
- tokens and manages a moving window (with user defined stride) for overflowing tokens.
- """
- batch_outputs = {}
- if answer is None:
- answer = [None] * len(table)
- for _table, _query, _answer in zip(table, query, answer):
- text = self.prepare_table_query(
- _table, _query, _answer, truncation_strategy=truncation_strategy, max_length=max_length
- )
-
- if self.do_lower_case:
- text = text.lower()
-
- tokens = self.tokenize(text)
- outputs = self.prepare_for_model(
- ids=self.convert_tokens_to_ids(tokens),
- add_special_tokens=add_special_tokens,
- padding=PaddingStrategy.DO_NOT_PAD.value, # we pad in batch afterwards
- truncation=truncation_strategy.value,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=None, # we pad in batch afterwards
- return_attention_mask=False, # we pad in batch afterwards
- return_token_type_ids=return_token_type_ids,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_length=return_length,
- return_tensors=None, # We convert the whole batch to tensors at the end
- prepend_batch_axis=False,
- verbose=verbose,
- )
-
- for key, value in outputs.items():
- if key not in batch_outputs:
- batch_outputs[key] = []
- batch_outputs[key].append(value)
-
- batch_outputs = self.pad(
- batch_outputs,
- padding=padding_strategy.value,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_attention_mask=return_attention_mask,
- )
-
- batch_outputs = BatchEncoding(batch_outputs, tensor_type=return_tensors)
-
- return batch_outputs
-
- @add_end_docstrings(ENCODE_KWARGS_DOCSTRING)
- def encode(
- self,
- table: "pd.DataFrame",
- query: Optional[TextInput] = None,
- answer: Optional[str] = None,
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str, TruncationStrategy, TapexTruncationStrategy] = None,
- max_length: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- **kwargs,
- ) -> List[int]:
- """
- Prepare a table, a string and possible answer for the model. This method does not return token type IDs,
- attention masks, etc. which are necessary for the model to work correctly. Use this method if you want to build
- your processing on your own, otherwise refer to `__call__`.
- """
- encoded_inputs = self.encode_plus(
- table,
- query=query,
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- return_tensors=return_tensors,
- **kwargs,
- )
-
- return encoded_inputs["input_ids"]
-
- @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, TAPEX_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING)
- def encode_plus(
- self,
- table: "pd.DataFrame",
- query: Optional[TextInput] = None,
- answer: Optional[str] = None,
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str] = None,
- max_length: Optional[int] = None,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- # Backward compatibility for 'truncation_strategy', 'pad_to_max_length'
- padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- verbose=verbose,
- **kwargs,
- )
-
- return self._encode_plus(
- table=table,
- query=query,
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding_strategy=padding_strategy,
- truncation_strategy=truncation_strategy,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
-
- def _encode_plus(
- self,
- table: "pd.DataFrame",
- query: Optional[TextInput] = None,
- answer: Optional[str] = None,
- add_special_tokens: bool = True,
- padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
- truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- if return_offsets_mapping:
- raise NotImplementedError(
- "return_offset_mapping is not available when using Python tokenizers. "
- "To use this feature, change your tokenizer to one deriving from "
- "transformers.PreTrainedTokenizerFast. "
- "More information on available tokenizers at "
- "https://github.com/huggingface/transformers/pull/2674"
- )
-
- text = self.prepare_table_query(
- table, query, answer, truncation_strategy=truncation_strategy, max_length=max_length
- )
-
- # if necessary, perform lower case
- if self.do_lower_case:
- text = text.lower()
-
- tokens = self.tokenize(text)
-
- return self.prepare_for_model(
- ids=self.convert_tokens_to_ids(tokens),
- add_special_tokens=add_special_tokens,
- padding=padding_strategy.value,
- truncation=truncation_strategy.value,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- prepend_batch_axis=True,
- return_attention_mask=return_attention_mask,
- return_token_type_ids=return_token_type_ids,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_length=return_length,
- verbose=verbose,
- )
-
- def target_call_func(
- self,
- answer: Union[str, List[str]],
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str, TruncationStrategy] = None,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- """
- The method tokenizes and prepares the answer label for the model.
-
- Args:
- answer (`str` or `List[str]`):
- Corresponding answer supervision to the queries for training the model.
- """
- is_batched = isinstance(answer, (list, tuple))
-
- if is_batched:
- return self.target_batch_encode_plus(
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
- else:
- return self.target_encode_plus(
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
-
- def target_batch_encode_plus(
- self,
- answer: List[str],
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str] = None,
- max_length: Optional[int] = None,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- """
- Prepare answer strings for the model.
-
- Args:
- answer `List[str]`:
- Corresponding answer supervision to the queries for training the model.
- """
- # Backward compatibility for 'truncation_strategy', 'pad_to_max_length'
- padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- verbose=verbose,
- **kwargs,
- )
-
- return self._target_batch_encode_plus(
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding_strategy=padding_strategy,
- truncation_strategy=truncation_strategy,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
-
- def _target_batch_encode_plus(
- self,
- answer: List[str],
- add_special_tokens: bool = True,
- padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
- truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- batch_outputs = {}
- for text in answer:
- if self.do_lower_case:
- text = text.lower()
-
- tokens = self.tokenize(text)
- outputs = self.prepare_for_model(
- ids=self.convert_tokens_to_ids(tokens),
- add_special_tokens=add_special_tokens,
- padding=PaddingStrategy.DO_NOT_PAD.value, # we pad in batch afterwards
- truncation=truncation_strategy.value,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=None, # we pad in batch afterwards
- return_attention_mask=False, # we pad in batch afterwards
- return_token_type_ids=return_token_type_ids,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_length=return_length,
- return_tensors=None, # We convert the whole batch to tensors at the end
- prepend_batch_axis=False,
- verbose=verbose,
- )
-
- for key, value in outputs.items():
- if key not in batch_outputs:
- batch_outputs[key] = []
- batch_outputs[key].append(value)
-
- batch_outputs = self.pad(
- batch_outputs,
- padding=padding_strategy.value,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_attention_mask=return_attention_mask,
- )
-
- batch_outputs = BatchEncoding(batch_outputs, tensor_type=return_tensors)
-
- return BatchEncoding(batch_outputs)
-
- def target_encode(
- self,
- answer: str,
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str, TruncationStrategy, TapexTruncationStrategy] = None,
- max_length: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- **kwargs,
- ) -> List[int]:
- """
- Prepare the answer string for the model. This method does not return token type IDs, attention masks, etc.
- which are necessary for the model to work correctly. Use this method if you want to build your processing on
- your own, otherwise refer to `__call__`.
-
- Args:
- answer `str`:
- Corresponding answer supervision to the queries for training the model
- """
- encoded_outputs = self.target_encode_plus(
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- return_tensors=return_tensors,
- **kwargs,
- )
-
- return encoded_outputs["input_ids"]
-
- def target_encode_plus(
- self,
- answer: str,
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str] = None,
- max_length: Optional[int] = None,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- """
- Prepare a answer string for the model.
-
- Args:
- answer `str`:
- Corresponding answer supervision to the queries for training the model.
- """
- # Backward compatibility for 'truncation_strategy', 'pad_to_max_length'
- padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- verbose=verbose,
- **kwargs,
- )
-
- return self._target_encode_plus(
- answer=answer,
- add_special_tokens=add_special_tokens,
- padding_strategy=padding_strategy,
- truncation_strategy=truncation_strategy,
- max_length=max_length,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- **kwargs,
- )
-
- def _target_encode_plus(
- self,
- answer: str,
- add_special_tokens: bool = True,
- padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
- truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- **kwargs,
- ) -> BatchEncoding:
- if return_offsets_mapping:
- raise NotImplementedError(
- "return_offset_mapping is not available when using Python tokenizers. "
- "To use this feature, change your tokenizer to one deriving from "
- "transformers.PreTrainedTokenizerFast. "
- "More information on available tokenizers at "
- "https://github.com/huggingface/transformers/pull/2674"
- )
-
- text = answer
-
- # if necessary, perform lower case
- if self.do_lower_case:
- text = text.lower()
-
- tokens = self.tokenize(text)
-
- return self.prepare_for_model(
- ids=self.convert_tokens_to_ids(tokens),
- add_special_tokens=add_special_tokens,
- padding=padding_strategy.value,
- truncation=truncation_strategy.value,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=pad_to_multiple_of,
- return_tensors=return_tensors,
- prepend_batch_axis=True,
- return_attention_mask=return_attention_mask,
- return_token_type_ids=return_token_type_ids,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_length=return_length,
- verbose=verbose,
- )
-
- def prepare_table_query(
- self,
- table,
- query,
- answer=None,
- truncation_strategy=Union[str, TruncationStrategy, TapexTruncationStrategy],
- max_length=None,
- ):
- """
- This method can be used to linearize a table and add a corresponding query.
-
- Optionally, it also handles truncation of the table (cells).
-
- An answer can be provided for more precise truncation.
- """
- if not table.empty:
- # step 1: create table dictionary
- table_content = {"header": list(table.columns), "rows": [list(row.values) for i, row in table.iterrows()]}
-
- # step 2: modify table internally
- # always truncate table cells based on self.max_cell_length
- # optionally truncate rows if truncation_strategy is set to it
- self.truncate_table_cells(table_content, query, answer)
- if truncation_strategy == TapexTruncationStrategy.DROP_ROWS_TO_FIT:
- self.truncate_table_rows(table_content, query, answer, max_length=max_length)
-
- # step 3: linearize table
- linear_table = self.table_linearize.process_table(table_content)
- else:
- linear_table = ""
-
- if linear_table == "":
- logger.warning(
- "You provide an empty table, or all cells contain much tokens (e.g., >= 1024 tokens). "
- + f"Please carefully check the corresponding table with the query : {query}."
- )
- if query == "":
- logger.warning("You provide nothing to query with respect to the table.")
- # step 4: concatenate query with linear_table
- separator = " " if query and linear_table else ""
- joint_input = (query + separator + linear_table) if query else linear_table
-
- return joint_input
-
- def truncate_table_cells(self, table_content: Dict, question: str, answer: List):
- # TODO (Qian): is it possible to revert the original cell if it is in the final answer?
- cell_mapping = {}
- for row in table_content["rows"]:
- for i, cell in enumerate(row):
- truncate_cell = self.truncate_cell(cell)
- if truncate_cell is not None:
- cell_mapping[cell] = truncate_cell
- row[i] = truncate_cell
-
- # modify the answer list
- if answer is not None:
- for i, case in enumerate(answer):
- if case in cell_mapping.keys():
- answer[i] = cell_mapping[case]
-
- def truncate_cell(self, cell_value):
- # do not process on these cases
- if isinstance(cell_value, int) or isinstance(cell_value, float):
- return cell_value
- if cell_value.strip() != "":
- try_tokens = self.tokenize(cell_value)
- if len(try_tokens) >= self.max_cell_length:
- retain_tokens = try_tokens[: self.max_cell_length]
- retain_cell_value = self.convert_tokens_to_string(retain_tokens)
- return retain_cell_value
- else:
- return None
- else:
- return cell_value
-
- def truncate_table_rows(
- self, table_content: Dict, question: str, answer: Optional[Union[str, List[str]]] = None, max_length=None
- ):
- """
- Args:
- table_content:
- {"header": xxx, "rows": xxx, "id" (Optionally): xxx}
-
- question:
- natural language sentence
-
- answer:
- if for training, is the supervision; otherwise will be empty
- """
- delete_ratio, remain_token_len = self.estimate_delete_ratio(table_content, question, max_length)
- # randomly delete unrelated rows
- self.delete_unrelated_rows(table_content, question, answer, delete_ratio)
- # guarantee the result < max_length
- maximum_keep_rows = 0
- for ind, row_example in enumerate(table_content["rows"]):
- value_string = self.table_linearize.process_row(row_example, ind + 1)
- value_token_len = len(self.tokenize(value_string))
- # over the size limit, and take action
- if value_token_len > remain_token_len:
- break
- remain_token_len -= value_token_len
- maximum_keep_rows += 1
- del table_content["rows"][maximum_keep_rows:]
-
- def estimate_delete_ratio(self, table_content: Dict, question: str, max_length=None):
- if "header" not in table_content or "rows" not in table_content:
- raise ValueError("The table content should contain both 'header' and 'rows' keys.")
- # calculate the tokens of header, special tokens will only be pre-prepended into question
- question_tokens = self.tokenize(question, add_special_tokens=True)
- # calculate the tokens of header
- header_string = self.table_linearize.process_header(table_content["header"])
- header_tokens = self.tokenize(header_string, add_special_tokens=False)
- # split all cell values into tokens and see how many can be accommodated
- used_token_len = len(question_tokens) + len(header_tokens)
- # remaining token space for rows
- remain_token_len = max_length - used_token_len
-
- value_string = ""
- for _, row_example in enumerate(table_content["rows"]):
- # use a general index to roughly estimate the overall token len
- value_string += self.table_linearize.process_row(row_example, 100) + " "
- value_token_len = len(self.tokenize(value_string))
-
- if value_token_len < remain_token_len:
- # no row will be deleted
- return 0.0, remain_token_len
- else:
- # calc a roughly delete rate
- return 1.0 - remain_token_len / value_token_len, remain_token_len
-
- def delete_unrelated_rows(self, table_content: Dict, question: str, answer: List, delete_ratio: float):
- """
- The argument answer is used only during training.
- """
- truncated_unrelated_indices = []
- related_indices = []
- if answer is None or len(answer) == 0:
- answer_set = set()
- else:
- answer_set = {ans_ex.lower() for ans_ex in answer}
- # add question key words into answer set
- if question is not None:
- answer_set.update(question.split())
- question_set = set(question.strip("?!.,").split(" "))
- row_max_len = len(table_content["rows"])
- for _row_idx, row in enumerate(table_content["rows"]):
- lower_row = {str(cell).lower() for cell in row}
- if len(lower_row & answer_set) == 0 and len(lower_row & question_set) == 0:
- truncated_unrelated_indices.append(_row_idx)
- else:
- # add neighbours to preserve information aggressively
- related_indices.extend([_row_idx - 2, _row_idx - 1, _row_idx, _row_idx + 1, _row_idx + 2])
-
- # remove the neighbours
- truncated_unrelated_indices = [
- _row_idx for _row_idx in truncated_unrelated_indices if _row_idx not in related_indices
- ]
- # select some cases to drop
- drop_items = min(len(truncated_unrelated_indices), int(len(table_content["rows"]) * delete_ratio))
- drop_row_indices = random.choices(truncated_unrelated_indices, k=drop_items)
-
- for _row_idx in reversed(range(row_max_len)):
- if _row_idx in drop_row_indices:
- del table_content["rows"][_row_idx]
-
- # only when the drop ratio is too large, logging for warning.
- if "id" in table_content and len(drop_row_indices) > 0:
- logger.warning("Delete {:.2f} rows in table {}".format(len(drop_row_indices), table_content["id"]))
diff --git a/transformers/models/deprecated/trajectory_transformer/__init__.py b/transformers/models/deprecated/trajectory_transformer/__init__.py
deleted file mode 100644
index b7af1bb48cb7d6a495611b0dadfc910779262813..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/trajectory_transformer/__init__.py
+++ /dev/null
@@ -1,63 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {
- "configuration_trajectory_transformer": [
- "TRAJECTORY_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "TrajectoryTransformerConfig",
- ],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_trajectory_transformer"] = [
- "TRAJECTORY_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TrajectoryTransformerModel",
- "TrajectoryTransformerPreTrainedModel",
- "load_tf_weights_in_trajectory_transformer",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_trajectory_transformer import (
- TRAJECTORY_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
- TrajectoryTransformerConfig,
- )
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_trajectory_transformer import (
- TRAJECTORY_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
- TrajectoryTransformerModel,
- TrajectoryTransformerPreTrainedModel,
- load_tf_weights_in_trajectory_transformer,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deprecated/trajectory_transformer/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/trajectory_transformer/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 69a410819ad52c1909aa0dff7641b10c321e9cce..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/trajectory_transformer/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/trajectory_transformer/__pycache__/configuration_trajectory_transformer.cpython-310.pyc b/transformers/models/deprecated/trajectory_transformer/__pycache__/configuration_trajectory_transformer.cpython-310.pyc
deleted file mode 100644
index 7922d84bc82ced25003ff6ffc35a83b6881f22a7..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/trajectory_transformer/__pycache__/configuration_trajectory_transformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/trajectory_transformer/__pycache__/convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/deprecated/trajectory_transformer/__pycache__/convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 017b8cba5419a997d907c6e2de918806e4ec2eab..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/trajectory_transformer/__pycache__/convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/trajectory_transformer/__pycache__/modeling_trajectory_transformer.cpython-310.pyc b/transformers/models/deprecated/trajectory_transformer/__pycache__/modeling_trajectory_transformer.cpython-310.pyc
deleted file mode 100644
index 283f51b117c0f40968c809f415ce98ec9e5cff72..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/trajectory_transformer/__pycache__/modeling_trajectory_transformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/trajectory_transformer/configuration_trajectory_transformer.py b/transformers/models/deprecated/trajectory_transformer/configuration_trajectory_transformer.py
deleted file mode 100644
index eccb71fcc429e742c2df458f5c831f8ca2df451d..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/trajectory_transformer/configuration_trajectory_transformer.py
+++ /dev/null
@@ -1,155 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The Trajectory Transformers paper authors and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TrajectoryTransformer model configuration"""
-
-from ....configuration_utils import PretrainedConfig
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from .._archive_maps import TRAJECTORY_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class TrajectoryTransformerConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`TrajectoryTransformerModel`]. It is used to
- instantiate an TrajectoryTransformer model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the
- TrajectoryTransformer
- [CarlCochet/trajectory-transformer-halfcheetah-medium-v2](https://huggingface.co/CarlCochet/trajectory-transformer-halfcheetah-medium-v2)
- architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 100):
- Vocabulary size of the TrajectoryTransformer model. Defines the number of different tokens that can be
- represented by the `trajectories` passed when calling [`TrajectoryTransformerModel`]
- action_weight (`int`, *optional*, defaults to 5):
- Weight of the action in the loss function
- reward_weight (`int`, *optional*, defaults to 1):
- Weight of the reward in the loss function
- value_weight (`int`, *optional*, defaults to 1):
- Weight of the value in the loss function
- block_size (`int`, *optional*, defaults to 249):
- Size of the blocks in the trajectory transformer.
- action_dim (`int`, *optional*, defaults to 6):
- Dimension of the action space.
- observation_dim (`int`, *optional*, defaults to 17):
- Dimension of the observation space.
- transition_dim (`int`, *optional*, defaults to 25):
- Dimension of the transition space.
- n_layer (`int`, *optional*, defaults to 4):
- Number of hidden layers in the Transformer encoder.
- n_head (`int`, *optional*, defaults to 4):
- Number of attention heads for each attention layer in the Transformer encoder.
- n_embd (`int`, *optional*, defaults to 128):
- Dimensionality of the embeddings and hidden states.
- resid_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- embd_pdrop (`int`, *optional*, defaults to 0.1):
- The dropout ratio for the embeddings.
- attn_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- kaiming_initializer_range (`float, *optional*, defaults to 1):
- A coefficient scaling the negative slope of the kaiming initializer rectifier for EinLinear layers.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- Example:
-
- ```python
- >>> from transformers import TrajectoryTransformerConfig, TrajectoryTransformerModel
-
- >>> # Initializing a TrajectoryTransformer CarlCochet/trajectory-transformer-halfcheetah-medium-v2 style configuration
- >>> configuration = TrajectoryTransformerConfig()
-
- >>> # Initializing a model (with random weights) from the CarlCochet/trajectory-transformer-halfcheetah-medium-v2 style configuration
- >>> model = TrajectoryTransformerModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "trajectory_transformer"
- keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {
- "hidden_size": "n_embd",
- "num_attention_heads": "n_head",
- "num_hidden_layers": "n_layer",
- }
-
- def __init__(
- self,
- vocab_size=100,
- action_weight=5,
- reward_weight=1,
- value_weight=1,
- block_size=249,
- action_dim=6,
- observation_dim=17,
- transition_dim=25,
- n_layer=4,
- n_head=4,
- n_embd=128,
- embd_pdrop=0.1,
- attn_pdrop=0.1,
- resid_pdrop=0.1,
- learning_rate=0.0006,
- max_position_embeddings=512,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- kaiming_initializer_range=1,
- use_cache=True,
- pad_token_id=1,
- bos_token_id=50256,
- eos_token_id=50256,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.action_weight = action_weight
- self.reward_weight = reward_weight
- self.value_weight = value_weight
- self.max_position_embeddings = max_position_embeddings
- self.block_size = block_size
- self.action_dim = action_dim
- self.observation_dim = observation_dim
- self.transition_dim = transition_dim
- self.learning_rate = learning_rate
- self.n_layer = n_layer
- self.n_head = n_head
- self.n_embd = n_embd
- self.embd_pdrop = embd_pdrop
- self.attn_pdrop = attn_pdrop
- self.resid_pdrop = resid_pdrop
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.kaiming_initializer_range = kaiming_initializer_range
- self.use_cache = use_cache
- super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
diff --git a/transformers/models/deprecated/trajectory_transformer/convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/deprecated/trajectory_transformer/convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index 622552fa78360826fc976d6f1d8c97fcc74a8a38..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/trajectory_transformer/convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The Trajectory Transformers paper authors and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TrajectoryTransformer pytorch checkpoint conversion"""
-
-import torch
-import trajectory.utils as utils
-
-from transformers import TrajectoryTransformerModel
-
-
-class Parser(utils.Parser):
- dataset: str = "halfcheetah-medium-expert-v2"
- config: str = "config.offline"
-
-
-def convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch(logbase, dataset, loadpath, epoch, device):
- """Converting Sequential blocks to ModuleList"""
-
- gpt, gpt_epoch = utils.load_model(logbase, dataset, loadpath, epoch=epoch, device=device)
- trajectory_transformer = TrajectoryTransformerModel(gpt.config)
-
- trajectory_transformer.tok_emb.load_state_dict(gpt.tok_emb.state_dict())
- trajectory_transformer.pos_emb = gpt.pos_emb
- trajectory_transformer.drop.load_state_dict(gpt.drop.state_dict())
- trajectory_transformer.ln_f.load_state_dict(gpt.ln_f.state_dict())
- trajectory_transformer.head.load_state_dict(gpt.head.state_dict())
-
- for i, block in enumerate(gpt.blocks):
- trajectory_transformer.blocks[i].ln1.load_state_dict(gpt.blocks[i].ln1.state_dict())
- trajectory_transformer.blocks[i].ln2.load_state_dict(gpt.blocks[i].ln2.state_dict())
- trajectory_transformer.blocks[i].attn.load_state_dict(gpt.blocks[i].attn.state_dict())
-
- trajectory_transformer.blocks[i].l1.load_state_dict(gpt.blocks[i].mlp[0].state_dict())
- trajectory_transformer.blocks[i].act.load_state_dict(gpt.blocks[i].mlp[1].state_dict())
- trajectory_transformer.blocks[i].l2.load_state_dict(gpt.blocks[i].mlp[2].state_dict())
- trajectory_transformer.blocks[i].drop.load_state_dict(gpt.blocks[i].mlp[3].state_dict())
-
- torch.save(trajectory_transformer.state_dict(), "pytorch_model.bin")
-
-
-if __name__ == "__main__":
- """
- To run this script you will need to install the original repository to run the original model. You can find it
- here: https://github.com/jannerm/trajectory-transformer From this repository code you can also download the
- original pytorch checkpoints.
-
- Run with the command:
-
- ```sh
- >>> python convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch.py --dataset
- ... --gpt_loadpath
- ```
- """
-
- args = Parser().parse_args("plan")
- convert_trajectory_transformer_original_pytorch_checkpoint_to_pytorch(
- args.logbase, args.dataset, args.gpt_loadpath, args.gpt_epoch, args.device
- )
diff --git a/transformers/models/deprecated/trajectory_transformer/modeling_trajectory_transformer.py b/transformers/models/deprecated/trajectory_transformer/modeling_trajectory_transformer.py
deleted file mode 100644
index 5c98aa45dc2739336b75e72167c739160cf5d11f..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/trajectory_transformer/modeling_trajectory_transformer.py
+++ /dev/null
@@ -1,606 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The Trajectory Transformers paper authors and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch TrajectoryTransformer model."""
-
-import math
-import os
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import functional as F
-
-from ....modeling_utils import PreTrainedModel
-from ....utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_trajectory_transformer import TrajectoryTransformerConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "CarlCochet/trajectory-transformer-halfcheetah-medium-v2"
-_CONFIG_FOR_DOC = "TrajectoryTransformerConfig"
-
-
-from .._archive_maps import TRAJECTORY_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-def load_tf_weights_in_trajectory_transformer(model, config, tf_checkpoint_path):
- """Load tf checkpoints in a pytorch model."""
- try:
- import re
-
- import numpy as np
- import tensorflow as tf
- except ImportError:
- logger.error(
- "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
- "https://www.tensorflow.org/install/ for installation instructions."
- )
- raise
- tf_path = os.path.abspath(tf_checkpoint_path)
- logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
- # Load weights from TF model
- init_vars = tf.train.list_variables(tf_path)
- names = []
- arrays = []
- for name, shape in init_vars:
- logger.info(f"Loading TF weight {name} with shape {shape}")
- array = tf.train.load_variable(tf_path, name)
- names.append(name)
- arrays.append(array)
-
- for name, array in zip(names, arrays):
- name = name.split("/")
- # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
- # which are not required for using pretrained model
- if any(
- n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
- for n in name
- ):
- logger.info(f"Skipping {'/'.join(name)}")
- continue
- pointer = model
- for m_name in name:
- if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
- scope_names = re.split(r"_(\d+)", m_name)
- else:
- scope_names = [m_name]
- if scope_names[0] == "kernel" or scope_names[0] == "gamma":
- pointer = getattr(pointer, "weight")
- elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
- pointer = getattr(pointer, "bias")
- elif scope_names[0] == "output_weights":
- pointer = getattr(pointer, "weight")
- elif scope_names[0] == "squad":
- pointer = getattr(pointer, "classifier")
- else:
- try:
- pointer = getattr(pointer, scope_names[0])
- except AttributeError:
- logger.info(f"Skipping {'/'.join(name)}")
- continue
- if len(scope_names) >= 2:
- num = int(scope_names[1])
- pointer = pointer[num]
- if m_name[-11:] == "_embeddings":
- pointer = getattr(pointer, "weight")
- elif m_name == "kernel":
- array = np.transpose(array)
- try:
- if pointer.shape != array.shape:
- raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
- except AssertionError as e:
- e.args += (pointer.shape, array.shape)
- raise
- logger.info(f"Initialize PyTorch weight {name}")
- pointer.data = torch.from_numpy(array)
- return model
-
-
-@dataclass
-class TrajectoryTransformerOutput(ModelOutput):
- """
- Base class for model's outputs that also contains a pooling of the last hidden states.
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
- Language modeling loss.
- logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
- past_key_values (`Tuple[Tuple[torch.Tensor]]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- Tuple of length `config.n_layers`, containing tuples of tensors of shape `(batch_size, num_heads,
- sequence_length, embed_size_per_head)`). Contains pre-computed hidden-states (key and values in the
- attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. GPT2Attentions weights after the attention softmax, used to compute the weighted average
- in the self-attention heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- logits: torch.FloatTensor = None
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-class TrajectoryTransformerPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = TrajectoryTransformerConfig
- load_tf_weights = load_tf_weights_in_trajectory_transformer
- base_model_prefix = "trajectory_transformer"
- main_input_name = "trajectories"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- if isinstance(module, (nn.Linear, nn.Embedding)):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if isinstance(module, nn.Linear) and module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
- elif isinstance(module, EinLinear):
- for i in range(module.n_models):
- nn.init.kaiming_uniform_(module.weight[i], a=math.sqrt(5) / self.config.kaiming_initializer_range)
- if module.bias is not None:
- fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight[i])
- bound = (1 / math.sqrt(fan_in)) * self.config.initializer_range
- nn.init.uniform_(module.bias[i], -bound, bound)
-
-
-TRAJECTORY_TRANSFORMER_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`TrajectoryTransformerConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-TRAJECTORY_TRANSFORMER_INPUTS_DOCSTRING = r"""
- Args:
- trajectories (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Batch of trajectories, where a trajectory is a sequence of states, actions and rewards.
- past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`, *optional*):
- Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
- their past given to this model should not be passed as `input_ids` as they have already been computed.
- targets (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Desired targets used to compute the loss.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class EinLinear(nn.Module):
- def __init__(self, n_models, in_features, out_features, bias):
- super().__init__()
- self.n_models = n_models
- self.out_features = out_features
- self.in_features = in_features
- self.weight = nn.Parameter(torch.Tensor(n_models, out_features, in_features))
- if bias:
- self.bias = nn.Parameter(torch.Tensor(n_models, out_features))
- else:
- self.register_parameter("bias", None)
-
- def reset_parameters(self):
- for i in range(self.n_models):
- nn.init.kaiming_uniform_(self.weight[i], a=math.sqrt(5))
- if self.bias is not None:
- fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight[i])
- bound = 1 / math.sqrt(fan_in)
- nn.init.uniform_(self.bias[i], -bound, bound)
-
- def forward(self, input):
- """
- Args:
- input (`torch.FloatTensor` of shape `(B, n_models, input_dim)`):
- The input to the layer.
- """
- # [ batch_size x n_models x output_dim ]
- output = torch.einsum("eoi,bei->beo", self.weight, input)
- if self.bias is not None:
- raise RuntimeError()
- return output
-
-
-class CausalSelfAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
-
- if config.n_embd % config.n_head != 0:
- raise ValueError(f"n_head ({config.n_head}) should be a divisor of n_embd ({config.n_embd})")
-
- # key, query, value projections for all heads
- self.key = nn.Linear(config.n_embd, config.n_embd)
- self.query = nn.Linear(config.n_embd, config.n_embd)
- self.value = nn.Linear(config.n_embd, config.n_embd)
-
- # regularization
- self.attn_drop = nn.Dropout(config.attn_pdrop)
- self.resid_drop = nn.Dropout(config.resid_pdrop)
-
- # output projection
- self.proj = nn.Linear(config.n_embd, config.n_embd)
-
- # causal mask to ensure that attention is only applied to the left in the input sequence
- self.register_buffer(
- "mask",
- torch.tril(torch.ones(config.block_size, config.block_size)).view(
- 1, 1, config.block_size, config.block_size
- ),
- persistent=False,
- )
-
- # mask previous value estimates
- joined_dim = config.observation_dim + config.action_dim + 2
- self.mask.squeeze()[:, joined_dim - 1 :: joined_dim] = 0
-
- self.n_head = config.n_head
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ):
- batch_size, sequence_length, embedding_dim = hidden_states.size()
-
- # calculate query, key, values for all heads in batch and move head forward to be the batch dim
- # [ batch_size x n_heads x sequence_length x head_dim ]
- key = (
- self.key(hidden_states)
- .view(batch_size, sequence_length, self.n_head, embedding_dim // self.n_head)
- .transpose(1, 2)
- )
- query = (
- self.query(hidden_states)
- .view(batch_size, sequence_length, self.n_head, embedding_dim // self.n_head)
- .transpose(1, 2)
- )
- value = (
- self.value(hidden_states)
- .view(batch_size, sequence_length, self.n_head, embedding_dim // self.n_head)
- .transpose(1, 2)
- )
-
- if layer_past is not None:
- past_key, past_value = layer_past
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- if use_cache is True:
- present = (key, value)
- else:
- present = None
-
- # causal self-attention
- # [ batch_size x n_heads x sequence_length x sequence_length ]
- attn_weights = (torch.matmul(query, key.transpose(-2, -1))) * (1.0 / math.sqrt(key.size(-1)))
- attn_weights = attn_weights.masked_fill(
- self.mask[:, :, :sequence_length, :sequence_length] == 0, torch.finfo(attn_weights.dtype).min
- )
- attn_weights = F.softmax(attn_weights, dim=-1)
- self._attn_map = attn_weights.clone()
- attn_weights = self.attn_drop(attn_weights)
-
- output = torch.matmul(attn_weights, value)
- # [ batch_size x sequence_length x embedding_dim ]
- # re-assemble all head outputs side by side
- output = output.transpose(1, 2).contiguous().view(batch_size, sequence_length, embedding_dim)
-
- # output projection
- output = self.resid_drop(self.proj(output))
-
- outputs = (output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
-
-class Block(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.ln1 = nn.LayerNorm(config.n_embd)
- self.ln2 = nn.LayerNorm(config.n_embd)
- self.attn = CausalSelfAttention(config)
-
- # MLP
- self.l1 = nn.Linear(config.n_embd, 4 * config.n_embd)
- self.act = nn.GELU()
- self.l2 = nn.Linear(4 * config.n_embd, config.n_embd)
- self.drop = nn.Dropout(config.resid_pdrop)
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ):
- residual = hidden_states
- hidden_states = self.ln1(hidden_states)
-
- attn_outputs = self.attn(
- hidden_states, layer_past=layer_past, use_cache=use_cache, output_attentions=output_attentions
- )
- attn_output = attn_outputs[0]
- outputs = attn_outputs[1:]
- hidden_states = attn_output + residual
-
- residual = hidden_states
- hidden_states = self.ln2(hidden_states)
- hidden_states = self.l1(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.l2(hidden_states)
- hidden_states = residual + self.drop(hidden_states)
-
- if use_cache:
- outputs = (hidden_states,) + outputs
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- return outputs
-
-
-@add_start_docstrings(
- "The bare TrajectoryTransformer Model transformer outputting raw hidden-states without any specific head on top.",
- TRAJECTORY_TRANSFORMER_START_DOCSTRING,
-)
-class TrajectoryTransformerModel(TrajectoryTransformerPreTrainedModel):
- """the full GPT language model, with a context size of block_size"""
-
- def __init__(self, config):
- super().__init__(config)
-
- # input embedding stem (+1 for stop token)
- self.tok_emb = nn.Embedding(config.vocab_size * config.transition_dim + 1, config.n_embd)
-
- self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
- self.drop = nn.Dropout(config.embd_pdrop)
- # transformer
- self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
- # decoder head
- self.ln_f = nn.LayerNorm(config.n_embd)
- self.head = EinLinear(config.transition_dim, config.n_embd, config.vocab_size + 1, bias=False)
-
- self.vocab_size = config.vocab_size
- self.stop_token = config.vocab_size * config.transition_dim
- self.block_size = config.block_size
-
- self.observation_dim = config.observation_dim
- self.action_dim = config.action_dim
- self.transition_dim = config.transition_dim
- self.embedding_dim = config.n_embd
-
- self.action_weight = config.action_weight
- self.reward_weight = config.reward_weight
- self.value_weight = config.value_weight
-
- self.gradient_checkpointing = False
-
- self.post_init()
-
- def get_block_size(self):
- return self.block_size
-
- def offset_tokens(self, trajectories):
- _, sequence_length = trajectories.shape
-
- n_states = int(np.ceil(sequence_length / self.transition_dim))
-
- offsets = torch.arange(self.transition_dim) * self.vocab_size
- offsets = offsets.repeat(n_states).to(trajectories.device)
-
- offset_trajectories = trajectories + offsets[:sequence_length]
- offset_trajectories[trajectories == self.vocab_size] = self.stop_token
- return offset_trajectories
-
- def pad_to_full_observation(self, hidden_states):
- batch_size, sequence_length, _ = hidden_states.shape
-
- n_pad = (self.transition_dim - sequence_length % self.transition_dim) % self.transition_dim
- padding = torch.zeros(batch_size, n_pad, self.embedding_dim, device=hidden_states.device)
-
- # [ batch_size x padded_sequence_length' x embedding_dim ]
- hidden_states_pad = torch.cat([hidden_states, padding], dim=1)
- hidden_states_pad = hidden_states_pad.view(-1, self.transition_dim, self.embedding_dim)
-
- return hidden_states_pad, n_pad
-
- @add_start_docstrings_to_model_forward(
- TRAJECTORY_TRANSFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length")
- )
- @replace_return_docstrings(output_type=TrajectoryTransformerOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- trajectories: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- targets: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], TrajectoryTransformerOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import TrajectoryTransformerModel
- >>> import torch
-
- >>> model = TrajectoryTransformerModel.from_pretrained(
- ... "CarlCochet/trajectory-transformer-halfcheetah-medium-v2"
- ... )
- >>> model.to(device)
- >>> model.eval()
-
- >>> observations_dim, action_dim, batch_size = 17, 6, 256
- >>> seq_length = observations_dim + action_dim + 1
-
- >>> trajectories = torch.LongTensor([np.random.permutation(self.seq_length) for _ in range(batch_size)]).to(
- ... device
- ... )
- >>> targets = torch.LongTensor([np.random.permutation(self.seq_length) for _ in range(batch_size)]).to(device)
-
- >>> outputs = model(
- ... trajectories,
- ... targets=targets,
- ... use_cache=True,
- ... output_attentions=True,
- ... output_hidden_states=True,
- ... return_dict=True,
- ... )
- ```
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- if past_key_values is None:
- past_key_values = tuple([None] * len(self.blocks))
-
- batch_size, sequence_length = trajectories.size()
-
- if sequence_length > self.block_size:
- raise ValueError("Cannot forward, model block size is exhausted.")
-
- offset_trajectories = self.offset_tokens(trajectories)
- # [ batch_size x sequence_length x embedding_dim ]
- # forward the GPT model
- token_embeddings = self.tok_emb(offset_trajectories) # each index maps to a (learnable) vector
- position_embeddings = self.pos_emb[:, :sequence_length, :] # each position maps to a (learnable) vector
-
- hidden_states = self.drop(token_embeddings + position_embeddings)
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- presents = () if use_cache else None
- all_self_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
-
- for i, (block, layer_past) in enumerate(zip(self.blocks, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- layer_past,
- use_cache,
- output_attentions,
- )
- else:
- outputs = block(hidden_states, layer_past, use_cache, output_attentions)
-
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
-
- # [ batch_size x sequence_length x embedding_dim ]
- hidden_state = self.ln_f(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- hidden_states_pad, n_pad = self.pad_to_full_observation(hidden_state)
-
- logits = self.head(hidden_states_pad)
- logits = logits.reshape(batch_size, sequence_length + n_pad, self.vocab_size + 1)
- logits = logits[:, :sequence_length]
-
- # if we are given some desired targets also calculate the loss
- if targets is not None:
- loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.view(-1), reduction="none")
- if self.action_weight != 1 or self.reward_weight != 1 or self.value_weight != 1:
- # make weights
- n_states = int(np.ceil(sequence_length / self.transition_dim))
- weights = torch.cat(
- [
- torch.ones(self.observation_dim, device=trajectories.device),
- torch.ones(self.action_dim, device=trajectories.device) * self.action_weight,
- torch.ones(1, device=trajectories.device) * self.reward_weight,
- torch.ones(1, device=trajectories.device) * self.value_weight,
- ]
- )
- weights = weights.repeat(n_states)
- weights = weights[1:].repeat(batch_size, 1)
- loss = loss * weights.view(-1)
- loss = (loss * attention_mask.view(-1)).mean()
- else:
- loss = None
-
- if not return_dict:
- return tuple(v for v in [loss, logits, presents, all_hidden_states, all_self_attentions] if v is not None)
-
- return TrajectoryTransformerOutput(
- loss=loss,
- logits=logits,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
diff --git a/transformers/models/deprecated/transfo_xl/__init__.py b/transformers/models/deprecated/transfo_xl/__init__.py
deleted file mode 100644
index f3674e19665ca74e1e6ee3ac92ca812e54580007..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/transfo_xl/__init__.py
+++ /dev/null
@@ -1,97 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available
-
-
-_import_structure = {
- "configuration_transfo_xl": ["TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP", "TransfoXLConfig"],
- "tokenization_transfo_xl": ["TransfoXLCorpus", "TransfoXLTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_transfo_xl"] = [
- "TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST",
- "AdaptiveEmbedding",
- "TransfoXLForSequenceClassification",
- "TransfoXLLMHeadModel",
- "TransfoXLModel",
- "TransfoXLPreTrainedModel",
- "load_tf_weights_in_transfo_xl",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_transfo_xl"] = [
- "TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFAdaptiveEmbedding",
- "TFTransfoXLForSequenceClassification",
- "TFTransfoXLLMHeadModel",
- "TFTransfoXLMainLayer",
- "TFTransfoXLModel",
- "TFTransfoXLPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
- from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_transfo_xl import (
- TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
- AdaptiveEmbedding,
- TransfoXLForSequenceClassification,
- TransfoXLLMHeadModel,
- TransfoXLModel,
- TransfoXLPreTrainedModel,
- load_tf_weights_in_transfo_xl,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_transfo_xl import (
- TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFAdaptiveEmbedding,
- TFTransfoXLForSequenceClassification,
- TFTransfoXLLMHeadModel,
- TFTransfoXLMainLayer,
- TFTransfoXLModel,
- TFTransfoXLPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deprecated/transfo_xl/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/transfo_xl/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index b5a5da076ea711607dcafcaac8e83b0f42bdde9a..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/transfo_xl/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/transfo_xl/__pycache__/configuration_transfo_xl.cpython-310.pyc b/transformers/models/deprecated/transfo_xl/__pycache__/configuration_transfo_xl.cpython-310.pyc
deleted file mode 100644
index a226e9a39df734ee1dec008e5fa98b2a3b110f2e..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/transfo_xl/__pycache__/configuration_transfo_xl.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/transfo_xl/__pycache__/convert_transfo_xl_original_tf_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/deprecated/transfo_xl/__pycache__/convert_transfo_xl_original_tf_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index b58110f749ca9d885a6e56ccf4b68e43c8d6f543..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/transfo_xl/__pycache__/convert_transfo_xl_original_tf_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/transfo_xl/__pycache__/modeling_tf_transfo_xl.cpython-310.pyc b/transformers/models/deprecated/transfo_xl/__pycache__/modeling_tf_transfo_xl.cpython-310.pyc
deleted file mode 100644
index 9c131ffe7d00928b1de09f498fa6d6976480cace..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/transfo_xl/__pycache__/modeling_tf_transfo_xl.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/transfo_xl/__pycache__/modeling_tf_transfo_xl_utilities.cpython-310.pyc b/transformers/models/deprecated/transfo_xl/__pycache__/modeling_tf_transfo_xl_utilities.cpython-310.pyc
deleted file mode 100644
index 5b1dee563edf19ee0526f0023a1624bf79262bc7..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/transfo_xl/__pycache__/modeling_tf_transfo_xl_utilities.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/transfo_xl/__pycache__/modeling_transfo_xl.cpython-310.pyc b/transformers/models/deprecated/transfo_xl/__pycache__/modeling_transfo_xl.cpython-310.pyc
deleted file mode 100644
index 2a008c8f396b762c7708b17502aa8ac760618021..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/transfo_xl/__pycache__/modeling_transfo_xl.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/transfo_xl/__pycache__/modeling_transfo_xl_utilities.cpython-310.pyc b/transformers/models/deprecated/transfo_xl/__pycache__/modeling_transfo_xl_utilities.cpython-310.pyc
deleted file mode 100644
index 28b90f92102abf62e5158d04be9bd7a7a4be99ca..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/transfo_xl/__pycache__/modeling_transfo_xl_utilities.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/transfo_xl/__pycache__/tokenization_transfo_xl.cpython-310.pyc b/transformers/models/deprecated/transfo_xl/__pycache__/tokenization_transfo_xl.cpython-310.pyc
deleted file mode 100644
index b3996ca86b3c83459cc1de33a9cd92e04811137b..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/transfo_xl/__pycache__/tokenization_transfo_xl.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/transfo_xl/configuration_transfo_xl.py b/transformers/models/deprecated/transfo_xl/configuration_transfo_xl.py
deleted file mode 100644
index 50bf94ae7ea3983d40dcf3d03c2e4e1027c23c80..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/transfo_xl/configuration_transfo_xl.py
+++ /dev/null
@@ -1,189 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Transformer XL configuration"""
-
-from ....configuration_utils import PretrainedConfig
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from .._archive_maps import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class TransfoXLConfig(PretrainedConfig):
- """
- This is the configuration class to store the configuration of a [`TransfoXLModel`] or a [`TFTransfoXLModel`]. It is
- used to instantiate a Transformer-XL model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the TransfoXL
- [transfo-xl/transfo-xl-wt103](https://huggingface.co/transfo-xl/transfo-xl-wt103) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 267735):
- Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`TransfoXLModel`] or [`TFTransfoXLModel`].
- cutoffs (`List[int]`, *optional*, defaults to `[20000, 40000, 200000]`):
- Cutoffs for the adaptive softmax.
- d_model (`int`, *optional*, defaults to 1024):
- Dimensionality of the model's hidden states.
- d_embed (`int`, *optional*, defaults to 1024):
- Dimensionality of the embeddings
- n_head (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- d_head (`int`, *optional*, defaults to 64):
- Dimensionality of the model's heads.
- d_inner (`int`, *optional*, defaults to 4096):
- Inner dimension in FF
- div_val (`int`, *optional*, defaults to 4):
- Divident value for adapative input and softmax
- pre_lnorm (`boolean`, *optional*, defaults to `False`):
- Whether or not to apply LayerNorm to the input instead of the output in the blocks.
- n_layer (`int`, *optional*, defaults to 18):
- Number of hidden layers in the Transformer encoder.
- mem_len (`int`, *optional*, defaults to 1600):
- Length of the retained previous heads.
- clamp_len (`int`, *optional*, defaults to 1000):
- Use the same pos embeddings after clamp_len.
- same_length (`boolean`, *optional*, defaults to `True`):
- Whether or not to use the same attn length for all tokens
- proj_share_all_but_first (`boolean`, *optional*, defaults to `True`):
- True to share all but first projs, False not to share.
- attn_type (`int`, *optional*, defaults to 0):
- Attention type. 0 for Transformer-XL, 1 for Shaw et al, 2 for Vaswani et al, 3 for Al Rfou et al.
- sample_softmax (`int`, *optional*, defaults to -1):
- Number of samples in the sampled softmax.
- adaptive (`boolean`, *optional*, defaults to `True`):
- Whether or not to use adaptive softmax.
- dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- dropatt (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- untie_r (`boolean`, *optional*, defaults to `True`):
- Whether ot not to untie relative position biases.
- init (`str`, *optional*, defaults to `"normal"`):
- Parameter initializer to use.
- init_range (`float`, *optional*, defaults to 0.01):
- Parameters initialized by U(-init_range, init_range).
- proj_init_std (`float`, *optional*, defaults to 0.01):
- Parameters initialized by N(0, init_std)
- init_std (`float`, *optional*, defaults to 0.02):
- Parameters initialized by N(0, init_std)
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
- The epsilon to use in the layer normalization layers
- eos_token_id (`int`, *optional*, defaults to 0):
- End of stream token id.
-
- Examples:
-
- ```python
- >>> from transformers import TransfoXLConfig, TransfoXLModel
-
- >>> # Initializing a Transformer XL configuration
- >>> configuration = TransfoXLConfig()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = TransfoXLModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "transfo-xl"
- keys_to_ignore_at_inference = ["mems"]
- attribute_map = {
- "n_token": "vocab_size",
- "hidden_size": "d_model",
- "num_attention_heads": "n_head",
- "num_hidden_layers": "n_layer",
- }
-
- def __init__(
- self,
- vocab_size=267735,
- cutoffs=[20000, 40000, 200000],
- d_model=1024,
- d_embed=1024,
- n_head=16,
- d_head=64,
- d_inner=4096,
- div_val=4,
- pre_lnorm=False,
- n_layer=18,
- mem_len=1600,
- clamp_len=1000,
- same_length=True,
- proj_share_all_but_first=True,
- attn_type=0,
- sample_softmax=-1,
- adaptive=True,
- dropout=0.1,
- dropatt=0.0,
- untie_r=True,
- init="normal",
- init_range=0.01,
- proj_init_std=0.01,
- init_std=0.02,
- layer_norm_epsilon=1e-5,
- eos_token_id=0,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.cutoffs = []
- self.cutoffs.extend(cutoffs)
- if proj_share_all_but_first:
- self.tie_projs = [False] + [True] * len(self.cutoffs)
- else:
- self.tie_projs = [False] + [False] * len(self.cutoffs)
- self.d_model = d_model
- self.d_embed = d_embed
- self.d_head = d_head
- self.d_inner = d_inner
- self.div_val = div_val
- self.pre_lnorm = pre_lnorm
- self.n_layer = n_layer
- self.n_head = n_head
- self.mem_len = mem_len
- self.same_length = same_length
- self.attn_type = attn_type
- self.clamp_len = clamp_len
- self.sample_softmax = sample_softmax
- self.adaptive = adaptive
- self.dropout = dropout
- self.dropatt = dropatt
- self.untie_r = untie_r
- self.init = init
- self.init_range = init_range
- self.proj_init_std = proj_init_std
- self.init_std = init_std
- self.layer_norm_epsilon = layer_norm_epsilon
- super().__init__(eos_token_id=eos_token_id, **kwargs)
-
- @property
- def max_position_embeddings(self):
- # Message copied from Transformer-XL documentation
- logger.info(f"The model {self.model_type} is one of the few models that has no sequence length limit.")
- return -1
-
- @max_position_embeddings.setter
- def max_position_embeddings(self, value):
- # Message copied from Transformer-XL documentation
- raise NotImplementedError(
- f"The model {self.model_type} is one of the few models that has no sequence length limit."
- )
diff --git a/transformers/models/deprecated/transfo_xl/convert_transfo_xl_original_tf_checkpoint_to_pytorch.py b/transformers/models/deprecated/transfo_xl/convert_transfo_xl_original_tf_checkpoint_to_pytorch.py
deleted file mode 100644
index d2693ac333b84b08769eb15a13a26dcf1a547267..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/transfo_xl/convert_transfo_xl_original_tf_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Transformer XL checkpoint and datasets."""
-
-
-import argparse
-import os
-import pickle
-import sys
-
-import torch
-
-from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
-from transformers.models.deprecated.transfo_xl import tokenization_transfo_xl as data_utils
-from transformers.models.deprecated.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
-from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging
-
-
-logging.set_verbosity_info()
-
-# We do this to be able to load python 2 datasets pickles
-# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
-data_utils.Vocab = data_utils.TransfoXLTokenizer
-data_utils.Corpus = data_utils.TransfoXLCorpus
-sys.modules["data_utils"] = data_utils
-sys.modules["vocabulary"] = data_utils
-
-
-def convert_transfo_xl_checkpoint_to_pytorch(
- tf_checkpoint_path, transfo_xl_config_file, pytorch_dump_folder_path, transfo_xl_dataset_file
-):
- if transfo_xl_dataset_file:
- # Convert a pre-processed corpus (see original TensorFlow repo)
- with open(transfo_xl_dataset_file, "rb") as fp:
- corpus = pickle.load(fp, encoding="latin1")
- # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
- pytorch_vocab_dump_path = pytorch_dump_folder_path + "/" + VOCAB_FILES_NAMES["pretrained_vocab_file"]
- print(f"Save vocabulary to {pytorch_vocab_dump_path}")
- corpus_vocab_dict = corpus.vocab.__dict__
- torch.save(corpus_vocab_dict, pytorch_vocab_dump_path)
-
- corpus_dict_no_vocab = corpus.__dict__
- corpus_dict_no_vocab.pop("vocab", None)
- pytorch_dataset_dump_path = pytorch_dump_folder_path + "/" + CORPUS_NAME
- print(f"Save dataset to {pytorch_dataset_dump_path}")
- torch.save(corpus_dict_no_vocab, pytorch_dataset_dump_path)
-
- if tf_checkpoint_path:
- # Convert a pre-trained TensorFlow model
- config_path = os.path.abspath(transfo_xl_config_file)
- tf_path = os.path.abspath(tf_checkpoint_path)
-
- print(f"Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.")
- # Initialise PyTorch model
- if transfo_xl_config_file == "":
- config = TransfoXLConfig()
- else:
- config = TransfoXLConfig.from_json_file(transfo_xl_config_file)
- print(f"Building PyTorch model from configuration: {config}")
- model = TransfoXLLMHeadModel(config)
-
- model = load_tf_weights_in_transfo_xl(model, config, tf_path)
- # Save pytorch-model
- pytorch_weights_dump_path = os.path.join(pytorch_dump_folder_path, WEIGHTS_NAME)
- pytorch_config_dump_path = os.path.join(pytorch_dump_folder_path, CONFIG_NAME)
- print(f"Save PyTorch model to {os.path.abspath(pytorch_weights_dump_path)}")
- torch.save(model.state_dict(), pytorch_weights_dump_path)
- print(f"Save configuration file to {os.path.abspath(pytorch_config_dump_path)}")
- with open(pytorch_config_dump_path, "w", encoding="utf-8") as f:
- f.write(config.to_json_string())
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- required=True,
- help="Path to the folder to store the PyTorch model or dataset/vocab.",
- )
- parser.add_argument(
- "--tf_checkpoint_path",
- default="",
- type=str,
- help="An optional path to a TensorFlow checkpoint path to be converted.",
- )
- parser.add_argument(
- "--transfo_xl_config_file",
- default="",
- type=str,
- help=(
- "An optional config json file corresponding to the pre-trained BERT model. \n"
- "This specifies the model architecture."
- ),
- )
- parser.add_argument(
- "--transfo_xl_dataset_file",
- default="",
- type=str,
- help="An optional dataset file to be converted in a vocabulary.",
- )
- args = parser.parse_args()
- convert_transfo_xl_checkpoint_to_pytorch(
- args.tf_checkpoint_path,
- args.transfo_xl_config_file,
- args.pytorch_dump_folder_path,
- args.transfo_xl_dataset_file,
- )
diff --git a/transformers/models/deprecated/transfo_xl/modeling_tf_transfo_xl.py b/transformers/models/deprecated/transfo_xl/modeling_tf_transfo_xl.py
deleted file mode 100644
index 27200a5d63f18b1d6457f1e303022f43b0d75d50..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/transfo_xl/modeling_tf_transfo_xl.py
+++ /dev/null
@@ -1,1122 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
- TF 2.0 Transformer XL model.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ....modeling_tf_utils import (
- TFModelInputType,
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ....tf_utils import shape_list, stable_softmax
-from ....utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_transfo_xl import TransfoXLConfig
-from .modeling_tf_transfo_xl_utilities import TFAdaptiveSoftmaxMask
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "transfo-xl/transfo-xl-wt103"
-_CONFIG_FOR_DOC = "TransfoXLConfig"
-
-
-from .._archive_maps import TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class TFPositionalEmbedding(keras.layers.Layer):
- def __init__(self, demb, **kwargs):
- super().__init__(**kwargs)
-
- self.inv_freq = 1 / (10000 ** (tf.range(0, demb, 2.0) / demb))
-
- def call(self, pos_seq, bsz=None):
- self.inv_freq = tf.cast(self.inv_freq, dtype=pos_seq.dtype)
- sinusoid_inp = tf.einsum("i,j->ij", pos_seq, self.inv_freq)
- pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], -1)
-
- if bsz is not None:
- return tf.tile(pos_emb[:, None, :], [1, bsz, 1])
- else:
- return pos_emb[:, None, :]
-
-
-class TFPositionwiseFF(keras.layers.Layer):
- def __init__(self, d_model, d_inner, dropout, pre_lnorm=False, layer_norm_epsilon=1e-5, init_std=0.02, **kwargs):
- super().__init__(**kwargs)
-
- self.d_model = d_model
- self.d_inner = d_inner
- self.dropout = dropout
-
- self.layer_1 = keras.layers.Dense(
- d_inner, kernel_initializer=get_initializer(init_std), activation=tf.nn.relu, name="CoreNet_._0"
- )
- self.drop_1 = keras.layers.Dropout(dropout)
- self.layer_2 = keras.layers.Dense(d_model, kernel_initializer=get_initializer(init_std), name="CoreNet_._3")
- self.drop_2 = keras.layers.Dropout(dropout)
-
- self.layer_norm = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layer_norm")
-
- self.pre_lnorm = pre_lnorm
-
- def call(self, inp, training=False):
- if self.pre_lnorm:
- # layer normalization + positionwise feed-forward
- core_out = self.layer_norm(inp)
- core_out = self.layer_1(core_out)
- core_out = self.drop_1(core_out, training=training)
- core_out = self.layer_2(core_out)
- core_out = self.drop_2(core_out, training=training)
-
- # residual connection
- output = core_out + inp
- else:
- # positionwise feed-forward
- core_out = self.layer_1(inp)
- core_out = self.drop_1(core_out, training=training)
- core_out = self.layer_2(core_out)
- core_out = self.drop_2(core_out, training=training)
-
- # residual connection + layer normalization
- output = self.layer_norm(inp + core_out)
-
- return output
-
-
-class TFRelPartialLearnableMultiHeadAttn(keras.layers.Layer):
- def __init__(
- self,
- n_head,
- d_model,
- d_head,
- dropout,
- dropatt=0.0,
- pre_lnorm=False,
- r_r_bias=None,
- r_w_bias=None,
- layer_norm_epsilon=1e-5,
- init_std=0.02,
- output_attentions=False,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.n_head = n_head
- self.d_model = d_model
- self.d_head = d_head
- self.dropout = dropout
- self.output_attentions = output_attentions
-
- self.qkv_net = keras.layers.Dense(
- 3 * n_head * d_head, kernel_initializer=get_initializer(init_std), use_bias=False, name="qkv_net"
- )
-
- self.drop = keras.layers.Dropout(dropout)
- self.dropatt = keras.layers.Dropout(dropatt)
- self.o_net = keras.layers.Dense(
- d_model, kernel_initializer=get_initializer(init_std), use_bias=False, name="o_net"
- )
-
- self.layer_norm = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layer_norm")
-
- self.scale = 1 / (d_head**0.5)
-
- self.pre_lnorm = pre_lnorm
-
- if r_r_bias is not None and r_w_bias is not None: # Biases are shared
- self.r_r_bias = r_r_bias
- self.r_w_bias = r_w_bias
- else:
- self.r_r_bias = None
- self.r_w_bias = None
-
- self.r_net = keras.layers.Dense(
- self.n_head * self.d_head, kernel_initializer=get_initializer(init_std), use_bias=False, name="r_net"
- )
-
- def build(self, input_shape):
- if self.r_r_bias is None or self.r_w_bias is None: # Biases are not shared
- self.r_r_bias = self.add_weight(
- shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_r_bias"
- )
- self.r_w_bias = self.add_weight(
- shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_w_bias"
- )
- super().build(input_shape)
-
- def _rel_shift(self, x):
- x_size = shape_list(x)
-
- x = tf.pad(x, [[0, 0], [1, 0], [0, 0], [0, 0]])
- x = tf.reshape(x, [x_size[1] + 1, x_size[0], x_size[2], x_size[3]])
- x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])
- x = tf.reshape(x, x_size)
-
- return x
-
- def call(self, w, r, attn_mask, mems, head_mask, output_attentions, training=False):
- qlen, rlen, bsz = shape_list(w)[0], shape_list(r)[0], shape_list(w)[1]
-
- if mems is not None:
- mems = tf.cast(mems, dtype=w.dtype)
- cat = tf.concat([mems, w], 0)
- if self.pre_lnorm:
- w_heads = self.qkv_net(self.layer_norm(cat))
- else:
- w_heads = self.qkv_net(cat)
- r_head_k = self.r_net(r)
-
- w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, axis=-1)
- w_head_q = w_head_q[-qlen:]
- else:
- if self.pre_lnorm:
- w_heads = self.qkv_net(self.layer_norm(w))
- else:
- w_heads = self.qkv_net(w)
- r_head_k = self.r_net(r)
-
- w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, axis=-1)
-
- klen = shape_list(w_head_k)[0]
-
- w_head_q = tf.reshape(w_head_q, (qlen, bsz, self.n_head, self.d_head)) # qlen x bsz x n_head x d_head
- w_head_k = tf.reshape(w_head_k, (klen, bsz, self.n_head, self.d_head)) # qlen x bsz x n_head x d_head
- w_head_v = tf.reshape(w_head_v, (klen, bsz, self.n_head, self.d_head)) # qlen x bsz x n_head x d_head
-
- r_head_k = tf.reshape(r_head_k, (rlen, self.n_head, self.d_head)) # qlen x n_head x d_head
-
- # compute attention score
- rw_head_q = w_head_q + self.r_w_bias # qlen x bsz x n_head x d_head
- AC = tf.einsum("ibnd,jbnd->ijbn", rw_head_q, w_head_k) # qlen x klen x bsz x n_head
-
- rr_head_q = w_head_q + self.r_r_bias
- BD = tf.einsum("ibnd,jnd->ijbn", rr_head_q, r_head_k) # qlen x klen x bsz x n_head
- BD = self._rel_shift(BD)
-
- # [qlen x klen x bsz x n_head]
- attn_score = AC + BD
- attn_score = attn_score * self.scale
-
- # compute attention probability
- if attn_mask is not None:
- attn_mask_t = attn_mask[:, :, None, None]
- attn_mask_t = tf.cast(attn_mask_t, dtype=attn_score.dtype)
- attn_score = attn_score * (1.0 - attn_mask_t) - 1e30 * attn_mask_t
-
- # [qlen x klen x bsz x n_head]
- attn_prob = stable_softmax(attn_score, axis=1)
- attn_prob = self.dropatt(attn_prob, training=training)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_prob = attn_prob * head_mask
-
- # compute attention vector
- attn_vec = tf.einsum("ijbn,jbnd->ibnd", attn_prob, w_head_v)
-
- # [qlen x bsz x n_head x d_head]
- attn_vec_sizes = shape_list(attn_vec)
- attn_vec = tf.reshape(attn_vec, (attn_vec_sizes[0], attn_vec_sizes[1], self.n_head * self.d_head))
-
- # linear projection
- attn_out = self.o_net(attn_vec)
- attn_out = self.drop(attn_out, training=training)
-
- if self.pre_lnorm:
- # residual connection
- outputs = [w + attn_out]
- else:
- # residual connection + layer normalization
- outputs = [self.layer_norm(w + attn_out)]
-
- if output_attentions:
- outputs.append(attn_prob)
-
- return outputs
-
-
-class TFRelPartialLearnableDecoderLayer(keras.layers.Layer):
- def __init__(
- self,
- n_head,
- d_model,
- d_head,
- d_inner,
- dropout,
- dropatt=0.0,
- pre_lnorm=False,
- r_w_bias=None,
- r_r_bias=None,
- layer_norm_epsilon=1e-5,
- init_std=0.02,
- output_attentions=False,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.dec_attn = TFRelPartialLearnableMultiHeadAttn(
- n_head,
- d_model,
- d_head,
- dropout,
- dropatt=dropatt,
- pre_lnorm=pre_lnorm,
- r_w_bias=r_w_bias,
- r_r_bias=r_r_bias,
- init_std=init_std,
- layer_norm_epsilon=layer_norm_epsilon,
- output_attentions=output_attentions,
- name="dec_attn",
- )
- self.pos_ff = TFPositionwiseFF(
- d_model,
- d_inner,
- dropout,
- pre_lnorm=pre_lnorm,
- init_std=init_std,
- layer_norm_epsilon=layer_norm_epsilon,
- name="pos_ff",
- )
-
- def call(self, dec_inp, r, dec_attn_mask, mems, head_mask, output_attentions, training=False):
- attn_outputs = self.dec_attn(dec_inp, r, dec_attn_mask, mems, head_mask, output_attentions, training=training)
- ff_output = self.pos_ff(attn_outputs[0], training=training)
-
- outputs = [ff_output] + attn_outputs[1:]
-
- return outputs
-
-
-class TFTransfoEmbeddings(keras.layers.Layer):
- def __init__(self, vocab_size, emb_size, init_std, **kwargs):
- super().__init__(**kwargs)
-
- self.vocab_size = vocab_size
- self.emb_size = emb_size
- self.init_std = init_std
-
- def build(self, input_shape):
- self.weight = self.add_weight(
- shape=(self.vocab_size, self.emb_size),
- initializer=get_initializer(self.init_std),
- name="embeddings",
- )
-
- super().build(input_shape)
-
- def call(self, inputs):
- return tf.gather(self.weight, inputs)
-
-
-class TFAdaptiveEmbedding(keras.layers.Layer):
- def __init__(self, n_token, d_embed, d_proj, cutoffs, div_val=1, init_std=0.02, sample_softmax=False, **kwargs):
- super().__init__(**kwargs)
-
- self.n_token = n_token
- self.d_embed = d_embed
- self.init_std = init_std
-
- self.cutoffs = cutoffs + [n_token]
- self.div_val = div_val
- self.d_proj = d_proj
-
- self.emb_scale = d_proj**0.5
-
- self.cutoff_ends = [0] + self.cutoffs
-
- self.emb_layers = []
- self.emb_projs = []
-
- if div_val == 1:
- raise NotImplementedError # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
- else:
- for i in range(len(self.cutoffs)):
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
- d_emb_i = d_embed // (div_val**i)
- self.emb_layers.append(
- TFTransfoEmbeddings(
- r_idx - l_idx,
- d_emb_i,
- init_std,
- name=f"emb_layers_._{i}",
- )
- )
-
- def build(self, input_shape):
- for i in range(len(self.cutoffs)):
- d_emb_i = self.d_embed // (self.div_val**i)
- self.emb_projs.append(
- self.add_weight(
- shape=(d_emb_i, self.d_proj),
- initializer=get_initializer(self.init_std),
- trainable=True,
- name=f"emb_projs_._{i}",
- )
- )
-
- super().build(input_shape)
-
- def call(self, inp):
- if self.div_val == 1:
- raise NotImplementedError # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
- else:
- inp_flat = tf.reshape(inp, (-1,))
- emb_flat = tf.zeros([shape_list(inp_flat)[0], self.d_proj])
- for i in range(len(self.cutoffs)):
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
-
- mask_i = (inp_flat >= l_idx) & (inp_flat < r_idx)
-
- inp_i = tf.boolean_mask(inp_flat, mask_i) - l_idx
- emb_i = self.emb_layers[i](inp_i)
- emb_i = tf.einsum("id,de->ie", emb_i, self.emb_projs[i])
-
- mask_idx = tf.where(mask_i)
- scatter = tf.scatter_nd(mask_idx, emb_i, shape_list(emb_flat))
- emb_flat = tf.cast(emb_flat, dtype=scatter.dtype)
- emb_flat += scatter
-
- embed_shape = shape_list(inp) + [self.d_proj]
- embed = tf.reshape(emb_flat, embed_shape)
-
- embed *= self.emb_scale
-
- return embed
-
-
-@keras_serializable
-class TFTransfoXLMainLayer(keras.layers.Layer):
- config_class = TransfoXLConfig
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.output_hidden_states = config.output_hidden_states
- self.output_attentions = config.output_attentions
- self.return_dict = config.use_return_dict
-
- self.n_token = config.vocab_size
-
- self.d_embed = config.d_embed
- self.d_model = config.d_model
- self.n_head = config.n_head
- self.d_head = config.d_head
- self.untie_r = config.untie_r
-
- self.word_emb = TFAdaptiveEmbedding(
- config.vocab_size,
- config.d_embed,
- config.d_model,
- config.cutoffs,
- div_val=config.div_val,
- init_std=config.init_std,
- name="word_emb",
- )
-
- self.drop = keras.layers.Dropout(config.dropout)
-
- self.n_layer = config.n_layer
- self.mem_len = config.mem_len
- self.attn_type = config.attn_type
-
- self.layers = []
- if config.attn_type == 0: # the default attention
- for i in range(config.n_layer):
- self.layers.append(
- TFRelPartialLearnableDecoderLayer(
- config.n_head,
- config.d_model,
- config.d_head,
- config.d_inner,
- config.dropout,
- dropatt=config.dropatt,
- pre_lnorm=config.pre_lnorm,
- r_w_bias=None if self.untie_r else self.r_w_bias,
- r_r_bias=None if self.untie_r else self.r_r_bias,
- layer_norm_epsilon=config.layer_norm_epsilon,
- init_std=config.init_std,
- output_attentions=self.output_attentions,
- name=f"layers_._{i}",
- )
- )
- else: # learnable embeddings and absolute embeddings
- raise NotImplementedError # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
-
- self.same_length = config.same_length
- self.clamp_len = config.clamp_len
-
- if self.attn_type == 0: # default attention
- self.pos_emb = TFPositionalEmbedding(self.d_model, name="pos_emb")
- else: # learnable embeddings and absolute embeddings
- raise NotImplementedError # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
-
- def build(self, input_shape):
- if not self.untie_r:
- self.r_w_bias = self.add_weight(
- shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_w_bias"
- )
- self.r_r_bias = self.add_weight(
- shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_r_bias"
- )
- super().build(input_shape)
-
- def get_input_embeddings(self):
- return self.word_emb
-
- def set_input_embeddings(self, value):
- raise NotImplementedError
-
- def backward_compatible(self):
- self.sample_softmax = -1
-
- def reset_memory_length(self, mem_len):
- self.mem_len = mem_len
-
- def _prune_heads(self, heads):
- raise NotImplementedError
-
- def init_mems(self, bsz):
- if self.mem_len > 0:
- mems = []
- for i in range(self.n_layer):
- empty = tf.zeros([self.mem_len, bsz, self.d_model])
- mems.append(empty)
-
- return mems
- else:
- return None
-
- def _update_mems(self, hids, mems, mlen, qlen):
- # does not deal with None
- if mems is None:
- return None
-
- # mems is not None
- assert len(hids) == len(mems), "len(hids) != len(mems)"
-
- # There are `mlen + qlen` steps that can be cached into mems
- new_mems = []
- end_idx = mlen + tf.math.maximum(0, qlen)
- beg_idx = tf.math.maximum(0, end_idx - tf.convert_to_tensor(self.mem_len))
- for i in range(len(hids)):
- mems[i] = tf.cast(mems[i], dtype=hids[i].dtype)
- cat = tf.concat([mems[i], hids[i]], axis=0)
- tf.stop_gradient(cat)
- new_mems.append(cat[beg_idx:end_idx])
-
- return new_mems
-
- @unpack_inputs
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- mems: List[tf.Tensor] | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ):
- # the original code for Transformer-XL used shapes [len, bsz] but we want a unified interface in the library
- # so we transpose here from shape [bsz, len] to shape [len, bsz]
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_ids = tf.transpose(input_ids, perm=(1, 0))
- qlen, bsz = shape_list(input_ids)
- elif inputs_embeds is not None:
- inputs_embeds = tf.transpose(inputs_embeds, perm=(1, 0, 2))
- qlen, bsz = shape_list(inputs_embeds)[:2]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if mems is None:
- mems = self.init_mems(bsz)
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] (a head_mask for each layer)
- # and head_mask is converted to shape [num_hidden_layers x qlen x klen x bsz x n_head]
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.n_layer
-
- if inputs_embeds is not None:
- word_emb = inputs_embeds
- else:
- word_emb = self.word_emb(input_ids)
-
- mlen = shape_list(mems[0])[0] if mems is not None else 0
- klen = mlen + qlen
-
- # Compute decoder attention mask
- all_ones = tf.ones([qlen, klen], dtype=tf.int32)
- upper_mask = 1 - tf.linalg.band_part(tf.ones([qlen, klen], dtype=tf.int32), -1, mlen)
- if self.same_length:
- mask_len = klen - self.mem_len
- mask_shift_len = qlen - tf.nn.relu(mask_len) # Lazy clamping of negatives to zero
-
- # Use an indicator variable instead of a conditional to keep the compiler happy
- lower_mask = tf.linalg.band_part(all_ones, -1, 0) - (
- tf.linalg.band_part(all_ones, mask_shift_len - 1, 0) * tf.cast(mask_shift_len != 0, tf.int32)
- )
- dec_attn_mask = upper_mask + lower_mask
- else:
- dec_attn_mask = upper_mask
-
- hids = []
- attentions = [] if output_attentions else None
- if self.attn_type == 0: # default
- pos_seq = tf.range(klen - 1, -1, -1.0)
- if self.clamp_len > 0:
- pos_seq = tf.minimum(pos_seq, self.clamp_len)
- pos_emb = self.pos_emb(pos_seq)
-
- core_out = self.drop(word_emb, training=training)
- pos_emb = self.drop(pos_emb, training=training)
-
- for i, layer in enumerate(self.layers):
- hids.append(core_out)
- mems_i = None if mems is None else mems[i]
- layer_outputs = layer(
- core_out,
- pos_emb,
- dec_attn_mask,
- mems_i,
- head_mask[i],
- output_attentions,
- training=training,
- )
- core_out = layer_outputs[0]
- if output_attentions:
- attentions.append(layer_outputs[1])
- else: # learnable embeddings and absolute embeddings
- raise NotImplementedError # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
-
- core_out = self.drop(core_out, training=training)
-
- new_mems = self._update_mems(hids, mems, mlen, qlen)
-
- # We transpose back here to shape [bsz, len, hidden_dim]
- core_out = tf.transpose(core_out, perm=(1, 0, 2))
-
- if output_hidden_states:
- # Transpose to library standard shape [bsz, len, hidden_dim] and add last layer
- hids = tuple(tf.transpose(t, perm=(1, 0, 2)) for t in hids)
- hids = hids + (core_out,)
- else:
- hids = None
- if output_attentions:
- # Transpose to library standard shape [bsz, n_heads, query_seq_len, key_seq_len]
- attentions = tuple(tf.transpose(t, perm=(2, 3, 0, 1)) for t in attentions)
-
- if not return_dict:
- return tuple(v for v in [core_out, new_mems, hids, attentions] if v is not None)
-
- return TFTransfoXLModelOutput(
- last_hidden_state=core_out,
- mems=new_mems,
- hidden_states=hids,
- attentions=attentions,
- )
-
-
-class TFTransfoXLPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = TransfoXLConfig
- base_model_prefix = "transformer"
-
-
-@dataclass
-class TFTransfoXLModelOutput(ModelOutput):
- """
- Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
-
- Args:
- last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- mems (`List[tf.Tensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks). Can be used (see `mems`
- input) to speed up sequential decoding. The token ids which have their past given to this model should not
- be passed as input ids as they have already been computed.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- last_hidden_state: tf.Tensor = None
- mems: List[tf.Tensor] = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-@dataclass
-class TFTransfoXLLMHeadModelOutput(ModelOutput):
- """
- Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
-
- Args:
- losses (`tf.Tensor` of shape *(batch_size, sequence_length-1)*, *optional*, returned when `labels` is provided):
- Language modeling losses (not reduced).
- prediction_scores (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token after SoftMax).
- mems (`List[tf.Tensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks). Can be used (see `mems`
- input) to speed up sequential decoding. The token ids which have their past given to this model should not
- be passed as input ids as they have already been computed.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- prediction_scores: tf.Tensor = None
- mems: List[tf.Tensor] = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-@dataclass
-class TFTransfoXLSequenceClassifierOutputWithPast(ModelOutput):
- """
- Base class for outputs of sentence classification models.
-
- Args:
- loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
- Classification (or regression if config.num_labels==1) loss.
- logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):
- Classification (or regression if config.num_labels==1) scores (before SoftMax).
- mems (`List[tf.Tensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks). Can be used (see `mems`
- input) to speed up sequential decoding. The token ids which have their past given to this model should not
- be passed as input ids as they have already been computed.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- loss: tf.Tensor | None = None
- logits: tf.Tensor = None
- mems: List[tf.Tensor] = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-TRANSFO_XL_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`TransfoXLConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-TRANSFO_XL_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- mems (`List[tf.Tensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model (see
- `mems` output below). Can be used to speed up sequential decoding. The token ids which have their mems
- given to this model should not be passed as `input_ids` as they have already been computed.
- head_mask (`tf.Tensor` or `Numpy array` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
- inputs_embeds (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
- TRANSFO_XL_START_DOCSTRING,
-)
-class TFTransfoXLModel(TFTransfoXLPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFTransfoXLMainLayer(config, name="transformer")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(TRANSFO_XL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTransfoXLModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- mems: List[tf.Tensor] | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: bool | None = None,
- output_hidden_states: bool | None = None,
- return_dict: bool | None = None,
- training: bool = False,
- ) -> TFTransfoXLModelOutput | Tuple[tf.Tensor]:
- outputs = self.transformer(
- input_ids=input_ids,
- mems=mems,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
-
-@add_start_docstrings(
- """
- The Transformer-XL Model with a language modeling head on top (adaptive softmax with weights tied to the adaptive
- input embeddings)
- """,
- TRANSFO_XL_START_DOCSTRING,
-)
-class TFTransfoXLLMHeadModel(TFTransfoXLPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.transformer = TFTransfoXLMainLayer(config, name="transformer")
- self.sample_softmax = config.sample_softmax
- assert self.sample_softmax <= 0, (
- "Sampling from the softmax is not implemented yet. Please look at issue: #3310:"
- " https://github.com/huggingface/transformers/issues/3310"
- )
-
- self.crit = TFAdaptiveSoftmaxMask(
- config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val, name="crit"
- )
-
- def _resize_token_embeddings(self, new_num_tokens):
- raise NotImplementedError()
-
- def get_output_embeddings(self):
- """Double-check if you are using adaptive softmax."""
- if len(self.crit.out_layers) > 0:
- return self.crit.out_layers[-1]
- return None
-
- def reset_memory_length(self, mem_len):
- self.transformer.reset_memory_length(mem_len)
-
- def init_mems(self, bsz):
- return self.transformer.init_mems(bsz)
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(TRANSFO_XL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTransfoXLLMHeadModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- mems: List[tf.Tensor] | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: bool | None = None,
- output_hidden_states: bool | None = None,
- return_dict: bool | None = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> TFTransfoXLLMHeadModelOutput | Tuple[tf.Tensor]:
- if input_ids is not None:
- bsz, tgt_len = shape_list(input_ids)[:2]
- else:
- bsz, tgt_len = shape_list(inputs_embeds)[:2]
-
- transformer_outputs = self.transformer(
- input_ids,
- mems,
- head_mask,
- inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict,
- training=training,
- )
-
- last_hidden = transformer_outputs[0]
- pred_hid = last_hidden[:, -tgt_len:]
-
- softmax_output = self.crit(pred_hid, labels, training=training)
- prediction_scores = softmax_output if labels is None else ()
-
- if not return_dict:
- return (prediction_scores,) + transformer_outputs[1:]
-
- return TFTransfoXLLMHeadModelOutput(
- prediction_scores=prediction_scores,
- mems=transformer_outputs.mems,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **model_kwargs):
- inputs = {}
-
- # if past is defined in model kwargs then use it for faster decoding
- if past_key_values:
- input_ids = tf.expand_dims(input_ids[:, -1], axis=-1)
- else:
- input_ids = input_ids
-
- return inputs
-
- # Adapted from the torch tie_weights function
- def tf_to_pt_weight_rename(self, tf_weight):
- if self.config.tie_word_embeddings and "crit.out_layers" in tf_weight:
- return tf_weight, tf_weight.replace("crit.out_layers", "transformer.word_emb.emb_layers")
- elif self.config.tie_projs and "crit.out_projs" in tf_weight:
- for i, tie_proj in enumerate(self.config.tie_projs):
- if tie_proj and self.config.div_val == 1 and self.config.d_model != self.config.d_embed:
- # self.crit.out_projs[i] = self.transformer.word_emb.emb_projs[0]
- return tf_weight, tf_weight.replace(f"crit.out_projs.{i}", "transformer.word_emb.emb_projs.0")
- elif tie_proj and self.config.div_val != 1:
- # self.crit.out_projs[i] = self.transformer.word_emb.emb_projs[i]
- return tf_weight, tf_weight.replace("crit.out_projs", "transformer.word_emb.emb_projs")
- else:
- return (tf_weight,)
-
-
-@add_start_docstrings(
- """
- The Transfo XL Model transformer with a sequence classification head on top (linear layer).
-
- [`TFTransfoXLForSequenceClassification`] uses the last token in order to do the classification, as other causal
- models (e.g. GPT-1,GPT-2) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- TRANSFO_XL_START_DOCSTRING,
-)
-class TFTransfoXLForSequenceClassification(TFTransfoXLPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
- self.score = keras.layers.Dense(
- config.num_labels,
- kernel_initializer=get_initializer(config.init_range),
- name="score",
- use_bias=False,
- )
- self.transformer = TFTransfoXLMainLayer(config, name="transformer")
-
- def get_output_embeddings(self):
- # Remove after transformers v4.32. Fix this model's `test_model_common_attributes` test too.
- logger.warning(
- "Sequence classification models do not have output embeddings. `.get_output_embeddings` will be removed "
- "in transformers v4.32."
- )
- return self.transformer.word_emb
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(TRANSFO_XL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTransfoXLSequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- mems: List[tf.Tensor] | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFTransfoXLSequenceClassifierOutputWithPast]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the cross entropy classification loss. Indices should be in `[0, ...,
- config.vocab_size - 1]`.
- """
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- mems=mems,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
- in_logits = None
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- sequence_lengths = (
- tf.argmax(tf.cast(tf.math.equal(input_ids, self.config.pad_token_id), input_ids.dtype), axis=-1)
- - 1
- )
- sequence_lengths = tf.where(sequence_lengths >= 0, sequence_lengths, input_ids.shape[-1] - 1)
- in_logits = tf.gather(logits, sequence_lengths, batch_dims=1, axis=1)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
- loss = None
-
- if labels is not None:
- if input_ids is not None:
- batch_size, sequence_length = shape_list(input_ids)[:2]
- else:
- batch_size, sequence_length = shape_list(inputs_embeds)[:2]
- assert (
- self.config.pad_token_id is not None or batch_size == 1
- ), "Cannot handle batch sizes > 1 if no padding token is defined."
-
- if not tf.is_tensor(sequence_lengths):
- in_logits = logits[0:batch_size, sequence_lengths]
-
- loss = self.hf_compute_loss(tf.reshape(labels, [-1, 1]), tf.reshape(in_logits, [-1, self.num_labels]))
-
- pooled_logits = in_logits if in_logits is not None else logits
-
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFTransfoXLSequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- mems=transformer_outputs.mems,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
diff --git a/transformers/models/deprecated/transfo_xl/modeling_tf_transfo_xl_utilities.py b/transformers/models/deprecated/transfo_xl/modeling_tf_transfo_xl_utilities.py
deleted file mode 100644
index ed1488d5595cb8f36eb540992fb4ca46534a60fb..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/transfo_xl/modeling_tf_transfo_xl_utilities.py
+++ /dev/null
@@ -1,179 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
- A TF 2.0 Adaptive Softmax for Transformer XL model.
-"""
-
-
-import tensorflow as tf
-
-from ....modeling_tf_utils import keras
-from ....tf_utils import shape_list
-
-
-class TFAdaptiveSoftmaxMask(keras.layers.Layer):
- def __init__(self, vocab_size, d_embed, d_proj, cutoffs, div_val=1, keep_order=False, **kwargs):
- super().__init__(**kwargs)
-
- self.vocab_size = vocab_size
- self.d_embed = d_embed
- self.d_proj = d_proj
-
- self.cutoffs = cutoffs + [vocab_size]
- self.cutoff_ends = [0] + self.cutoffs
- self.div_val = div_val
-
- self.shortlist_size = self.cutoffs[0]
- self.n_clusters = len(self.cutoffs) - 1
- self.head_size = self.shortlist_size + self.n_clusters
- self.keep_order = keep_order
-
- self.out_layers = []
- self.out_projs = []
-
- def build(self, input_shape):
- if self.n_clusters > 0:
- self.cluster_weight = self.add_weight(
- shape=(self.n_clusters, self.d_embed), initializer="zeros", trainable=True, name="cluster_weight"
- )
- self.cluster_bias = self.add_weight(
- shape=(self.n_clusters,), initializer="zeros", trainable=True, name="cluster_bias"
- )
-
- if self.div_val == 1:
- for i in range(len(self.cutoffs)):
- if self.d_proj != self.d_embed:
- weight = self.add_weight(
- shape=(self.d_embed, self.d_proj),
- initializer="zeros",
- trainable=True,
- name=f"out_projs_._{i}",
- )
- self.out_projs.append(weight)
- else:
- self.out_projs.append(None)
- weight = self.add_weight(
- shape=(self.vocab_size, self.d_embed),
- initializer="zeros",
- trainable=True,
- name=f"out_layers_._{i}_._weight",
- )
- bias = self.add_weight(
- shape=(self.vocab_size,),
- initializer="zeros",
- trainable=True,
- name=f"out_layers_._{i}_._bias",
- )
- self.out_layers.append((weight, bias))
- else:
- for i in range(len(self.cutoffs)):
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
- d_emb_i = self.d_embed // (self.div_val**i)
-
- weight = self.add_weight(
- shape=(d_emb_i, self.d_proj), initializer="zeros", trainable=True, name=f"out_projs_._{i}"
- )
- self.out_projs.append(weight)
- weight = self.add_weight(
- shape=(r_idx - l_idx, d_emb_i),
- initializer="zeros",
- trainable=True,
- name=f"out_layers_._{i}_._weight",
- )
- bias = self.add_weight(
- shape=(r_idx - l_idx,),
- initializer="zeros",
- trainable=True,
- name=f"out_layers_._{i}_._bias",
- )
- self.out_layers.append((weight, bias))
- super().build(input_shape)
-
- @staticmethod
- def _logit(x, W, b, proj=None):
- y = x
- if proj is not None:
- y = tf.einsum("ibd,ed->ibe", y, proj)
- return tf.einsum("ibd,nd->ibn", y, W) + b
-
- @staticmethod
- def _gather_logprob(logprob, target):
- lp_size = shape_list(logprob)
- r = tf.range(lp_size[0], dtype=target.dtype)
- idx = tf.stack([r, target], 1)
- return tf.gather_nd(logprob, idx)
-
- def call(self, hidden, target, return_mean=True, training=False):
- head_logprob = 0
- if self.n_clusters == 0:
- output = self._logit(hidden, self.out_layers[0][0], self.out_layers[0][1], self.out_projs[0])
- if target is not None:
- loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=target, logits=output)
- out = tf.nn.log_softmax(output, axis=-1)
- else:
- hidden_sizes = shape_list(hidden)
- out = []
- loss = tf.zeros(hidden_sizes[:2])
- for i in range(len(self.cutoffs)):
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
- if target is not None:
- mask = (target >= l_idx) & (target < r_idx)
- mask_idx = tf.where(mask)
- cur_target = tf.boolean_mask(target, mask) - l_idx
-
- if self.div_val == 1:
- cur_W = self.out_layers[0][0][l_idx:r_idx]
- cur_b = self.out_layers[0][1][l_idx:r_idx]
- else:
- cur_W = self.out_layers[i][0]
- cur_b = self.out_layers[i][1]
-
- if i == 0:
- cur_W = tf.concat([cur_W, self.cluster_weight], 0)
- cur_b = tf.concat([cur_b, self.cluster_bias], 0)
-
- head_logit = self._logit(hidden, cur_W, cur_b, self.out_projs[0])
- head_logprob = tf.nn.log_softmax(head_logit)
- out.append(head_logprob[..., : self.cutoffs[0]])
- if target is not None:
- cur_head_logprob = tf.boolean_mask(head_logprob, mask)
- cur_logprob = self._gather_logprob(cur_head_logprob, cur_target)
- else:
- tail_logit = self._logit(hidden, cur_W, cur_b, self.out_projs[i])
- tail_logprob = tf.nn.log_softmax(tail_logit)
- cluster_prob_idx = self.cutoffs[0] + i - 1 # No probability for the head cluster
- logprob_i = head_logprob[..., cluster_prob_idx, None] + tail_logprob
- out.append(logprob_i)
- if target is not None:
- cur_head_logprob = tf.boolean_mask(head_logprob, mask)
- cur_tail_logprob = tf.boolean_mask(tail_logprob, mask)
- cur_logprob = self._gather_logprob(cur_tail_logprob, cur_target)
- cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
- if target is not None:
- loss += tf.scatter_nd(mask_idx, -cur_logprob, shape_list(loss))
- out = tf.concat(out, axis=-1)
-
- if target is not None:
- if return_mean:
- loss = tf.reduce_mean(loss)
- # Add the training-time loss value to the layer using `self.add_loss()`.
- self.add_loss(loss)
-
- # Log the loss as a metric (we could log arbitrary metrics,
- # including different metrics for training and inference.
- self.add_metric(loss, name=self.name, aggregation="mean" if return_mean else "")
-
- return out
diff --git a/transformers/models/deprecated/transfo_xl/modeling_transfo_xl.py b/transformers/models/deprecated/transfo_xl/modeling_transfo_xl.py
deleted file mode 100644
index 897a3899c74cbd84713d96e3dad90cce2411dd17..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/transfo_xl/modeling_transfo_xl.py
+++ /dev/null
@@ -1,1295 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
- PyTorch Transformer XL model. Adapted from https://github.com/kimiyoung/transformer-xl. In particular
- https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/mem_transformer.py
-"""
-import warnings
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ....modeling_utils import PreTrainedModel
-from ....utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_transfo_xl import TransfoXLConfig
-from .modeling_transfo_xl_utilities import ProjectedAdaptiveLogSoftmax
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "transfo-xl/transfo-xl-wt103"
-_CONFIG_FOR_DOC = "TransfoXLConfig"
-
-
-from .._archive_maps import TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-def build_tf_to_pytorch_map(model, config):
- """
- A map of modules from TF to PyTorch. This time I use a map to keep the PyTorch model as identical to the original
- PyTorch model as possible.
- """
- tf_to_pt_map = {}
-
- if hasattr(model, "transformer"):
- # We are loading in a TransfoXLLMHeadModel => we will load also the Adaptive Softmax
- tf_to_pt_map.update(
- {
- "transformer/adaptive_softmax/cutoff_0/cluster_W": model.crit.cluster_weight,
- "transformer/adaptive_softmax/cutoff_0/cluster_b": model.crit.cluster_bias,
- }
- )
- for i, (out_l, proj_l, tie_proj) in enumerate(
- zip(model.crit.out_layers, model.crit.out_projs, config.tie_projs)
- ):
- layer_str = f"transformer/adaptive_softmax/cutoff_{i}/"
- if config.tie_word_embeddings:
- tf_to_pt_map.update({layer_str + "b": out_l.bias})
- else:
- raise NotImplementedError
- # I don't think this is implemented in the TF code
- tf_to_pt_map.update({layer_str + "lookup_table": out_l.weight, layer_str + "b": out_l.bias})
- if not tie_proj:
- tf_to_pt_map.update({layer_str + "proj": proj_l})
- # Now load the rest of the transformer
- model = model.transformer
-
- # Embeddings
- for i, (embed_l, proj_l) in enumerate(zip(model.word_emb.emb_layers, model.word_emb.emb_projs)):
- layer_str = f"transformer/adaptive_embed/cutoff_{i}/"
- tf_to_pt_map.update({layer_str + "lookup_table": embed_l.weight, layer_str + "proj_W": proj_l})
-
- # Transformer blocks
- for i, b in enumerate(model.layers):
- layer_str = f"transformer/layer_{i}/"
- tf_to_pt_map.update(
- {
- layer_str + "rel_attn/LayerNorm/gamma": b.dec_attn.layer_norm.weight,
- layer_str + "rel_attn/LayerNorm/beta": b.dec_attn.layer_norm.bias,
- layer_str + "rel_attn/o/kernel": b.dec_attn.o_net.weight,
- layer_str + "rel_attn/qkv/kernel": b.dec_attn.qkv_net.weight,
- layer_str + "rel_attn/r/kernel": b.dec_attn.r_net.weight,
- layer_str + "ff/LayerNorm/gamma": b.pos_ff.layer_norm.weight,
- layer_str + "ff/LayerNorm/beta": b.pos_ff.layer_norm.bias,
- layer_str + "ff/layer_1/kernel": b.pos_ff.CoreNet[0].weight,
- layer_str + "ff/layer_1/bias": b.pos_ff.CoreNet[0].bias,
- layer_str + "ff/layer_2/kernel": b.pos_ff.CoreNet[3].weight,
- layer_str + "ff/layer_2/bias": b.pos_ff.CoreNet[3].bias,
- }
- )
-
- # Relative positioning biases
- if config.untie_r:
- r_r_list = []
- r_w_list = []
- for b in model.layers:
- r_r_list.append(b.dec_attn.r_r_bias)
- r_w_list.append(b.dec_attn.r_w_bias)
- else:
- r_r_list = [model.r_r_bias]
- r_w_list = [model.r_w_bias]
- tf_to_pt_map.update({"transformer/r_r_bias": r_r_list, "transformer/r_w_bias": r_w_list})
- return tf_to_pt_map
-
-
-def load_tf_weights_in_transfo_xl(model, config, tf_path):
- """Load tf checkpoints in a pytorch model"""
- try:
- import numpy as np
- import tensorflow as tf
- except ImportError:
- logger.error(
- "Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
- "https://www.tensorflow.org/install/ for installation instructions."
- )
- raise
- # Build TF to PyTorch weights loading map
- tf_to_pt_map = build_tf_to_pytorch_map(model, config)
-
- # Load weights from TF model
- init_vars = tf.train.list_variables(tf_path)
- tf_weights = {}
- for name, shape in init_vars:
- logger.info(f"Loading TF weight {name} with shape {shape}")
- array = tf.train.load_variable(tf_path, name)
- tf_weights[name] = array
-
- for name, pointer in tf_to_pt_map.items():
- assert name in tf_weights
- array = tf_weights[name]
- # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
- # which are not required for using pretrained model
- if "kernel" in name or "proj" in name:
- array = np.transpose(array)
- if ("r_r_bias" in name or "r_w_bias" in name) and len(pointer) > 1:
- # Here we will split the TF weights
- assert len(pointer) == array.shape[0]
- for i, p_i in enumerate(pointer):
- arr_i = array[i, ...]
- try:
- assert p_i.shape == arr_i.shape
- except AssertionError as e:
- e.args += (p_i.shape, arr_i.shape)
- raise
- logger.info(f"Initialize PyTorch weight {name} for layer {i}")
- p_i.data = torch.from_numpy(arr_i)
- else:
- try:
- assert (
- pointer.shape == array.shape
- ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched"
- except AssertionError as e:
- e.args += (pointer.shape, array.shape)
- raise
- logger.info(f"Initialize PyTorch weight {name}")
- pointer.data = torch.from_numpy(array)
- tf_weights.pop(name, None)
- tf_weights.pop(name + "/Adam", None)
- tf_weights.pop(name + "/Adam_1", None)
-
- logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}")
- return model
-
-
-class PositionalEmbedding(nn.Module):
- def __init__(self, demb):
- super().__init__()
-
- self.demb = demb
-
- inv_freq = 1 / (10000 ** (torch.arange(0.0, demb, 2.0) / demb))
- self.register_buffer("inv_freq", inv_freq)
-
- def forward(self, pos_seq, bsz=None):
- sinusoid_inp = torch.outer(pos_seq, self.inv_freq)
- pos_emb = torch.cat([sinusoid_inp.sin(), sinusoid_inp.cos()], dim=-1)
-
- if bsz is not None:
- return pos_emb[:, None, :].expand(-1, bsz, -1)
- else:
- return pos_emb[:, None, :]
-
-
-class PositionwiseFF(nn.Module):
- def __init__(self, d_model, d_inner, dropout, pre_lnorm=False, layer_norm_epsilon=1e-5):
- super().__init__()
-
- self.d_model = d_model
- self.d_inner = d_inner
- self.dropout = dropout
-
- self.CoreNet = nn.Sequential(
- nn.Linear(d_model, d_inner),
- nn.ReLU(inplace=True),
- nn.Dropout(dropout),
- nn.Linear(d_inner, d_model),
- nn.Dropout(dropout),
- )
-
- self.layer_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
-
- self.pre_lnorm = pre_lnorm
-
- def forward(self, inp):
- if self.pre_lnorm:
- # layer normalization + positionwise feed-forward
- core_out = self.CoreNet(self.layer_norm(inp))
-
- # residual connection
- output = core_out + inp
- else:
- # positionwise feed-forward
- core_out = self.CoreNet(inp)
-
- # residual connection + layer normalization
- output = self.layer_norm(inp + core_out)
-
- return output
-
-
-class RelPartialLearnableMultiHeadAttn(nn.Module):
- def __init__(
- self,
- n_head,
- d_model,
- d_head,
- dropout,
- dropatt=0,
- pre_lnorm=False,
- r_r_bias=None,
- r_w_bias=None,
- layer_norm_epsilon=1e-5,
- ):
- super().__init__()
-
- self.n_head = n_head
- self.d_model = d_model
- self.d_head = d_head
- self.dropout = dropout
-
- self.qkv_net = nn.Linear(d_model, 3 * n_head * d_head, bias=False)
-
- self.drop = nn.Dropout(dropout)
- self.dropatt = nn.Dropout(dropatt)
- self.o_net = nn.Linear(n_head * d_head, d_model, bias=False)
-
- self.layer_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
-
- self.scale = 1 / (d_head**0.5)
-
- self.pre_lnorm = pre_lnorm
-
- if r_r_bias is None or r_w_bias is None: # Biases are not shared
- self.r_r_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
- self.r_w_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
- else:
- self.r_r_bias = r_r_bias
- self.r_w_bias = r_w_bias
-
- self.r_net = nn.Linear(self.d_model, self.n_head * self.d_head, bias=False)
-
- def _rel_shift(self, x):
- zero_pad_shape = (x.size(0), 1) + x.size()[2:]
- zero_pad = torch.zeros(zero_pad_shape, device=x.device, dtype=x.dtype)
- x_padded = torch.cat([zero_pad, x], dim=1)
-
- x_padded_shape = (x.size(1) + 1, x.size(0)) + x.size()[2:]
- x_padded = x_padded.view(*x_padded_shape)
-
- x = x_padded[1:].view_as(x)
-
- return x
-
- def forward(self, w, r, attn_mask=None, mems=None, head_mask=None, output_attentions=False):
- qlen, rlen, bsz = w.size(0), r.size(0), w.size(1)
-
- if mems is not None:
- cat = torch.cat([mems, w], 0)
- if self.pre_lnorm:
- w_heads = self.qkv_net(self.layer_norm(cat))
- else:
- w_heads = self.qkv_net(cat)
- r_head_k = self.r_net(r)
-
- w_head_q, w_head_k, w_head_v = torch.chunk(w_heads, 3, dim=-1)
- w_head_q = w_head_q[-qlen:]
- else:
- if self.pre_lnorm:
- w_heads = self.qkv_net(self.layer_norm(w))
- else:
- w_heads = self.qkv_net(w)
- r_head_k = self.r_net(r)
-
- w_head_q, w_head_k, w_head_v = torch.chunk(w_heads, 3, dim=-1)
-
- klen = w_head_k.size(0)
-
- w_head_q = w_head_q.view(qlen, bsz, self.n_head, self.d_head) # qlen x bsz x n_head x d_head
- w_head_k = w_head_k.view(klen, bsz, self.n_head, self.d_head) # qlen x bsz x n_head x d_head
- w_head_v = w_head_v.view(klen, bsz, self.n_head, self.d_head) # qlen x bsz x n_head x d_head
-
- r_head_k = r_head_k.view(rlen, self.n_head, self.d_head) # qlen x n_head x d_head
-
- # compute attention score
- rw_head_q = w_head_q + self.r_w_bias # qlen x bsz x n_head x d_head
- AC = torch.einsum("ibnd,jbnd->ijbn", (rw_head_q, w_head_k)) # qlen x klen x bsz x n_head
-
- rr_head_q = w_head_q + self.r_r_bias
- BD = torch.einsum("ibnd,jnd->ijbn", (rr_head_q, r_head_k)) # qlen x klen x bsz x n_head
- BD = self._rel_shift(BD)
-
- # [qlen x klen x bsz x n_head]
- attn_score = AC + BD
- attn_score.mul_(self.scale)
-
- mask_value = torch.finfo(attn_score.dtype).min
-
- # compute attention probability
- if attn_mask is not None and torch.sum(attn_mask).item():
- attn_mask = attn_mask == 1 # Switch to bool
- if attn_mask.dim() == 2:
- attn_score = (
- attn_score.float().masked_fill(attn_mask[None, :, :, None], mask_value).type_as(attn_score)
- )
- elif attn_mask.dim() == 3:
- attn_score = attn_score.float().masked_fill(attn_mask[:, :, :, None], mask_value).type_as(attn_score)
-
- # [qlen x klen x bsz x n_head]
- attn_prob = nn.functional.softmax(attn_score, dim=1)
- attn_prob = self.dropatt(attn_prob)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_prob = attn_prob * head_mask
-
- # compute attention vector
- attn_vec = torch.einsum("ijbn,jbnd->ibnd", (attn_prob, w_head_v))
-
- # [qlen x bsz x n_head x d_head]
- attn_vec = attn_vec.contiguous().view(attn_vec.size(0), attn_vec.size(1), self.n_head * self.d_head)
-
- # linear projection
- attn_out = self.o_net(attn_vec)
- attn_out = self.drop(attn_out)
-
- if self.pre_lnorm:
- # residual connection
- outputs = [w + attn_out]
- else:
- # residual connection + layer normalization
- outputs = [self.layer_norm(w + attn_out)]
-
- if output_attentions:
- outputs.append(attn_prob)
-
- return outputs
-
-
-class RelPartialLearnableDecoderLayer(nn.Module):
- def __init__(self, n_head, d_model, d_head, d_inner, dropout, layer_norm_epsilon=1e-5, **kwargs):
- super().__init__()
-
- self.dec_attn = RelPartialLearnableMultiHeadAttn(
- n_head, d_model, d_head, dropout, layer_norm_epsilon=layer_norm_epsilon, **kwargs
- )
- self.pos_ff = PositionwiseFF(
- d_model, d_inner, dropout, pre_lnorm=kwargs.get("pre_lnorm"), layer_norm_epsilon=layer_norm_epsilon
- )
-
- def forward(self, dec_inp, r, dec_attn_mask=None, mems=None, head_mask=None, output_attentions=False):
- attn_outputs = self.dec_attn(
- dec_inp,
- r,
- attn_mask=dec_attn_mask,
- mems=mems,
- head_mask=head_mask,
- output_attentions=output_attentions,
- )
- ff_output = self.pos_ff(attn_outputs[0])
-
- outputs = [ff_output] + attn_outputs[1:]
-
- return outputs
-
-
-class AdaptiveEmbedding(nn.Module):
- def __init__(self, n_token, d_embed, d_proj, cutoffs, div_val=1, sample_softmax=False):
- super().__init__()
-
- self.n_token = n_token
- self.d_embed = d_embed
-
- self.cutoffs = cutoffs + [n_token]
- self.div_val = div_val
- self.d_proj = d_proj
-
- self.emb_scale = d_proj**0.5
-
- self.cutoff_ends = [0] + self.cutoffs
-
- self.emb_layers = nn.ModuleList()
- self.emb_projs = nn.ParameterList()
- if div_val == 1:
- self.emb_layers.append(nn.Embedding(n_token, d_embed, sparse=sample_softmax > 0))
- if d_proj != d_embed:
- self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_embed)))
- else:
- for i in range(len(self.cutoffs)):
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
- d_emb_i = d_embed // (div_val**i)
- self.emb_layers.append(nn.Embedding(r_idx - l_idx, d_emb_i))
- self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_emb_i)))
-
- def forward(self, inp):
- if self.div_val == 1:
- embed = self.emb_layers[0](inp)
- if self.d_proj != self.d_embed:
- embed = nn.functional.linear(embed, self.emb_projs[0])
- else:
- param = next(self.parameters())
- inp_flat = inp.view(-1)
- emb_flat = torch.zeros([inp_flat.size(0), self.d_proj], dtype=param.dtype, device=param.device)
- for i in range(len(self.cutoffs)):
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
-
- mask_i = (inp_flat >= l_idx) & (inp_flat < r_idx)
- indices_i = mask_i.nonzero().squeeze()
-
- if indices_i.numel() == 0:
- continue
-
- inp_i = inp_flat.index_select(0, indices_i) - l_idx
- emb_i = self.emb_layers[i](inp_i)
- emb_i = nn.functional.linear(emb_i, self.emb_projs[i])
-
- emb_flat.index_copy_(0, indices_i, emb_i)
-
- embed_shape = inp.size() + (self.d_proj,)
- embed = emb_flat.view(embed_shape)
-
- embed.mul_(self.emb_scale)
-
- return embed
-
-
-class TransfoXLPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = TransfoXLConfig
- load_tf_weights = load_tf_weights_in_transfo_xl
- base_model_prefix = "transformer"
-
- def _init_weight(self, weight):
- if self.config.init == "uniform":
- nn.init.uniform_(weight, -self.config.init_range, self.config.init_range)
- elif self.config.init == "normal":
- nn.init.normal_(weight, 0.0, self.config.init_std)
-
- def _init_bias(self, bias):
- nn.init.constant_(bias, 0.0)
-
- def _init_weights(self, m):
- """Initialize the weights."""
- classname = m.__class__.__name__
- if classname.find("Linear") != -1:
- if hasattr(m, "weight") and m.weight is not None:
- self._init_weight(m.weight)
- if hasattr(m, "bias") and m.bias is not None:
- self._init_bias(m.bias)
- elif classname.find("AdaptiveEmbedding") != -1:
- if hasattr(m, "emb_projs"):
- for i in range(len(m.emb_projs)):
- if m.emb_projs[i] is not None:
- nn.init.normal_(m.emb_projs[i], 0.0, self.config.proj_init_std)
- elif classname.find("Embedding") != -1:
- if hasattr(m, "weight"):
- self._init_weight(m.weight)
- elif classname.find("ProjectedAdaptiveLogSoftmax") != -1:
- if hasattr(m, "cluster_weight") and m.cluster_weight is not None:
- self._init_weight(m.cluster_weight)
- if hasattr(m, "cluster_bias") and m.cluster_bias is not None:
- self._init_bias(m.cluster_bias)
- if hasattr(m, "out_projs"):
- for i in range(len(m.out_projs)):
- if m.out_projs[i] is not None:
- nn.init.normal_(m.out_projs[i], 0.0, self.config.proj_init_std)
- elif classname.find("LayerNorm") != -1:
- if hasattr(m, "weight"):
- nn.init.normal_(m.weight, 1.0, self.config.init_std)
- if hasattr(m, "bias") and m.bias is not None:
- self._init_bias(m.bias)
- else:
- if hasattr(m, "r_emb"):
- self._init_weight(m.r_emb)
- if hasattr(m, "r_w_bias"):
- self._init_weight(m.r_w_bias)
- if hasattr(m, "r_r_bias"):
- self._init_weight(m.r_r_bias)
- if hasattr(m, "r_bias"):
- self._init_bias(m.r_bias)
-
- def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, layer: Optional[int] = -1):
- """
- Resize input token embeddings matrix of the model if new_num_tokens != config.vocab_size. Take care of tying
- weights embeddings afterwards if the model class has a *tie_weights()* method.
-
- Arguments:
- new_num_tokens: (*optional*) int:
- New number of tokens in the embedding matrix. Increasing the size will add newly initialized vectors at
- the end. Reducing the size will remove vectors from the end. If not provided or None: does nothing and
- just returns a pointer to the input tokens `torch.nn.Embeddings` Module of the model.
- layer: (*optional*) int:
- Layer of the *AdaptiveEmbedding* where the resizing should be done. Per default the last layer will be
- resized. Be aware that when resizing other than the last layer, you have to ensure that the new
- token(s) in the tokenizer are at the corresponding position.
-
- Return: `torch.nn.Embeddings` Pointer to the input tokens Embeddings Module of the model
- """
- base_model = getattr(self, self.base_model_prefix, self) # get the base model if needed
-
- if new_num_tokens is None:
- return self.get_input_embeddings()
-
- new_num_tokens_layer, layer = self._get_new_num_tokens_layer(new_num_tokens, layer)
- assert new_num_tokens_layer > 0, "The size of the new embedding layer cannot be 0 or less"
- model_embeds = base_model._resize_token_embeddings(new_num_tokens_layer, layer)
-
- # Update base model and current model config
- self.config.vocab_size = new_num_tokens
- base_model.vocab_size = new_num_tokens
- base_model.n_token = new_num_tokens
-
- new_embedding_shapes = self._get_embedding_shapes()
- self._resize_cutoffs(new_num_tokens, new_num_tokens_layer, new_embedding_shapes, layer)
-
- # Tie weights again if needed
- self.tie_weights()
-
- return model_embeds
-
- def _get_new_num_tokens_layer(self, new_num_tokens, layer):
- embeddings = self.get_input_embeddings()
- if layer == -1:
- layer = len(embeddings.emb_layers) - 1
- assert 0 <= layer <= len(embeddings.emb_layers) - 1
-
- new_num_tokens_layer = (
- new_num_tokens
- - sum([emb.weight.shape[0] for emb in embeddings.emb_layers[:layer]])
- - sum([emb.weight.shape[0] for emb in embeddings.emb_layers[layer + 1 :]])
- )
- return new_num_tokens_layer, layer
-
- def _get_embedding_shapes(self):
- embeddings = self.get_input_embeddings()
- return [emb.weight.shape[0] for emb in embeddings.emb_layers]
-
- def _resize_token_embeddings(self, new_num_tokens, layer=-1):
- embeddings = self.get_input_embeddings()
- if new_num_tokens is None:
- return embeddings
- new_embeddings_layer = self._get_resized_embeddings(embeddings.emb_layers[layer], new_num_tokens)
- embeddings.emb_layers[layer] = new_embeddings_layer
-
- self.set_input_embeddings(embeddings)
-
- return self.get_input_embeddings()
-
- def _resize_cutoffs(self, new_num_tokens, new_emb_size, new_embedding_shapes, layer):
- embeddings = self.get_input_embeddings()
-
- for i in range(layer, len(embeddings.cutoffs)):
- embeddings.cutoffs[i] = sum(new_embedding_shapes[: i + 1])
-
- embeddings.cutoff_ends = [0] + embeddings.cutoffs
- embeddings.n_token = new_num_tokens
-
- self.config.cutoffs = embeddings.cutoffs[:-1]
-
- return embeddings.cutoffs
-
-
-@dataclass
-class TransfoXLModelOutput(ModelOutput):
- """
- Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- mems (`List[torch.FloatTensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks). Can be used (see `mems`
- input) to speed up sequential decoding. The token ids which have their past given to this model should not
- be passed as input ids as they have already been computed.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- last_hidden_state: torch.FloatTensor
- mems: List[torch.FloatTensor] = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class TransfoXLSequenceClassifierOutputWithPast(ModelOutput):
- """
- Base class for outputs of sentence classification models.
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
- Classification (or regression if config.num_labels==1) loss.
- logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Classification (or regression if config.num_labels==1) scores (before SoftMax).
- mems (`List[torch.FloatTensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks). Can be used (see `mems`
- input) to speed up sequential decoding. The token ids which have their past given to this model should not
- be passed as input ids as they have already been computed.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- logits: torch.FloatTensor = None
- mems: List[torch.FloatTensor] = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class TransfoXLLMHeadModelOutput(ModelOutput):
- """
- Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
-
- Args:
- losses (`torch.FloatTensor` of shape *(batch_size, sequence_length-1)*, *optional*, returned when `labels` is provided):
- Language modeling losses (not reduced).
- prediction_scores (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token after SoftMax).
- mems (`List[torch.FloatTensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks). Can be used (see `mems`
- input) to speed up sequential decoding. The token ids which have their past given to this model should not
- be passed as input ids as they have already been computed.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- loss (`torch.FloatTensor` of shape `()`, *optional*, returned when `labels` is provided)
- Reduced language modeling loss.
- """
-
- losses: Optional[torch.FloatTensor] = None
- prediction_scores: torch.FloatTensor = None
- mems: List[torch.FloatTensor] = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
- loss: Optional[torch.FloatTensor] = None
-
- @property
- def logits(self):
- # prediction scores are the output of the adaptive softmax, see
- # the file `modeling_transfo_xl_utilities`. Since the adaptive
- # softmax returns the log softmax value, `self.prediction_scores`
- # are strictly speaking not exactly `logits`, but behave the same
- # way logits do.
- return self.prediction_scores
-
-
-TRANSFO_XL_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`TransfoXLConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-TRANSFO_XL_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- mems (`List[torch.FloatTensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model (see
- `mems` output below). Can be used to speed up sequential decoding. The token ids which have their mems
- given to this model should not be passed as `input_ids` as they have already been computed.
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
- TRANSFO_XL_START_DOCSTRING,
-)
-class TransfoXLModel(TransfoXLPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.n_token = config.vocab_size
-
- self.d_embed = config.d_embed
- self.d_model = config.d_model
- self.n_head = config.n_head
- self.d_head = config.d_head
-
- self.word_emb = AdaptiveEmbedding(
- config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val
- )
-
- self.drop = nn.Dropout(config.dropout)
-
- self.n_layer = config.n_layer
- self.mem_len = config.mem_len
- self.attn_type = config.attn_type
-
- if not config.untie_r:
- self.r_w_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
- self.r_r_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
-
- self.layers = nn.ModuleList()
- if config.attn_type == 0: # the default attention
- for i in range(config.n_layer):
- self.layers.append(
- RelPartialLearnableDecoderLayer(
- config.n_head,
- config.d_model,
- config.d_head,
- config.d_inner,
- config.dropout,
- dropatt=config.dropatt,
- pre_lnorm=config.pre_lnorm,
- r_w_bias=None if config.untie_r else self.r_w_bias,
- r_r_bias=None if config.untie_r else self.r_r_bias,
- layer_norm_epsilon=config.layer_norm_epsilon,
- )
- )
- else: # learnable embeddings and absolute embeddings are not used in our pretrained checkpoints
- raise NotImplementedError # Removed them to avoid maintaining dead code
-
- self.same_length = config.same_length
- self.clamp_len = config.clamp_len
-
- if self.attn_type == 0: # default attention
- self.pos_emb = PositionalEmbedding(self.d_model)
- else: # learnable embeddings and absolute embeddings
- raise NotImplementedError # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.word_emb
-
- def set_input_embeddings(self, new_embeddings):
- self.word_emb = new_embeddings
-
- def backward_compatible(self):
- self.sample_softmax = -1
-
- def reset_memory_length(self, mem_len):
- self.mem_len = mem_len
-
- def _prune_heads(self, heads):
- logger.info("Head pruning is not implemented for Transformer-XL model")
- pass
-
- def init_mems(self, bsz):
- if self.mem_len > 0:
- mems = []
- param = next(self.parameters())
- for i in range(self.n_layer):
- empty = torch.zeros(self.mem_len, bsz, self.config.d_model, dtype=param.dtype, device=param.device)
- mems.append(empty)
-
- return mems
- else:
- return None
-
- def _update_mems(self, hids, mems, mlen, qlen):
- # does not deal with None
- if mems is None:
- return None
-
- # mems is not None
- assert len(hids) == len(mems), "len(hids) != len(mems)"
-
- # There are `mlen + qlen` steps that can be cached into mems
- with torch.no_grad():
- new_mems = []
- end_idx = mlen + max(0, qlen)
- beg_idx = max(0, end_idx - self.mem_len)
- for i in range(len(hids)):
- cat = torch.cat([mems[i], hids[i]], dim=0)
- new_mems.append(cat[beg_idx:end_idx].detach())
-
- return new_mems
-
- @add_start_docstrings_to_model_forward(TRANSFO_XL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TransfoXLModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- mems: Optional[List[torch.FloatTensor]] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TransfoXLModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # the original code for Transformer-XL used shapes [len, bsz] but we want a unified interface in the library
- # so we transpose here from shape [bsz, len] to shape [len, bsz]
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_ids = input_ids.transpose(0, 1).contiguous()
- qlen, bsz = input_ids.size()
- elif inputs_embeds is not None:
- inputs_embeds = inputs_embeds.transpose(0, 1).contiguous()
- qlen, bsz = inputs_embeds.shape[0], inputs_embeds.shape[1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if mems is None:
- mems = self.init_mems(bsz)
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] (a head_mask for each layer)
- # and head_mask is converted to shape [num_hidden_layers x qlen x klen x bsz x n_head]
- if head_mask is not None:
- if head_mask.dim() == 1:
- head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(0).unsqueeze(0)
- head_mask = head_mask.expand(self.n_layer, -1, -1, -1, -1)
- elif head_mask.dim() == 2:
- head_mask = head_mask.unsqueeze(1).unsqueeze(1).unsqueeze(1)
- head_mask = head_mask.to(
- dtype=next(self.parameters()).dtype
- ) # switch to float if need + fp16 compatibility
- else:
- head_mask = [None] * self.n_layer
-
- if inputs_embeds is not None:
- word_emb = inputs_embeds
- else:
- word_emb = self.word_emb(input_ids)
-
- mlen = mems[0].size(0) if mems is not None else 0
- klen = mlen + qlen
- if self.same_length:
- all_ones = word_emb.new_ones((qlen, klen), dtype=torch.bool)
- mask_len = klen - self.mem_len
- if mask_len > 0:
- mask_shift_len = qlen - mask_len
- else:
- mask_shift_len = qlen
- dec_attn_mask = (torch.triu(all_ones, 1 + mlen) + torch.tril(all_ones, -mask_shift_len))[:, :, None] # -1
- else:
- dec_attn_mask = torch.triu(word_emb.new_ones((qlen, klen), dtype=torch.bool), diagonal=1 + mlen)[
- :, :, None
- ]
-
- hids = []
- attentions = [] if output_attentions else None
- if self.attn_type == 0: # default
- pos_seq = torch.arange(klen - 1, -1, -1.0, device=word_emb.device, dtype=torch.int64).type_as(
- dtype=word_emb.dtype
- )
- if self.clamp_len > 0:
- pos_seq.clamp_(max=self.clamp_len)
- pos_emb = self.pos_emb(pos_seq)
-
- core_out = self.drop(word_emb)
- pos_emb = self.drop(pos_emb)
-
- for i, layer in enumerate(self.layers):
- hids.append(core_out)
- mems_i = None if mems is None else mems[i]
- layer_outputs = layer(
- core_out,
- pos_emb,
- dec_attn_mask=dec_attn_mask,
- mems=mems_i,
- head_mask=head_mask[i],
- output_attentions=output_attentions,
- )
- core_out = layer_outputs[0]
- if output_attentions:
- attentions.append(layer_outputs[1])
- else: # learnable embeddings and absolute embeddings
- raise NotImplementedError # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
-
- core_out = self.drop(core_out)
-
- new_mems = self._update_mems(hids, mems, mlen, qlen)
-
- if output_hidden_states:
- # Add last layer and transpose to library standard shape [bsz, len, hidden_dim]
- hids.append(core_out)
- hids = tuple(t.transpose(0, 1).contiguous() for t in hids)
- else:
- hids = None
- if output_attentions:
- # Transpose to library standard shape [bsz, n_heads, query_seq_len, key_seq_len]
- attentions = tuple(t.permute(2, 3, 0, 1).contiguous() for t in attentions)
- # We transpose back here to shape [bsz, len, hidden_dim]
- core_out = core_out.transpose(0, 1).contiguous()
-
- if not return_dict:
- return tuple(v for v in [core_out, new_mems, hids, attentions] if v is not None)
-
- return TransfoXLModelOutput(
- last_hidden_state=core_out,
- mems=new_mems,
- hidden_states=hids,
- attentions=attentions,
- )
-
-
-@add_start_docstrings(
- """
- The Transformer-XL Model with a language modeling head on top (adaptive softmax with weights tied to the adaptive
- input embeddings)
- """,
- TRANSFO_XL_START_DOCSTRING,
-)
-class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
- _tied_weights_keys = [r"crit\.out_projs\.\d+", r"crit\.out_layers\.\d+\.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.transformer = TransfoXLModel(config)
- self.sample_softmax = config.sample_softmax
- self.trainer_compatible = getattr(config, "trainer_compatible", False)
-
- if not self.trainer_compatible:
- warnings.warn(
- "The output of TransfoXL will be updated in v5 to support a single loss as first argument. In order "
- "to use that updated output, please specify `trainer_compatible=True` as your configuration"
- " attribute.",
- DeprecationWarning,
- )
-
- assert self.sample_softmax <= 0, (
- "Sampling from the softmax is not implemented yet. Please look at issue: #3310:"
- " https://github.com/huggingface/transformers/issues/3310"
- )
-
- self.crit = ProjectedAdaptiveLogSoftmax(
- config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def tie_weights(self):
- """
- Run this to be sure output and input (adaptive) softmax weights are tied
- """
-
- if self.config.tie_word_embeddings:
- for i in range(len(self.crit.out_layers)):
- self._tie_or_clone_weights(self.crit.out_layers[i], self.transformer.word_emb.emb_layers[i])
- if self.config.tie_projs:
- for i, tie_proj in enumerate(self.config.tie_projs):
- if tie_proj and self.config.div_val == 1 and self.config.d_model != self.config.d_embed:
- if self.config.torchscript:
- self.crit.out_projs[i] = nn.Parameter(self.transformer.word_emb.emb_projs[0].clone())
- else:
- self.crit.out_projs[i] = self.transformer.word_emb.emb_projs[0]
- elif tie_proj and self.config.div_val != 1:
- if self.config.torchscript:
- self.crit.out_projs[i] = nn.Parameter(self.transformer.word_emb.emb_projs[i].clone())
- else:
- self.crit.out_projs[i] = self.transformer.word_emb.emb_projs[i]
-
- def reset_memory_length(self, mem_len):
- self.transformer.reset_memory_length(mem_len)
-
- def init_mems(self, bsz):
- return self.transformer.init_mems(bsz)
-
- @add_start_docstrings_to_model_forward(TRANSFO_XL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TransfoXLLMHeadModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- mems: Optional[List[torch.FloatTensor]] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TransfoXLLMHeadModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- if input_ids is not None:
- bsz, tgt_len = input_ids.size(0), input_ids.size(1)
- elif inputs_embeds is not None:
- bsz, tgt_len = inputs_embeds.size(0), inputs_embeds.size(1)
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- transformer_outputs = self.transformer(
- input_ids,
- mems=mems,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden = transformer_outputs[0]
- pred_hid = last_hidden[:, -tgt_len:]
-
- if labels is not None:
- # Prevents all labels being -100 and throwing an error
- # when backwarding the loss
- miss_valid_label = labels[0, 1:].sum() == (labels.size(1) - 1) * -100
- if miss_valid_label:
- # Sets an token, just to prevent loss from being NaN
- labels[0, 1] = self.config.eos_token_id
-
- softmax_output = self.crit(pred_hid, labels)
- prediction_scores = softmax_output.view(bsz, tgt_len, -1) if labels is None else ()
-
- if labels is not None:
- losses = softmax_output.view(bsz, tgt_len - 1)
- # Avoids from incorporating padding (-100) tokens into loss value
- loss = losses[losses != 0].mean()
- else:
- losses, loss = None, None
-
- if not return_dict:
- if self.trainer_compatible:
- output = (prediction_scores, losses) if losses is not None else (prediction_scores,)
- output += transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
- else:
- output = (prediction_scores, *transformer_outputs[1:])
- output = ((losses,) + output) if losses is not None else output
- return (output + (loss,)) if loss is not None else output
-
- return TransfoXLLMHeadModelOutput(
- loss=loss,
- prediction_scores=prediction_scores,
- losses=losses,
- mems=transformer_outputs.mems,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def get_output_embeddings(self):
- """Double-check if you are using adaptive softmax."""
- if self.sample_softmax > 0:
- return self.out_layer
- else:
- return self.crit.out_layers[-1]
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **model_kwargs):
- inputs = {}
-
- # if past is defined in model kwargs then use it for faster decoding
- if past_key_values:
- inputs["mems"] = past_key_values
- inputs["input_ids"] = input_ids[:, -1].unsqueeze(-1)
- else:
- inputs["input_ids"] = input_ids
-
- return inputs
-
- def _resize_cutoffs(self, new_num_tokens, new_emb_size, new_embedding_shapes, layer):
- new_cutoffs = super()._resize_cutoffs(new_num_tokens, new_emb_size, new_embedding_shapes, layer)
-
- self.crit.cutoffs = new_cutoffs
- self.crit.cutoff_ends = [0] + new_cutoffs
- self.crit.n_token = new_num_tokens
-
- @staticmethod
- def _reorder_cache(mems: List[torch.Tensor], beam_idx: torch.Tensor) -> List[torch.Tensor]:
- """
- This function is used to re-order the `mems` cache if [`~PreTrainedModel.beam_search`] or
- [`~PreTrainedModel.beam_sample`] is called. This is required to match `mems` with the correct beam_idx at every
- generation step.
- """
- return [layer_past.index_select(1, beam_idx.to(layer_past.device)) for layer_past in mems]
-
-
-@add_start_docstrings(
- """
- The Transformer-XL Model transformer with a sequence classification head on top (linear layer).
-
- [`TransfoXLForSequenceClassification`] uses the last token in order to do the classification, as other causal
- models (e.g. GPT-1) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- TRANSFO_XL_START_DOCSTRING,
-)
-class TransfoXLForSequenceClassification(TransfoXLPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = TransfoXLModel(config)
- self.score = nn.Linear(config.d_embed, self.num_labels, bias=False)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(TRANSFO_XL_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TransfoXLSequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- mems: Optional[List[torch.FloatTensor]] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TransfoXLSequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- mems=mems,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size, sequence_length = input_ids.shape[:2]
- else:
- batch_size, sequence_length = inputs_embeds.shape[:2]
-
- assert (
- self.config.pad_token_id is not None or batch_size == 1
- ), "Cannot handle batch sizes > 1 if no padding token is defined."
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
-
- pooled_logits = logits[range(batch_size), sequence_lengths]
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TransfoXLSequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- mems=transformer_outputs.mems,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
diff --git a/transformers/models/deprecated/transfo_xl/modeling_transfo_xl_utilities.py b/transformers/models/deprecated/transfo_xl/modeling_transfo_xl_utilities.py
deleted file mode 100644
index addf2a08372bc00a377ab7410d977c31fb1d48eb..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/transfo_xl/modeling_transfo_xl_utilities.py
+++ /dev/null
@@ -1,252 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
- Utilities for PyTorch Transformer XL model. Directly adapted from https://github.com/kimiyoung/transformer-xl.
-"""
-
-
-import torch
-from torch import nn
-
-
-# CUDA_MAJOR = int(torch.version.cuda.split('.')[0])
-# CUDA_MINOR = int(torch.version.cuda.split('.')[1])
-
-
-class ProjectedAdaptiveLogSoftmax(nn.Module):
- def __init__(self, n_token, d_embed, d_proj, cutoffs, div_val=1, keep_order=False):
- super().__init__()
-
- self.n_token = n_token
- self.d_embed = d_embed
- self.d_proj = d_proj
-
- self.cutoffs = cutoffs + [n_token]
- self.cutoff_ends = [0] + self.cutoffs
- self.div_val = div_val
-
- self.shortlist_size = self.cutoffs[0]
- self.n_clusters = len(self.cutoffs) - 1
- self.head_size = self.shortlist_size + self.n_clusters
-
- if self.n_clusters > 0:
- self.cluster_weight = nn.Parameter(torch.zeros(self.n_clusters, self.d_embed))
- self.cluster_bias = nn.Parameter(torch.zeros(self.n_clusters))
-
- self.out_layers = nn.ModuleList()
- self.out_projs = nn.ParameterList()
-
- if div_val == 1:
- for i in range(len(self.cutoffs)):
- if d_proj != d_embed:
- self.out_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_embed)))
- else:
- self.out_projs.append(None)
-
- self.out_layers.append(nn.Linear(d_embed, n_token))
- else:
- for i in range(len(self.cutoffs)):
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
- d_emb_i = d_embed // (div_val**i)
-
- self.out_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_emb_i)))
-
- self.out_layers.append(nn.Linear(d_emb_i, r_idx - l_idx))
-
- self.keep_order = keep_order
-
- def _compute_logit(self, hidden, weight, bias, proj):
- if proj is None:
- logit = nn.functional.linear(hidden, weight, bias=bias)
- else:
- # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
- proj_hid = nn.functional.linear(hidden, proj.t().contiguous())
- logit = nn.functional.linear(proj_hid, weight, bias=bias)
- # else:
- # logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
- # if bias is not None:
- # logit = logit + bias
-
- return logit
-
- def forward(self, hidden, labels=None, keep_order=False):
- """
- Params:
- hidden :: [len*bsz x d_proj]
- labels :: [len*bsz]
-
- Return:
- if labels is None: out :: [len*bsz x n_tokens] log probabilities of tokens over the vocabulary else: out ::
- [(len-1)*bsz] Negative log likelihood. We could replace this implementation by the native PyTorch one if
- theirs had an option to set bias on all clusters in the native one. here:
- https://github.com/pytorch/pytorch/blob/dbe6a7a9ff1a364a8706bf5df58a1ca96d2fd9da/torch/nn/modules/adaptive.py#L138
- """
-
- if labels is not None:
- # Shift so that tokens < n predict n
- hidden = hidden[..., :-1, :].contiguous()
- labels = labels[..., 1:].contiguous()
- hidden = hidden.view(-1, hidden.size(-1))
- labels = labels.view(-1)
- if hidden.size(0) != labels.size(0):
- raise RuntimeError("Input and labels should have the same size in the batch dimension.")
- else:
- hidden = hidden.view(-1, hidden.size(-1))
-
- if self.n_clusters == 0:
- logit = self._compute_logit(hidden, self.out_layers[0].weight, self.out_layers[0].bias, self.out_projs[0])
- if labels is not None:
- mask = labels != -100
- out = torch.zeros_like(labels, dtype=hidden.dtype, device=hidden.device)
- out[mask] = (
- -nn.functional.log_softmax(logit, dim=-1)[mask].gather(1, labels[mask].unsqueeze(1)).squeeze(1)
- )
- else:
- out = nn.functional.log_softmax(logit, dim=-1)
- else:
- # construct weights and biases
- weights, biases = [], []
- for i in range(len(self.cutoffs)):
- if self.div_val == 1:
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
- weight_i = self.out_layers[0].weight[l_idx:r_idx]
- bias_i = self.out_layers[0].bias[l_idx:r_idx]
- else:
- weight_i = self.out_layers[i].weight
- bias_i = self.out_layers[i].bias
-
- if i == 0:
- weight_i = torch.cat([weight_i, self.cluster_weight], dim=0)
- bias_i = torch.cat([bias_i, self.cluster_bias], dim=0)
-
- weights.append(weight_i)
- biases.append(bias_i)
-
- head_weight, head_bias, head_proj = weights[0], biases[0], self.out_projs[0]
-
- head_logit = self._compute_logit(hidden, head_weight, head_bias, head_proj)
- head_logprob = nn.functional.log_softmax(head_logit, dim=1)
-
- if labels is None:
- out = hidden.new_empty((head_logit.size(0), self.n_token))
- else:
- out = torch.zeros_like(labels, dtype=hidden.dtype, device=hidden.device)
-
- offset = 0
- cutoff_values = [0] + self.cutoffs
- for i in range(len(cutoff_values) - 1):
- l_idx, r_idx = cutoff_values[i], cutoff_values[i + 1]
-
- if labels is not None:
- mask_i = (labels >= l_idx) & (labels < r_idx)
- indices_i = mask_i.nonzero().squeeze()
-
- if indices_i.numel() == 0:
- continue
-
- target_i = labels.index_select(0, indices_i) - l_idx
- head_logprob_i = head_logprob.index_select(0, indices_i)
- hidden_i = hidden.index_select(0, indices_i)
- else:
- hidden_i = hidden
-
- if i == 0:
- if labels is not None:
- logprob_i = head_logprob_i.gather(1, target_i[:, None]).squeeze(1)
- else:
- out[:, : self.cutoffs[0]] = head_logprob[:, : self.cutoffs[0]]
- else:
- weight_i, bias_i, proj_i = weights[i], biases[i], self.out_projs[i]
-
- tail_logit_i = self._compute_logit(hidden_i, weight_i, bias_i, proj_i)
- tail_logprob_i = nn.functional.log_softmax(tail_logit_i, dim=1)
- cluster_prob_idx = self.cutoffs[0] + i - 1 # No probability for the head cluster
- if labels is not None:
- logprob_i = head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
- 1, target_i[:, None]
- ).squeeze(1)
- else:
- logprob_i = head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
- out[:, l_idx:r_idx] = logprob_i
-
- if labels is not None:
- if (hasattr(self, "keep_order") and self.keep_order) or keep_order:
- out.index_copy_(0, indices_i, -logprob_i)
- else:
- out[offset : offset + logprob_i.size(0)].copy_(-logprob_i)
- offset += logprob_i.size(0)
-
- return out
-
- def log_prob(self, hidden):
- r"""
- Computes log probabilities for all \\(n\_classes\\) From:
- https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/adaptive.p
-
- Args:
- hidden (Tensor): a minibatch of example
-
- Returns:
- log-probabilities of for each class \\(c\\) in range \\(0 <= c <= n\_classes\\), where \\(n\_classes\\) is
- a parameter passed to `AdaptiveLogSoftmaxWithLoss` constructor. Shape:
-
- - Input: \\((N, in\_features)\\)
- - Output: \\((N, n\_classes)\\)
- """
- if self.n_clusters == 0:
- logit = self._compute_logit(hidden, self.out_layers[0].weight, self.out_layers[0].bias, self.out_projs[0])
- return nn.functional.log_softmax(logit, dim=-1)
- else:
- # construct weights and biases
- weights, biases = [], []
- for i in range(len(self.cutoffs)):
- if self.div_val == 1:
- l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
- weight_i = self.out_layers[0].weight[l_idx:r_idx]
- bias_i = self.out_layers[0].bias[l_idx:r_idx]
- else:
- weight_i = self.out_layers[i].weight
- bias_i = self.out_layers[i].bias
-
- if i == 0:
- weight_i = torch.cat([weight_i, self.cluster_weight], dim=0)
- bias_i = torch.cat([bias_i, self.cluster_bias], dim=0)
-
- weights.append(weight_i)
- biases.append(bias_i)
-
- head_weight, head_bias, head_proj = weights[0], biases[0], self.out_projs[0]
- head_logit = self._compute_logit(hidden, head_weight, head_bias, head_proj)
-
- out = hidden.new_empty((head_logit.size(0), self.n_token))
- head_logprob = nn.functional.log_softmax(head_logit, dim=1)
-
- cutoff_values = [0] + self.cutoffs
- for i in range(len(cutoff_values) - 1):
- start_idx, stop_idx = cutoff_values[i], cutoff_values[i + 1]
-
- if i == 0:
- out[:, : self.cutoffs[0]] = head_logprob[:, : self.cutoffs[0]]
- else:
- weight_i, bias_i, proj_i = weights[i], biases[i], self.out_projs[i]
-
- tail_logit_i = self._compute_logit(hidden, weight_i, bias_i, proj_i)
- tail_logprob_i = nn.functional.log_softmax(tail_logit_i, dim=1)
-
- logprob_i = head_logprob[:, -i] + tail_logprob_i
- out[:, start_idx, stop_idx] = logprob_i
-
- return out
diff --git a/transformers/models/deprecated/transfo_xl/tokenization_transfo_xl.py b/transformers/models/deprecated/transfo_xl/tokenization_transfo_xl.py
deleted file mode 100644
index 7290a7a83b85660069bc2b88e4ba2734114f3f9b..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/transfo_xl/tokenization_transfo_xl.py
+++ /dev/null
@@ -1,819 +0,0 @@
-# coding=utf-8
-# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
- Tokenization classes for Transformer XL model. Adapted from https://github.com/kimiyoung/transformer-xl.
-"""
-
-
-import glob
-import os
-import pickle
-import re
-from collections import Counter, OrderedDict
-from typing import List, Optional, Tuple
-
-import numpy as np
-
-from ....tokenization_utils import PreTrainedTokenizer
-from ....utils import (
- cached_file,
- is_sacremoses_available,
- is_torch_available,
- logging,
- requires_backends,
- strtobool,
- torch_only_method,
-)
-
-
-if is_sacremoses_available():
- import sacremoses as sm
-
-
-if is_torch_available():
- import torch
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {
- "pretrained_vocab_file": "vocab.pkl",
- "pretrained_vocab_file_torch": "vocab.bin",
- "vocab_file": "vocab.txt",
-}
-
-
-PRETRAINED_CORPUS_ARCHIVE_MAP = {
- "transfo-xl/transfo-xl-wt103": "https://huggingface.co/transfo-xl/transfo-xl-wt103/resolve/main/corpus.bin",
-}
-CORPUS_NAME = "corpus.bin"
-
-MATCH_NUMBERS = r"(?<=\d)[,.](?=\d)", r" @\g<0>@ "
-DETOKENIZE_NUMBERS = [(r" @\,@ ", r","), (r" @\.@ ", r".")]
-
-
-def tokenize_numbers(text_array: List[str]) -> List[str]:
- """
- Splits large comma-separated numbers and floating point values. This is done by replacing commas with ' @,@ ' and
- dots with ' @.@ '.
-
- Args:
- text_array: An already tokenized text as list.
-
- Returns:
- A list of strings with tokenized numbers.
-
- Example:
-
- ```python
- >>> tokenize_numbers(["$", "5,000", "1.73", "m"])
- ['$', '5', '@,@', '000', '1', '@.@', '73', 'm']
- ```"""
- tokenized = []
- for i in range(len(text_array)):
- reg, sub = MATCH_NUMBERS
- replaced = re.sub(reg, sub, text_array[i]).split()
- tokenized.extend(replaced)
-
- return tokenized
-
-
-def detokenize_numbers(text: str) -> str:
- """
- Inverts the operation of *tokenize_numbers*. This is replacing ' @,@ ' and ' @.@' by ',' and '.'.
-
- Args:
- text: A string where the number should be detokenized.
-
- Returns:
- A detokenized string.
-
- Example:
-
- ```python
- >>> detokenize_numbers("$ 5 @,@ 000 1 @.@ 73 m")
- '$ 5,000 1.73 m'
- ```"""
- for reg, sub in DETOKENIZE_NUMBERS:
- text = re.sub(reg, sub, text)
- return text
-
-
-class TransfoXLTokenizer(PreTrainedTokenizer):
- """
- Construct a Transformer-XL tokenizer adapted from Vocab class in [the original
- code](https://github.com/kimiyoung/transformer-xl). The Transformer-XL tokenizer is a word-level tokenizer (no
- sub-word tokenization).
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- special (`List[str]`, *optional*):
- A list of special tokens (to be treated by the original implementation of this tokenizer).
- min_freq (`int`, *optional*, defaults to 0):
- The minimum number of times a token has to be present in order to be kept in the vocabulary (otherwise it
- will be mapped to `unk_token`).
- max_size (`int`, *optional*):
- The maximum size of the vocabulary. If left unset, it will default to the size of the vocabulary found
- after excluding the tokens according to the `min_freq` rule.
- lower_case (`bool`, *optional*, defaults to `False`):
- Whether or not to lowercase the input when tokenizing.
- delimiter (`str`, *optional*):
- The delimiter used between tokens.
- vocab_file (`str`, *optional*):
- File containing the vocabulary (from the original implementation).
- pretrained_vocab_file (`str`, *optional*):
- File containing the vocabulary as saved with the `save_pretrained()` method.
- never_split (`List[str]`, *optional*):
- List of tokens that should never be split. If no list is specified, will simply use the existing special
- tokens.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sequence token.
- additional_special_tokens (`List[str]`, *optional*, defaults to `['']`):
- A list of additional special tokens (for the HuggingFace functionality).
- language (`str`, *optional*, defaults to `"en"`):
- The language of this tokenizer (used for mose preprocessing).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids"]
-
- def __init__(
- self,
- special=None,
- min_freq=0,
- max_size=None,
- lower_case=False,
- delimiter=None,
- vocab_file=None,
- pretrained_vocab_file: str = None,
- never_split=None,
- unk_token="",
- eos_token="",
- additional_special_tokens=[""],
- language="en",
- **kwargs,
- ):
- logger.error(
- "`TransfoXL` was deprecated due to security issues linked to `pickle.load` in `TransfoXLTokenizer`. "
- "See more details on this model's documentation page: "
- "`https://github.com/huggingface/transformers/blob/main/docs/source/en/model_doc/transfo-xl.md`."
- )
-
- requires_backends(self, "sacremoses")
- if special is None:
- special = []
- self.counter = Counter()
- self.special = special
- self.min_freq = min_freq
- self.max_size = max_size
- self.lower_case = lower_case
- self.delimiter = delimiter
- self.vocab_file = vocab_file
- self.punctuation_symbols = '!"#$%&()*+,-./\\:;<=>?@[\\]^_`{|}~'
- self.punction_without_space_before_pattern = re.compile(rf"[^\s][{self.punctuation_symbols}]")
- self.punctuation_with_space_around_pattern = self._compile_space_around_punctuation_pattern()
- self.language = language
- self.moses_punct_normalizer = sm.MosesPunctNormalizer(language)
- self.moses_tokenizer = sm.MosesTokenizer(language)
- self.moses_detokenizer = sm.MosesDetokenizer(language)
- self.idx2sym = []
- self.sym2idx = OrderedDict()
- # This try... catch... is not beautiful but honestly this tokenizer was not made to be used
- # in a library like ours, at all.
- try:
- vocab_dict = None
- if pretrained_vocab_file is not None:
- # Priority on pickle files (support PyTorch and TF)
- if not strtobool(os.environ.get("TRUST_REMOTE_CODE", "False")):
- raise ValueError(
- "This part uses `pickle.load` which is insecure and will execute arbitrary code that is "
- "potentially malicious. It's recommended to never unpickle data that could have come from an "
- "untrusted source, or that could have been tampered with. If you already verified the pickle "
- "data and decided to use it, you can set the environment variable "
- "`TRUST_REMOTE_CODE` to `True` to allow it."
- )
- with open(pretrained_vocab_file, "rb") as f:
- vocab_dict = pickle.load(f)
-
- # Loading a torch-saved transfo-xl vocab dict with pickle results in an integer
- # Entering this if statement means that we tried to load a torch-saved file with pickle, and we failed.
- # We therefore load it with torch, if it's available.
- if isinstance(vocab_dict, int):
- if not is_torch_available():
- raise ImportError(
- "Not trying to load dict with PyTorch as you need to install pytorch to load "
- "from a PyTorch pretrained vocabulary, "
- "or activate it with environment variables USE_TORCH=1 and USE_TF=0."
- )
- vocab_dict = torch.load(pretrained_vocab_file)
-
- if vocab_dict is not None:
- for key, value in vocab_dict.items():
- if key not in self.__dict__ or key in ["sym2idx", "idx2sym"]:
- self.__dict__[key] = value
- elif vocab_file is not None:
- self.build_vocab()
-
- except Exception as e:
- raise ValueError(
- f"Unable to parse file {pretrained_vocab_file}. Unknown format. "
- "If you tried to load a model saved through TransfoXLTokenizerFast, "
- "please note they are not compatible."
- ) from e
-
- if vocab_file is not None:
- self.build_vocab()
-
- super().__init__(
- special=special,
- min_freq=min_freq,
- max_size=max_size,
- lower_case=lower_case,
- delimiter=delimiter,
- vocab_file=vocab_file,
- pretrained_vocab_file=pretrained_vocab_file,
- never_split=never_split,
- unk_token=unk_token,
- eos_token=eos_token,
- additional_special_tokens=additional_special_tokens,
- language=language,
- **kwargs,
- )
-
- # these are not required to initialize the parent class as only used when tokenizing.
- if never_split is None:
- never_split = self.all_special_tokens
- self.never_split = never_split
-
- @property
- def do_lower_case(self):
- return self.lower_case
-
- def _compile_space_around_punctuation_pattern(self):
- look_ahead_for_special_token = f"(?=[{self.punctuation_symbols}])"
- look_ahead_to_match_all_except_space = r"(?=[^\s])"
- return re.compile(r"" + look_ahead_for_special_token + look_ahead_to_match_all_except_space)
-
- def count_file(self, path, verbose=False, add_eos=False):
- if verbose:
- logger.info(f"counting file {path} ...")
- assert os.path.exists(path), f"Input file {path} not found"
-
- sents = []
- with open(path, "r", encoding="utf-8") as f:
- for idx, line in enumerate(f):
- if verbose and idx > 0 and idx % 500000 == 0:
- logger.info(f" line {idx}")
- symbols = self.tokenize(line, add_eos=add_eos)
- self.counter.update(symbols)
- sents.append(symbols)
-
- return sents
-
- def count_sents(self, sents, verbose=False):
- """
- sents : a list of sentences, each a list of tokenized symbols
- """
- if verbose:
- logger.info(f"counting {len(sents)} sents ...")
- for idx, symbols in enumerate(sents):
- if verbose and idx > 0 and idx % 500000 == 0:
- logger.info(f" line {idx}")
- self.counter.update(symbols)
-
- def _build_from_file(self, vocab_file):
- self.idx2sym = []
- self.sym2idx = OrderedDict()
-
- with open(vocab_file, "r", encoding="utf-8") as f:
- for line in f:
- symb = line.strip().split()[0]
- self.add_symbol(symb)
- if "" in self.sym2idx:
- self.unk_idx = self.sym2idx[""]
- elif "" in self.sym2idx:
- self.unk_idx = self.sym2idx[""]
- else:
- raise ValueError("Token not in vocabulary and no token in vocabulary for replacement.")
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory,
- (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["pretrained_vocab_file"],
- )
- else:
- vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
- with open(vocab_file, "wb") as f:
- pickle.dump(self.__dict__, f)
- return (vocab_file,)
-
- def build_vocab(self):
- if self.vocab_file:
- logger.info(f"building vocab from {self.vocab_file}")
- self._build_from_file(self.vocab_file)
- logger.info(f"Final vocab size {len(self.sym2idx)}")
- else:
- logger.info(f"building vocab with min_freq={self.min_freq}, max_size={self.max_size}")
- self.idx2sym = []
- self.sym2idx = OrderedDict()
-
- for sym in self.special:
- self.add_special(sym)
-
- for sym, cnt in self.counter.most_common(self.max_size):
- if cnt < self.min_freq:
- break
- self.add_symbol(sym)
-
- logger.info(f"Final vocab size {len(self.sym2idx)} from {len(self.counter)} unique tokens")
-
- @torch_only_method
- def encode_file(self, path, ordered=False, verbose=False, add_eos=True, add_double_eos=False):
- if verbose:
- logger.info(f"encoding file {path} ...")
- assert os.path.exists(path), f"Output file {path} not found"
- encoded = []
- with open(path, "r", encoding="utf-8") as f:
- for idx, line in enumerate(f):
- if verbose and idx > 0 and idx % 500000 == 0:
- logger.info(f" line {idx}")
- symbols = self.tokenize(line, add_eos=add_eos, add_double_eos=add_double_eos)
- encoded.append(self.convert_to_tensor(symbols))
-
- if ordered:
- encoded = torch.cat(encoded)
-
- return encoded
-
- @torch_only_method
- def encode_sents(self, sents, ordered=False, verbose=False):
- if verbose:
- logger.info(f"encoding {len(sents)} sents ...")
- encoded = []
- for idx, symbols in enumerate(sents):
- if verbose and idx > 0 and idx % 500000 == 0:
- logger.info(f" line {idx}")
- encoded.append(self.convert_to_tensor(symbols))
-
- if ordered:
- encoded = torch.cat(encoded)
-
- return encoded
-
- def add_special(self, sym):
- if sym not in self.sym2idx:
- self.idx2sym.append(sym)
- self.sym2idx[sym] = len(self.idx2sym) - 1
- setattr(self, f"{sym.strip('<>')}_idx", self.sym2idx[sym])
-
- def add_symbol(self, sym):
- if sym not in self.sym2idx:
- self.idx2sym.append(sym)
- self.sym2idx[sym] = len(self.idx2sym) - 1
-
- def move_added_token(self, token: str, target_idx: int):
- """
- Moves an added token to a specific position in the vocab. This method should be used when resizing an embedding
- layer other than the last one in the `AdaptiveEmbedding` in order to move the token in the tokenizer from the
- default position (at the very end) to the desired one.
-
- Args:
- token: The token to move to a specific position in the vocab.
- target_idx: The position where the token should be moved to.
- """
- assert token in self.added_tokens_encoder, "Token which should be moved has to be an added token"
- assert token not in self.idx2sym, "Token which should be moved is already in vocab"
-
- # Insert sym into vocab
- self.idx2sym.insert(target_idx, token)
- self.sym2idx[token] = target_idx
-
- # Shift following indices in sym2idx
- for idx in range(target_idx + 1, len(self.idx2sym)):
- current_sym = self.idx2sym[idx]
- self.sym2idx[current_sym] = idx
-
- # Delete token from added_tokens
- old_index = self._added_tokens_encoder.pop(token)
- self._added_tokens_decoder.pop(old_index)
-
- def moses_punct_norm(self, text):
- return self.moses_punct_normalizer.normalize(text)
-
- def moses_tokenize(self, text):
- return self.moses_tokenizer.tokenize(
- text, aggressive_dash_splits=True, return_str=False, escape=False, protected_patterns=self.never_split
- )
-
- def moses_pipeline(self, text: str) -> List[str]:
- """
- Does basic tokenization using [`sacremoses.MosesPunctNormalizer`] and [`sacremoses.MosesTokenizer`] with
- *aggressive_dash_splits=True* (see [`sacremoses.tokenize.MosesTokenizer.tokenize`]). Additionally, large
- comma-separated numbers and floating point values are split. E.g. "23,000 people are 1.80m tall" -> "23 @,@ 000
- people are 1 @.@ 80m tall"
-
- Args:
- text: Text to be tokenize
-
- Returns:
- A list of tokenized string
-
- Example:
-
- ```python
- >>> tokenizer = TransfoXLTokenizer.from_pretrained("transfo-xl/transfo-xl-wt103")
- >>> tokenizer.moses_pipeline("23,000 people are 1.80 m tall")
- ['23', '@,@', '000', 'people', 'are', '1', '@.@', '80', 'm', 'tall']
- ```"""
- text = self.moses_punct_norm(text)
- text = self.moses_tokenize(text)
- text = tokenize_numbers(text)
- return text
-
- def _convert_id_to_token(self, idx):
- """Converts an id in a token (BPE) using the vocab."""
- assert 0 <= idx < len(self), f"Index {idx} out of vocabulary range"
- return self.idx2sym[idx]
-
- def _convert_token_to_id(self, sym):
- """Converts a token (str) in an id using the vocab."""
- if sym in self.sym2idx:
- return self.sym2idx[sym]
- else:
- # logger.info(f'encounter unk {sym}')
- # assert '' not in sym
- if hasattr(self, "unk_idx"):
- return self.sym2idx.get(sym, self.unk_idx)
- # Backward compatibility with pre-trained models
- elif "" in self.sym2idx:
- return self.sym2idx[""]
- elif "" in self.sym2idx:
- return self.sym2idx[""]
- else:
- raise ValueError("Token not in vocabulary and no token in vocabulary for replacement.")
-
- def convert_tokens_to_string(self, tokens):
- """
- Converts a sequence of tokens (string) in a single string. Additionally, the split numbers are converted back
- into it's original form.
- """
- out_string = self.moses_detokenizer.detokenize(tokens)
- return detokenize_numbers(out_string).strip()
-
- @torch_only_method
- def convert_to_tensor(self, symbols):
- return torch.LongTensor(self.convert_tokens_to_ids(symbols))
-
- @property
- def vocab_size(self):
- return len(self.idx2sym)
-
- def get_vocab(self):
- vocab = self.sym2idx.copy()
- vocab.update(self.added_tokens_encoder)
- return vocab
-
- def _tokenize(self, line, add_eos=False, add_double_eos=False):
- line = line.strip()
- # convert to lower case
- if self.lower_case:
- line = line.lower()
-
- # empty delimiter '' will evaluate False
- if self.delimiter == "":
- symbols = line
- else:
- symbols = self.moses_pipeline(line)
-
- if add_double_eos: # lm1b
- return [""] + symbols + [""]
- elif add_eos:
- return symbols + [""]
- else:
- return symbols
-
-
-class LMOrderedIterator(object):
- def __init__(self, data, bsz, bptt, device="cpu", ext_len=None):
- """
- data -- LongTensor -- the LongTensor is strictly ordered
- """
- self.bsz = bsz
- self.bptt = bptt
- self.ext_len = ext_len if ext_len is not None else 0
-
- self.device = device
-
- # Work out how cleanly we can divide the dataset into bsz parts.
- self.n_step = data.size(0) // bsz
-
- # Trim off any extra elements that wouldn't cleanly fit (remainders).
- data = data.narrow(0, 0, self.n_step * bsz)
-
- # Evenly divide the data across the bsz batches.
- self.data = data.view(bsz, -1).t().contiguous().to(device)
-
- # Number of mini-batches
- self.n_batch = (self.n_step + self.bptt - 1) // self.bptt
-
- def get_batch(self, i, bptt=None):
- if bptt is None:
- bptt = self.bptt
- seq_len = min(bptt, self.data.size(0) - 1 - i)
-
- end_idx = i + seq_len
- beg_idx = max(0, i - self.ext_len)
-
- data = self.data[beg_idx:end_idx]
- target = self.data[i + 1 : i + 1 + seq_len]
-
- data_out = data.transpose(0, 1).contiguous().to(self.device)
- target_out = target.transpose(0, 1).contiguous().to(self.device)
-
- return data_out, target_out, seq_len
-
- def get_fixlen_iter(self, start=0):
- for i in range(start, self.data.size(0) - 1, self.bptt):
- yield self.get_batch(i)
-
- def get_varlen_iter(self, start=0, std=5, min_len=5, max_deviation=3):
- max_len = self.bptt + max_deviation * std
- i = start
- while True:
- bptt = self.bptt if np.random.random() < 0.95 else self.bptt / 2.0
- bptt = min(max_len, max(min_len, int(np.random.normal(bptt, std))))
- data, target, seq_len = self.get_batch(i, bptt)
- i += seq_len
- yield data, target, seq_len
- if i >= self.data.size(0) - 2:
- break
-
- def __iter__(self):
- return self.get_fixlen_iter()
-
-
-class LMShuffledIterator(object):
- def __init__(self, data, bsz, bptt, device="cpu", ext_len=None, shuffle=False):
- """
- data -- list[LongTensor] -- there is no order among the LongTensors
- """
- self.data = data
-
- self.bsz = bsz
- self.bptt = bptt
- self.ext_len = ext_len if ext_len is not None else 0
-
- self.device = device
- self.shuffle = shuffle
-
- def get_sent_stream(self):
- # index iterator
- epoch_indices = np.random.permutation(len(self.data)) if self.shuffle else np.array(range(len(self.data)))
-
- # sentence iterator
- for idx in epoch_indices:
- yield self.data[idx]
-
- @torch_only_method
- def stream_iterator(self, sent_stream):
- # streams for each data in the batch
- streams = [None] * self.bsz
-
- data = torch.LongTensor(self.bptt, self.bsz)
- target = torch.LongTensor(self.bptt, self.bsz)
-
- n_retain = 0
-
- while True:
- # data : [n_retain+bptt x bsz]
- # target : [bptt x bsz]
- data[n_retain:].fill_(-1)
- target.fill_(-1)
-
- valid_batch = True
-
- for i in range(self.bsz):
- n_filled = 0
- try:
- while n_filled < self.bptt:
- if streams[i] is None or len(streams[i]) <= 1:
- streams[i] = next(sent_stream)
- # number of new tokens to fill in
- n_new = min(len(streams[i]) - 1, self.bptt - n_filled)
- # first n_retain tokens are retained from last batch
- data[n_retain + n_filled : n_retain + n_filled + n_new, i] = streams[i][:n_new]
- target[n_filled : n_filled + n_new, i] = streams[i][1 : n_new + 1]
- streams[i] = streams[i][n_new:]
- n_filled += n_new
- except StopIteration:
- valid_batch = False
- break
-
- if not valid_batch:
- return
-
- data_out = data.transpose(0, 1).contiguous().to(self.device)
- target_out = target.transpose(0, 1).contiguous().to(self.device)
-
- yield data_out, target_out, self.bptt
-
- n_retain = min(data.size(0), self.ext_len)
- if n_retain > 0:
- data[:n_retain] = data[-n_retain:]
- data.resize_(n_retain + self.bptt, data.size(1))
-
- def __iter__(self):
- # sent_stream is an iterator
- sent_stream = self.get_sent_stream()
-
- for batch in self.stream_iterator(sent_stream):
- yield batch
-
-
-class LMMultiFileIterator(LMShuffledIterator):
- def __init__(self, paths, vocab, bsz, bptt, device="cpu", ext_len=None, shuffle=False):
- self.paths = paths
- self.vocab = vocab
-
- self.bsz = bsz
- self.bptt = bptt
- self.ext_len = ext_len if ext_len is not None else 0
-
- self.device = device
- self.shuffle = shuffle
-
- def get_sent_stream(self, path):
- sents = self.vocab.encode_file(path, add_double_eos=True)
- if self.shuffle:
- np.random.shuffle(sents)
- sent_stream = iter(sents)
-
- return sent_stream
-
- def __iter__(self):
- if self.shuffle:
- np.random.shuffle(self.paths)
-
- for path in self.paths:
- # sent_stream is an iterator
- sent_stream = self.get_sent_stream(path)
- for batch in self.stream_iterator(sent_stream):
- yield batch
-
-
-class TransfoXLCorpus(object):
- @classmethod
- @torch_only_method
- def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None, *inputs, **kwargs):
- """
- Instantiate a pre-processed corpus.
- """
- vocab = TransfoXLTokenizer.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
- is_local = os.path.isdir(pretrained_model_name_or_path)
- # redirect to the cache, if necessary
- try:
- resolved_corpus_file = cached_file(pretrained_model_name_or_path, CORPUS_NAME, cache_dir=cache_dir)
- except EnvironmentError:
- logger.error(
- f"Corpus '{pretrained_model_name_or_path}' was not found in corpus list"
- f" ({', '.join(PRETRAINED_CORPUS_ARCHIVE_MAP.keys())}. We assumed '{pretrained_model_name_or_path}'"
- f" was a path or url but couldn't find files {CORPUS_NAME} at this path or url."
- )
- return None
- if is_local:
- logger.info(f"loading corpus file {resolved_corpus_file}")
- else:
- logger.info(f"loading corpus file {CORPUS_NAME} from cache at {resolved_corpus_file}")
-
- # Instantiate tokenizer.
- corpus = cls(*inputs, **kwargs)
- corpus_dict = torch.load(resolved_corpus_file)
- for key, value in corpus_dict.items():
- corpus.__dict__[key] = value
- corpus.vocab = vocab
- if corpus.train is not None:
- corpus.train = torch.tensor(corpus.train, dtype=torch.long)
- if corpus.valid is not None:
- corpus.valid = torch.tensor(corpus.valid, dtype=torch.long)
- if corpus.test is not None:
- corpus.test = torch.tensor(corpus.test, dtype=torch.long)
- return corpus
-
- def __init__(self, *args, **kwargs):
- self.vocab = TransfoXLTokenizer(*args, **kwargs)
- self.dataset = None
- self.train = None
- self.valid = None
- self.test = None
-
- def build_corpus(self, path, dataset):
- self.dataset = dataset
-
- if self.dataset in ["ptb", "wt2", "enwik8", "text8"]:
- self.vocab.count_file(os.path.join(path, "train.txt"))
- self.vocab.count_file(os.path.join(path, "valid.txt"))
- self.vocab.count_file(os.path.join(path, "test.txt"))
- elif self.dataset == "wt103":
- self.vocab.count_file(os.path.join(path, "train.txt"))
- elif self.dataset == "lm1b":
- train_path_pattern = os.path.join(
- path,
- "1-billion-word-language-modeling-benchmark-r13output",
- "training-monolingual.tokenized.shuffled",
- "news.en-*",
- )
- train_paths = glob.glob(train_path_pattern)
- # the vocab will load from file when build_vocab() is called
-
- self.vocab.build_vocab()
-
- if self.dataset in ["ptb", "wt2", "wt103"]:
- self.train = self.vocab.encode_file(os.path.join(path, "train.txt"), ordered=True)
- self.valid = self.vocab.encode_file(os.path.join(path, "valid.txt"), ordered=True)
- self.test = self.vocab.encode_file(os.path.join(path, "test.txt"), ordered=True)
- elif self.dataset in ["enwik8", "text8"]:
- self.train = self.vocab.encode_file(os.path.join(path, "train.txt"), ordered=True, add_eos=False)
- self.valid = self.vocab.encode_file(os.path.join(path, "valid.txt"), ordered=True, add_eos=False)
- self.test = self.vocab.encode_file(os.path.join(path, "test.txt"), ordered=True, add_eos=False)
- elif self.dataset == "lm1b":
- self.train = train_paths
- self.valid = self.vocab.encode_file(os.path.join(path, "valid.txt"), ordered=False, add_double_eos=True)
- self.test = self.vocab.encode_file(os.path.join(path, "test.txt"), ordered=False, add_double_eos=True)
-
- def get_iterator(self, split, *args, **kwargs):
- if split == "train":
- if self.dataset in ["ptb", "wt2", "wt103", "enwik8", "text8"]:
- data_iter = LMOrderedIterator(self.train, *args, **kwargs)
- elif self.dataset == "lm1b":
- kwargs["shuffle"] = True
- data_iter = LMMultiFileIterator(self.train, self.vocab, *args, **kwargs)
- elif split in ["valid", "test"]:
- data = self.valid if split == "valid" else self.test
- if self.dataset in ["ptb", "wt2", "wt103", "enwik8", "text8"]:
- data_iter = LMOrderedIterator(data, *args, **kwargs)
- elif self.dataset == "lm1b":
- data_iter = LMShuffledIterator(data, *args, **kwargs)
- else:
- data_iter = None
- raise ValueError(f"Split not recognized: {split}")
-
- return data_iter
-
-
-@torch_only_method
-def get_lm_corpus(datadir, dataset):
- fn = os.path.join(datadir, "cache.pt")
- fn_pickle = os.path.join(datadir, "cache.pkl")
- if os.path.exists(fn):
- logger.info("Loading cached dataset...")
- corpus = torch.load(fn_pickle)
- elif os.path.exists(fn):
- logger.info("Loading cached dataset from pickle...")
- if not strtobool(os.environ.get("TRUST_REMOTE_CODE", "False")):
- raise ValueError(
- "This part uses `pickle.load` which is insecure and will execute arbitrary code that is potentially "
- "malicious. It's recommended to never unpickle data that could have come from an untrusted source, or "
- "that could have been tampered with. If you already verified the pickle data and decided to use it, "
- "you can set the environment variable `TRUST_REMOTE_CODE` to `True` to allow it."
- )
- with open(fn, "rb") as fp:
- corpus = pickle.load(fp)
- else:
- logger.info(f"Producing dataset {dataset}...")
- kwargs = {}
- if dataset in ["wt103", "wt2"]:
- kwargs["special"] = [""]
- kwargs["lower_case"] = False
- elif dataset == "ptb":
- kwargs["special"] = [""]
- kwargs["lower_case"] = True
- elif dataset == "lm1b":
- kwargs["special"] = []
- kwargs["lower_case"] = False
- kwargs["vocab_file"] = os.path.join(datadir, "1b_word_vocab.txt")
- elif dataset in ["enwik8", "text8"]:
- pass
-
- corpus = TransfoXLCorpus(datadir, dataset, **kwargs)
- torch.save(corpus, fn)
-
- return corpus
diff --git a/transformers/models/deprecated/van/__init__.py b/transformers/models/deprecated/van/__init__.py
deleted file mode 100644
index 2db730984ffa031458589f1cc6c6c1944eba0e91..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/van/__init__.py
+++ /dev/null
@@ -1,54 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {"configuration_van": ["VAN_PRETRAINED_CONFIG_ARCHIVE_MAP", "VanConfig"]}
-
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_van"] = [
- "VAN_PRETRAINED_MODEL_ARCHIVE_LIST",
- "VanForImageClassification",
- "VanModel",
- "VanPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_van import VAN_PRETRAINED_CONFIG_ARCHIVE_MAP, VanConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_van import (
- VAN_PRETRAINED_MODEL_ARCHIVE_LIST,
- VanForImageClassification,
- VanModel,
- VanPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
diff --git a/transformers/models/deprecated/van/__pycache__/__init__.cpython-310.pyc b/transformers/models/deprecated/van/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index ffdbaa3c29d7dc35df35ebb4d924d278263e0320..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/van/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/van/__pycache__/configuration_van.cpython-310.pyc b/transformers/models/deprecated/van/__pycache__/configuration_van.cpython-310.pyc
deleted file mode 100644
index 41e4454fd80c1721545d43719d844766ef65a665..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/van/__pycache__/configuration_van.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/van/__pycache__/convert_van_to_pytorch.cpython-310.pyc b/transformers/models/deprecated/van/__pycache__/convert_van_to_pytorch.cpython-310.pyc
deleted file mode 100644
index a6500e58fb56253d91164acc0e8aff63f0903615..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/van/__pycache__/convert_van_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/van/__pycache__/modeling_van.cpython-310.pyc b/transformers/models/deprecated/van/__pycache__/modeling_van.cpython-310.pyc
deleted file mode 100644
index da20bd23dbb44bd49520e1f865d144eb4dabaca5..0000000000000000000000000000000000000000
Binary files a/transformers/models/deprecated/van/__pycache__/modeling_van.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deprecated/van/configuration_van.py b/transformers/models/deprecated/van/configuration_van.py
deleted file mode 100644
index f58d0215694a934fd9b654d1dc25e705fd78ccb3..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/van/configuration_van.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" VAN model configuration"""
-
-from ....configuration_utils import PretrainedConfig
-from ....utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from .._archive_maps import VAN_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class VanConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`VanModel`]. It is used to instantiate a VAN model
- according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the VAN
- [Visual-Attention-Network/van-base](https://huggingface.co/Visual-Attention-Network/van-base) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- patch_sizes (`List[int]`, *optional*, defaults to `[7, 3, 3, 3]`):
- Patch size to use in each stage's embedding layer.
- strides (`List[int]`, *optional*, defaults to `[4, 2, 2, 2]`):
- Stride size to use in each stage's embedding layer to downsample the input.
- hidden_sizes (`List[int]`, *optional*, defaults to `[64, 128, 320, 512]`):
- Dimensionality (hidden size) at each stage.
- depths (`List[int]`, *optional*, defaults to `[3, 3, 12, 3]`):
- Depth (number of layers) for each stage.
- mlp_ratios (`List[int]`, *optional*, defaults to `[8, 8, 4, 4]`):
- The expansion ratio for mlp layer at each stage.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in each layer. If string, `"gelu"`, `"relu"`,
- `"selu"` and `"gelu_new"` are supported.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-06):
- The epsilon used by the layer normalization layers.
- layer_scale_init_value (`float`, *optional*, defaults to 0.01):
- The initial value for layer scaling.
- drop_path_rate (`float`, *optional*, defaults to 0.0):
- The dropout probability for stochastic depth.
- dropout_rate (`float`, *optional*, defaults to 0.0):
- The dropout probability for dropout.
-
- Example:
- ```python
- >>> from transformers import VanModel, VanConfig
-
- >>> # Initializing a VAN van-base style configuration
- >>> configuration = VanConfig()
- >>> # Initializing a model from the van-base style configuration
- >>> model = VanModel(configuration)
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "van"
-
- def __init__(
- self,
- image_size=224,
- num_channels=3,
- patch_sizes=[7, 3, 3, 3],
- strides=[4, 2, 2, 2],
- hidden_sizes=[64, 128, 320, 512],
- depths=[3, 3, 12, 3],
- mlp_ratios=[8, 8, 4, 4],
- hidden_act="gelu",
- initializer_range=0.02,
- layer_norm_eps=1e-6,
- layer_scale_init_value=1e-2,
- drop_path_rate=0.0,
- dropout_rate=0.0,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.image_size = image_size
- self.num_channels = num_channels
- self.patch_sizes = patch_sizes
- self.strides = strides
- self.hidden_sizes = hidden_sizes
- self.depths = depths
- self.mlp_ratios = mlp_ratios
- self.hidden_act = hidden_act
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.layer_scale_init_value = layer_scale_init_value
- self.drop_path_rate = drop_path_rate
- self.dropout_rate = dropout_rate
diff --git a/transformers/models/deprecated/van/convert_van_to_pytorch.py b/transformers/models/deprecated/van/convert_van_to_pytorch.py
deleted file mode 100644
index 20492e42be2043d50e39b7573fc4e9fca05c7d32..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/van/convert_van_to_pytorch.py
+++ /dev/null
@@ -1,291 +0,0 @@
-# coding=utf-8
-# Copyright 2022 BNRist (Tsinghua University), TKLNDST (Nankai University) and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert VAN checkpoints from the original repository.
-
-URL: https://github.com/Visual-Attention-Network/VAN-Classification"""
-
-
-import argparse
-import json
-import sys
-from dataclasses import dataclass, field
-from functools import partial
-from pathlib import Path
-from typing import List
-
-import torch
-import torch.nn as nn
-from huggingface_hub import cached_download, hf_hub_download
-from torch import Tensor
-
-from transformers import AutoImageProcessor, VanConfig, VanForImageClassification
-from transformers.models.deprecated.van.modeling_van import VanLayerScaling
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-@dataclass
-class Tracker:
- module: nn.Module
- traced: List[nn.Module] = field(default_factory=list)
- handles: list = field(default_factory=list)
-
- def _forward_hook(self, m, inputs: Tensor, outputs: Tensor):
- has_not_submodules = len(list(m.modules())) == 1 or isinstance(m, nn.Conv2d) or isinstance(m, nn.BatchNorm2d)
- if has_not_submodules:
- if not isinstance(m, VanLayerScaling):
- self.traced.append(m)
-
- def __call__(self, x: Tensor):
- for m in self.module.modules():
- self.handles.append(m.register_forward_hook(self._forward_hook))
- self.module(x)
- [x.remove() for x in self.handles]
- return self
-
- @property
- def parametrized(self):
- # check the len of the state_dict keys to see if we have learnable params
- return list(filter(lambda x: len(list(x.state_dict().keys())) > 0, self.traced))
-
-
-@dataclass
-class ModuleTransfer:
- src: nn.Module
- dest: nn.Module
- verbose: int = 0
- src_skip: List = field(default_factory=list)
- dest_skip: List = field(default_factory=list)
-
- def __call__(self, x: Tensor):
- """
- Transfer the weights of `self.src` to `self.dest` by performing a forward pass using `x` as input. Under the
- hood we tracked all the operations in both modules.
- """
- dest_traced = Tracker(self.dest)(x).parametrized
- src_traced = Tracker(self.src)(x).parametrized
-
- src_traced = list(filter(lambda x: type(x) not in self.src_skip, src_traced))
- dest_traced = list(filter(lambda x: type(x) not in self.dest_skip, dest_traced))
-
- if len(dest_traced) != len(src_traced):
- raise Exception(
- f"Numbers of operations are different. Source module has {len(src_traced)} operations while"
- f" destination module has {len(dest_traced)}."
- )
-
- for dest_m, src_m in zip(dest_traced, src_traced):
- dest_m.load_state_dict(src_m.state_dict())
- if self.verbose == 1:
- print(f"Transfered from={src_m} to={dest_m}")
-
-
-def copy_parameters(from_model: nn.Module, our_model: nn.Module) -> nn.Module:
- # nn.Parameter cannot be tracked by the Tracker, thus we need to manually convert them
- from_state_dict = from_model.state_dict()
- our_state_dict = our_model.state_dict()
- config = our_model.config
- all_keys = []
- for stage_idx in range(len(config.hidden_sizes)):
- for block_id in range(config.depths[stage_idx]):
- from_key = f"block{stage_idx + 1}.{block_id}.layer_scale_1"
- to_key = f"van.encoder.stages.{stage_idx}.layers.{block_id}.attention_scaling.weight"
-
- all_keys.append((from_key, to_key))
- from_key = f"block{stage_idx + 1}.{block_id}.layer_scale_2"
- to_key = f"van.encoder.stages.{stage_idx}.layers.{block_id}.mlp_scaling.weight"
-
- all_keys.append((from_key, to_key))
-
- for from_key, to_key in all_keys:
- our_state_dict[to_key] = from_state_dict.pop(from_key)
-
- our_model.load_state_dict(our_state_dict)
- return our_model
-
-
-def convert_weight_and_push(
- name: str,
- config: VanConfig,
- checkpoint: str,
- from_model: nn.Module,
- save_directory: Path,
- push_to_hub: bool = True,
-):
- print(f"Downloading weights for {name}...")
- checkpoint_path = cached_download(checkpoint)
- print(f"Converting {name}...")
- from_state_dict = torch.load(checkpoint_path)["state_dict"]
- from_model.load_state_dict(from_state_dict)
- from_model.eval()
- with torch.no_grad():
- our_model = VanForImageClassification(config).eval()
- module_transfer = ModuleTransfer(src=from_model, dest=our_model)
- x = torch.randn((1, 3, 224, 224))
- module_transfer(x)
- our_model = copy_parameters(from_model, our_model)
-
- if not torch.allclose(from_model(x), our_model(x).logits):
- raise ValueError("The model logits don't match the original one.")
-
- checkpoint_name = name
- print(checkpoint_name)
-
- if push_to_hub:
- our_model.push_to_hub(
- repo_path_or_name=save_directory / checkpoint_name,
- commit_message="Add model",
- use_temp_dir=True,
- )
-
- # we can use the convnext one
- image_processor = AutoImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k")
- image_processor.push_to_hub(
- repo_path_or_name=save_directory / checkpoint_name,
- commit_message="Add image processor",
- use_temp_dir=True,
- )
-
- print(f"Pushed {checkpoint_name}")
-
-
-def convert_weights_and_push(save_directory: Path, model_name: str = None, push_to_hub: bool = True):
- filename = "imagenet-1k-id2label.json"
- num_labels = 1000
-
- repo_id = "huggingface/label-files"
- num_labels = num_labels
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
-
- id2label = id2label
- label2id = {v: k for k, v in id2label.items()}
-
- ImageNetPreTrainedConfig = partial(VanConfig, num_labels=num_labels, id2label=id2label, label2id=label2id)
-
- names_to_config = {
- "van-tiny": ImageNetPreTrainedConfig(
- hidden_sizes=[32, 64, 160, 256],
- depths=[3, 3, 5, 2],
- mlp_ratios=[8, 8, 4, 4],
- ),
- "van-small": ImageNetPreTrainedConfig(
- hidden_sizes=[64, 128, 320, 512],
- depths=[2, 2, 4, 2],
- mlp_ratios=[8, 8, 4, 4],
- ),
- "van-base": ImageNetPreTrainedConfig(
- hidden_sizes=[64, 128, 320, 512],
- depths=[3, 3, 12, 3],
- mlp_ratios=[8, 8, 4, 4],
- ),
- "van-large": ImageNetPreTrainedConfig(
- hidden_sizes=[64, 128, 320, 512],
- depths=[3, 5, 27, 3],
- mlp_ratios=[8, 8, 4, 4],
- ),
- }
-
- names_to_original_models = {
- "van-tiny": van_tiny,
- "van-small": van_small,
- "van-base": van_base,
- "van-large": van_large,
- }
-
- names_to_original_checkpoints = {
- "van-tiny": (
- "https://huggingface.co/Visual-Attention-Network/VAN-Tiny-original/resolve/main/van_tiny_754.pth.tar"
- ),
- "van-small": (
- "https://huggingface.co/Visual-Attention-Network/VAN-Small-original/resolve/main/van_small_811.pth.tar"
- ),
- "van-base": (
- "https://huggingface.co/Visual-Attention-Network/VAN-Base-original/resolve/main/van_base_828.pth.tar"
- ),
- "van-large": (
- "https://huggingface.co/Visual-Attention-Network/VAN-Large-original/resolve/main/van_large_839.pth.tar"
- ),
- }
-
- if model_name:
- convert_weight_and_push(
- model_name,
- names_to_config[model_name],
- checkpoint=names_to_original_checkpoints[model_name],
- from_model=names_to_original_models[model_name](),
- save_directory=save_directory,
- push_to_hub=push_to_hub,
- )
- else:
- for model_name, config in names_to_config.items():
- convert_weight_and_push(
- model_name,
- config,
- checkpoint=names_to_original_checkpoints[model_name],
- from_model=names_to_original_models[model_name](),
- save_directory=save_directory,
- push_to_hub=push_to_hub,
- )
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model-name",
- default=None,
- type=str,
- help=(
- "The name of the model you wish to convert, it must be one of the supported resnet* architecture,"
- " currently: van-tiny/small/base/large. If `None`, all of them will the converted."
- ),
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=Path,
- required=True,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--van_dir",
- required=True,
- type=Path,
- help=(
- "A path to VAN's original implementation directory. You can download from here:"
- " https://github.com/Visual-Attention-Network/VAN-Classification"
- ),
- )
- parser.add_argument(
- "--push_to_hub",
- default=True,
- type=bool,
- required=False,
- help="If True, push model and image processor to the hub.",
- )
-
- args = parser.parse_args()
- pytorch_dump_folder_path: Path = args.pytorch_dump_folder_path
- pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
- van_dir = args.van_dir
- # append the path to the parents to maskformer dir
- sys.path.append(str(van_dir.parent))
- from van.models.van import van_base, van_large, van_small, van_tiny
-
- convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)
diff --git a/transformers/models/deprecated/van/modeling_van.py b/transformers/models/deprecated/van/modeling_van.py
deleted file mode 100644
index 6fa2b73482e358d38efc2c08be55a15112f71671..0000000000000000000000000000000000000000
--- a/transformers/models/deprecated/van/modeling_van.py
+++ /dev/null
@@ -1,541 +0,0 @@
-# coding=utf-8
-# Copyright 2022 BNRist (Tsinghua University), TKLNDST (Nankai University) and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Visual Attention Network (VAN) model."""
-
-import math
-from collections import OrderedDict
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ....activations import ACT2FN
-from ....modeling_outputs import (
- BaseModelOutputWithNoAttention,
- BaseModelOutputWithPoolingAndNoAttention,
- ImageClassifierOutputWithNoAttention,
-)
-from ....modeling_utils import PreTrainedModel
-from ....utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_van import VanConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "VanConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "Visual-Attention-Network/van-base"
-_EXPECTED_OUTPUT_SHAPE = [1, 512, 7, 7]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "Visual-Attention-Network/van-base"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from .._archive_maps import VAN_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.convnext.modeling_convnext.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.convnext.modeling_convnext.ConvNextDropPath with ConvNext->Van
-class VanDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-class VanOverlappingPatchEmbedder(nn.Module):
- """
- Downsamples the input using a patchify operation with a `stride` of 4 by default making adjacent windows overlap by
- half of the area. From [PVTv2: Improved Baselines with Pyramid Vision
- Transformer](https://arxiv.org/abs/2106.13797).
- """
-
- def __init__(self, in_channels: int, hidden_size: int, patch_size: int = 7, stride: int = 4):
- super().__init__()
- self.convolution = nn.Conv2d(
- in_channels, hidden_size, kernel_size=patch_size, stride=stride, padding=patch_size // 2
- )
- self.normalization = nn.BatchNorm2d(hidden_size)
-
- def forward(self, input: torch.Tensor) -> torch.Tensor:
- hidden_state = self.convolution(input)
- hidden_state = self.normalization(hidden_state)
- return hidden_state
-
-
-class VanMlpLayer(nn.Module):
- """
- MLP with depth-wise convolution, from [PVTv2: Improved Baselines with Pyramid Vision
- Transformer](https://arxiv.org/abs/2106.13797).
- """
-
- def __init__(
- self,
- in_channels: int,
- hidden_size: int,
- out_channels: int,
- hidden_act: str = "gelu",
- dropout_rate: float = 0.5,
- ):
- super().__init__()
- self.in_dense = nn.Conv2d(in_channels, hidden_size, kernel_size=1)
- self.depth_wise = nn.Conv2d(hidden_size, hidden_size, kernel_size=3, padding=1, groups=hidden_size)
- self.activation = ACT2FN[hidden_act]
- self.dropout1 = nn.Dropout(dropout_rate)
- self.out_dense = nn.Conv2d(hidden_size, out_channels, kernel_size=1)
- self.dropout2 = nn.Dropout(dropout_rate)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- hidden_state = self.in_dense(hidden_state)
- hidden_state = self.depth_wise(hidden_state)
- hidden_state = self.activation(hidden_state)
- hidden_state = self.dropout1(hidden_state)
- hidden_state = self.out_dense(hidden_state)
- hidden_state = self.dropout2(hidden_state)
- return hidden_state
-
-
-class VanLargeKernelAttention(nn.Module):
- """
- Basic Large Kernel Attention (LKA).
- """
-
- def __init__(self, hidden_size: int):
- super().__init__()
- self.depth_wise = nn.Conv2d(hidden_size, hidden_size, kernel_size=5, padding=2, groups=hidden_size)
- self.depth_wise_dilated = nn.Conv2d(
- hidden_size, hidden_size, kernel_size=7, dilation=3, padding=9, groups=hidden_size
- )
- self.point_wise = nn.Conv2d(hidden_size, hidden_size, kernel_size=1)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- hidden_state = self.depth_wise(hidden_state)
- hidden_state = self.depth_wise_dilated(hidden_state)
- hidden_state = self.point_wise(hidden_state)
- return hidden_state
-
-
-class VanLargeKernelAttentionLayer(nn.Module):
- """
- Computes attention using Large Kernel Attention (LKA) and attends the input.
- """
-
- def __init__(self, hidden_size: int):
- super().__init__()
- self.attention = VanLargeKernelAttention(hidden_size)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- attention = self.attention(hidden_state)
- attended = hidden_state * attention
- return attended
-
-
-class VanSpatialAttentionLayer(nn.Module):
- """
- Van spatial attention layer composed by projection (via conv) -> act -> Large Kernel Attention (LKA) attention ->
- projection (via conv) + residual connection.
- """
-
- def __init__(self, hidden_size: int, hidden_act: str = "gelu"):
- super().__init__()
- self.pre_projection = nn.Sequential(
- OrderedDict(
- [
- ("conv", nn.Conv2d(hidden_size, hidden_size, kernel_size=1)),
- ("act", ACT2FN[hidden_act]),
- ]
- )
- )
- self.attention_layer = VanLargeKernelAttentionLayer(hidden_size)
- self.post_projection = nn.Conv2d(hidden_size, hidden_size, kernel_size=1)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- residual = hidden_state
- hidden_state = self.pre_projection(hidden_state)
- hidden_state = self.attention_layer(hidden_state)
- hidden_state = self.post_projection(hidden_state)
- hidden_state = hidden_state + residual
- return hidden_state
-
-
-class VanLayerScaling(nn.Module):
- """
- Scales the inputs by a learnable parameter initialized by `initial_value`.
- """
-
- def __init__(self, hidden_size: int, initial_value: float = 1e-2):
- super().__init__()
- self.weight = nn.Parameter(initial_value * torch.ones((hidden_size)), requires_grad=True)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- # unsqueezing for broadcasting
- hidden_state = self.weight.unsqueeze(-1).unsqueeze(-1) * hidden_state
- return hidden_state
-
-
-class VanLayer(nn.Module):
- """
- Van layer composed by normalization layers, large kernel attention (LKA) and a multi layer perceptron (MLP).
- """
-
- def __init__(
- self,
- config: VanConfig,
- hidden_size: int,
- mlp_ratio: int = 4,
- drop_path_rate: float = 0.5,
- ):
- super().__init__()
- self.drop_path = VanDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
- self.pre_normomalization = nn.BatchNorm2d(hidden_size)
- self.attention = VanSpatialAttentionLayer(hidden_size, config.hidden_act)
- self.attention_scaling = VanLayerScaling(hidden_size, config.layer_scale_init_value)
- self.post_normalization = nn.BatchNorm2d(hidden_size)
- self.mlp = VanMlpLayer(
- hidden_size, hidden_size * mlp_ratio, hidden_size, config.hidden_act, config.dropout_rate
- )
- self.mlp_scaling = VanLayerScaling(hidden_size, config.layer_scale_init_value)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- residual = hidden_state
- # attention
- hidden_state = self.pre_normomalization(hidden_state)
- hidden_state = self.attention(hidden_state)
- hidden_state = self.attention_scaling(hidden_state)
- hidden_state = self.drop_path(hidden_state)
- # residual connection
- hidden_state = residual + hidden_state
- residual = hidden_state
- # mlp
- hidden_state = self.post_normalization(hidden_state)
- hidden_state = self.mlp(hidden_state)
- hidden_state = self.mlp_scaling(hidden_state)
- hidden_state = self.drop_path(hidden_state)
- # residual connection
- hidden_state = residual + hidden_state
- return hidden_state
-
-
-class VanStage(nn.Module):
- """
- VanStage, consisting of multiple layers.
- """
-
- def __init__(
- self,
- config: VanConfig,
- in_channels: int,
- hidden_size: int,
- patch_size: int,
- stride: int,
- depth: int,
- mlp_ratio: int = 4,
- drop_path_rate: float = 0.0,
- ):
- super().__init__()
- self.embeddings = VanOverlappingPatchEmbedder(in_channels, hidden_size, patch_size, stride)
- self.layers = nn.Sequential(
- *[
- VanLayer(
- config,
- hidden_size,
- mlp_ratio=mlp_ratio,
- drop_path_rate=drop_path_rate,
- )
- for _ in range(depth)
- ]
- )
- self.normalization = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- hidden_state = self.embeddings(hidden_state)
- hidden_state = self.layers(hidden_state)
- # rearrange b c h w -> b (h w) c
- batch_size, hidden_size, height, width = hidden_state.shape
- hidden_state = hidden_state.flatten(2).transpose(1, 2)
- hidden_state = self.normalization(hidden_state)
- # rearrange b (h w) c- > b c h w
- hidden_state = hidden_state.view(batch_size, height, width, hidden_size).permute(0, 3, 1, 2)
- return hidden_state
-
-
-class VanEncoder(nn.Module):
- """
- VanEncoder, consisting of multiple stages.
- """
-
- def __init__(self, config: VanConfig):
- super().__init__()
- self.stages = nn.ModuleList([])
- patch_sizes = config.patch_sizes
- strides = config.strides
- hidden_sizes = config.hidden_sizes
- depths = config.depths
- mlp_ratios = config.mlp_ratios
- drop_path_rates = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
-
- for num_stage, (patch_size, stride, hidden_size, depth, mlp_expantion, drop_path_rate) in enumerate(
- zip(patch_sizes, strides, hidden_sizes, depths, mlp_ratios, drop_path_rates)
- ):
- is_first_stage = num_stage == 0
- in_channels = hidden_sizes[num_stage - 1]
- if is_first_stage:
- in_channels = config.num_channels
- self.stages.append(
- VanStage(
- config,
- in_channels,
- hidden_size,
- patch_size=patch_size,
- stride=stride,
- depth=depth,
- mlp_ratio=mlp_expantion,
- drop_path_rate=drop_path_rate,
- )
- )
-
- def forward(
- self,
- hidden_state: torch.Tensor,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple, BaseModelOutputWithNoAttention]:
- all_hidden_states = () if output_hidden_states else None
-
- for _, stage_module in enumerate(self.stages):
- hidden_state = stage_module(hidden_state)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_state,)
-
- if not return_dict:
- return tuple(v for v in [hidden_state, all_hidden_states] if v is not None)
-
- return BaseModelOutputWithNoAttention(last_hidden_state=hidden_state, hidden_states=all_hidden_states)
-
-
-class VanPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = VanConfig
- base_model_prefix = "van"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- nn.init.trunc_normal_(module.weight, std=self.config.initializer_range)
- if isinstance(module, nn.Linear) and module.bias is not None:
- nn.init.constant_(module.bias, 0)
- elif isinstance(module, nn.LayerNorm):
- nn.init.constant_(module.bias, 0)
- nn.init.constant_(module.weight, 1.0)
- elif isinstance(module, nn.Conv2d):
- fan_out = module.kernel_size[0] * module.kernel_size[1] * module.out_channels
- fan_out //= module.groups
- module.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
- if module.bias is not None:
- module.bias.data.zero_()
-
-
-VAN_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`VanConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-VAN_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`ConvNextImageProcessor.__call__`] for details.
-
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all stages. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare VAN model outputting raw features without any specific head on top. Note, VAN does not have an embedding"
- " layer.",
- VAN_START_DOCSTRING,
-)
-class VanModel(VanPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.config = config
- self.encoder = VanEncoder(config)
- # final layernorm layer
- self.layernorm = nn.LayerNorm(config.hidden_sizes[-1], eps=config.layer_norm_eps)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(VAN_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPoolingAndNoAttention,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor],
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPoolingAndNoAttention]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- encoder_outputs = self.encoder(
- pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- last_hidden_state = encoder_outputs[0]
- # global average pooling, n c w h -> n c
- pooled_output = last_hidden_state.mean(dim=[-2, -1])
-
- if not return_dict:
- return (last_hidden_state, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPoolingAndNoAttention(
- last_hidden_state=last_hidden_state,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- VAN Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
- ImageNet.
- """,
- VAN_START_DOCSTRING,
-)
-class VanForImageClassification(VanPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.van = VanModel(config)
- # Classifier head
- self.classifier = (
- nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(VAN_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=ImageClassifierOutputWithNoAttention,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.van(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
-
- pooled_output = outputs.pooler_output if return_dict else outputs[1]
-
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.config.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.config.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
diff --git a/transformers/models/depth_anything/__init__.py b/transformers/models/depth_anything/__init__.py
deleted file mode 100644
index 0d0ea5a514a836646026eeddf232b92c702f567c..0000000000000000000000000000000000000000
--- a/transformers/models/depth_anything/__init__.py
+++ /dev/null
@@ -1,56 +0,0 @@
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...file_utils import _LazyModule, is_torch_available
-from ...utils import OptionalDependencyNotAvailable
-
-
-_import_structure = {
- "configuration_depth_anything": ["DEPTH_ANYTHING_PRETRAINED_CONFIG_ARCHIVE_MAP", "DepthAnythingConfig"]
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_depth_anything"] = [
- "DEPTH_ANYTHING_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DepthAnythingForDepthEstimation",
- "DepthAnythingPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_depth_anything import DEPTH_ANYTHING_PRETRAINED_CONFIG_ARCHIVE_MAP, DepthAnythingConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_depth_anything import (
- DEPTH_ANYTHING_PRETRAINED_MODEL_ARCHIVE_LIST,
- DepthAnythingForDepthEstimation,
- DepthAnythingPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/depth_anything/__pycache__/__init__.cpython-310.pyc b/transformers/models/depth_anything/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 7372e8e88ba57a0cf2158ef6fd7ca1b257a78852..0000000000000000000000000000000000000000
Binary files a/transformers/models/depth_anything/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/depth_anything/__pycache__/configuration_depth_anything.cpython-310.pyc b/transformers/models/depth_anything/__pycache__/configuration_depth_anything.cpython-310.pyc
deleted file mode 100644
index 985f120fb0cbe7d492e873164972eb3604407088..0000000000000000000000000000000000000000
Binary files a/transformers/models/depth_anything/__pycache__/configuration_depth_anything.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/depth_anything/__pycache__/convert_depth_anything_to_hf.cpython-310.pyc b/transformers/models/depth_anything/__pycache__/convert_depth_anything_to_hf.cpython-310.pyc
deleted file mode 100644
index 47202003d09b04f6b965c27d345a7751b44386f2..0000000000000000000000000000000000000000
Binary files a/transformers/models/depth_anything/__pycache__/convert_depth_anything_to_hf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/depth_anything/__pycache__/modeling_depth_anything.cpython-310.pyc b/transformers/models/depth_anything/__pycache__/modeling_depth_anything.cpython-310.pyc
deleted file mode 100644
index 246b0c179e962bf9cfc4903f76675ec0346b6935..0000000000000000000000000000000000000000
Binary files a/transformers/models/depth_anything/__pycache__/modeling_depth_anything.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/depth_anything/configuration_depth_anything.py b/transformers/models/depth_anything/configuration_depth_anything.py
deleted file mode 100644
index 3d58a3874eedf3603819ecb96ed47fc9c12086fd..0000000000000000000000000000000000000000
--- a/transformers/models/depth_anything/configuration_depth_anything.py
+++ /dev/null
@@ -1,145 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DepthAnything model configuration"""
-
-import copy
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-from ..auto.configuration_auto import CONFIG_MAPPING
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DEPTH_ANYTHING_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DepthAnythingConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DepthAnythingModel`]. It is used to instantiate an DepthAnything
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the DepthAnything
- [LiheYoung/depth-anything-small-hf](https://huggingface.co/LiheYoung/depth-anything-small-hf) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- backbone_config (`Union[Dict[str, Any], PretrainedConfig]`, *optional*):
- The configuration of the backbone model. Only used in case `is_hybrid` is `True` or in case you want to
- leverage the [`AutoBackbone`] API.
- backbone (`str`, *optional*):
- Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
- will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
- is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
- use_pretrained_backbone (`bool`, *optional*, defaults to `False`):
- Whether to use pretrained weights for the backbone.
- patch_size (`int`, *optional*, defaults to 14):
- The size of the patches to extract from the backbone features.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- reassemble_hidden_size (`int`, *optional*, defaults to 384):
- The number of input channels of the reassemble layers.
- reassemble_factors (`List[int]`, *optional*, defaults to `[4, 2, 1, 0.5]`):
- The up/downsampling factors of the reassemble layers.
- neck_hidden_sizes (`List[str]`, *optional*, defaults to `[48, 96, 192, 384]`):
- The hidden sizes to project to for the feature maps of the backbone.
- fusion_hidden_size (`int`, *optional*, defaults to 64):
- The number of channels before fusion.
- head_in_index (`int`, *optional*, defaults to -1):
- The index of the features to use in the depth estimation head.
- head_hidden_size (`int`, *optional*, defaults to 32):
- The number of output channels in the second convolution of the depth estimation head.
-
- Example:
-
- ```python
- >>> from transformers import DepthAnythingConfig, DepthAnythingForDepthEstimation
-
- >>> # Initializing a DepthAnything small style configuration
- >>> configuration = DepthAnythingConfig()
-
- >>> # Initializing a model from the DepthAnything small style configuration
- >>> model = DepthAnythingForDepthEstimation(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "depth_anything"
-
- def __init__(
- self,
- backbone_config=None,
- backbone=None,
- use_pretrained_backbone=False,
- patch_size=14,
- initializer_range=0.02,
- reassemble_hidden_size=384,
- reassemble_factors=[4, 2, 1, 0.5],
- neck_hidden_sizes=[48, 96, 192, 384],
- fusion_hidden_size=64,
- head_in_index=-1,
- head_hidden_size=32,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- if use_pretrained_backbone:
- raise ValueError("Pretrained backbones are not supported yet.")
-
- if backbone_config is not None and backbone is not None:
- raise ValueError("You can't specify both `backbone` and `backbone_config`.")
-
- if backbone_config is None and backbone is None:
- logger.info("`backbone_config` is `None`. Initializing the config with the default `Dinov2` backbone.")
- backbone_config = CONFIG_MAPPING["dinov2"](
- image_size=518,
- hidden_size=384,
- num_attention_heads=6,
- out_indices=[9, 10, 11, 12],
- apply_layernorm=True,
- reshape_hidden_states=False,
- )
- elif isinstance(backbone_config, dict):
- backbone_model_type = backbone_config.get("model_type")
- config_class = CONFIG_MAPPING[backbone_model_type]
- backbone_config = config_class.from_dict(backbone_config)
-
- self.backbone_config = backbone_config
- self.backbone = backbone
- self.use_pretrained_backbone = use_pretrained_backbone
- self.reassemble_hidden_size = reassemble_hidden_size
- self.patch_size = patch_size
- self.initializer_range = initializer_range
- self.reassemble_factors = reassemble_factors
- self.neck_hidden_sizes = neck_hidden_sizes
- self.fusion_hidden_size = fusion_hidden_size
- self.head_in_index = head_in_index
- self.head_hidden_size = head_hidden_size
-
- def to_dict(self):
- """
- Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. Returns:
- `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
- """
- output = copy.deepcopy(self.__dict__)
-
- if output["backbone_config"] is not None:
- output["backbone_config"] = self.backbone_config.to_dict()
-
- output["model_type"] = self.__class__.model_type
- return output
diff --git a/transformers/models/depth_anything/convert_depth_anything_to_hf.py b/transformers/models/depth_anything/convert_depth_anything_to_hf.py
deleted file mode 100644
index 022a66c0d609cd2ae2d9893cda93c2e52b6996f3..0000000000000000000000000000000000000000
--- a/transformers/models/depth_anything/convert_depth_anything_to_hf.py
+++ /dev/null
@@ -1,299 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Depth Anything checkpoints from the original repository. URL:
-https://github.com/LiheYoung/Depth-Anything"""
-
-
-import argparse
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from transformers import DepthAnythingConfig, DepthAnythingForDepthEstimation, Dinov2Config, DPTImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_dpt_config(model_name):
- if "small" in model_name:
- backbone_config = Dinov2Config.from_pretrained(
- "facebook/dinov2-small", out_indices=[9, 10, 11, 12], apply_layernorm=True, reshape_hidden_states=False
- )
- fusion_hidden_size = 64
- neck_hidden_sizes = [48, 96, 192, 384]
- elif "base" in model_name:
- backbone_config = Dinov2Config.from_pretrained(
- "facebook/dinov2-base", out_indices=[9, 10, 11, 12], apply_layernorm=True, reshape_hidden_states=False
- )
- fusion_hidden_size = 128
- neck_hidden_sizes = [96, 192, 384, 768]
- elif "large" in model_name:
- backbone_config = Dinov2Config.from_pretrained(
- "facebook/dinov2-large", out_indices=[21, 22, 23, 24], apply_layernorm=True, reshape_hidden_states=False
- )
- fusion_hidden_size = 256
- neck_hidden_sizes = [256, 512, 1024, 1024]
- else:
- raise NotImplementedError("To do")
-
- config = DepthAnythingConfig(
- reassemble_hidden_size=backbone_config.hidden_size,
- patch_size=backbone_config.patch_size,
- backbone_config=backbone_config,
- fusion_hidden_size=fusion_hidden_size,
- neck_hidden_sizes=neck_hidden_sizes,
- )
-
- return config
-
-
-def create_rename_keys(config):
- rename_keys = []
-
- # fmt: off
- # stem
- rename_keys.append(("pretrained.cls_token", "backbone.embeddings.cls_token"))
- rename_keys.append(("pretrained.mask_token", "backbone.embeddings.mask_token"))
- rename_keys.append(("pretrained.pos_embed", "backbone.embeddings.position_embeddings"))
- rename_keys.append(("pretrained.patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight"))
- rename_keys.append(("pretrained.patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias"))
-
- # Transfomer encoder
- for i in range(config.backbone_config.num_hidden_layers):
- rename_keys.append((f"pretrained.blocks.{i}.ls1.gamma", f"backbone.encoder.layer.{i}.layer_scale1.lambda1"))
- rename_keys.append((f"pretrained.blocks.{i}.ls2.gamma", f"backbone.encoder.layer.{i}.layer_scale2.lambda1"))
- rename_keys.append((f"pretrained.blocks.{i}.norm1.weight", f"backbone.encoder.layer.{i}.norm1.weight"))
- rename_keys.append((f"pretrained.blocks.{i}.norm1.bias", f"backbone.encoder.layer.{i}.norm1.bias"))
- rename_keys.append((f"pretrained.blocks.{i}.norm2.weight", f"backbone.encoder.layer.{i}.norm2.weight"))
- rename_keys.append((f"pretrained.blocks.{i}.norm2.bias", f"backbone.encoder.layer.{i}.norm2.bias"))
- rename_keys.append((f"pretrained.blocks.{i}.mlp.fc1.weight", f"backbone.encoder.layer.{i}.mlp.fc1.weight"))
- rename_keys.append((f"pretrained.blocks.{i}.mlp.fc1.bias", f"backbone.encoder.layer.{i}.mlp.fc1.bias"))
- rename_keys.append((f"pretrained.blocks.{i}.mlp.fc2.weight", f"backbone.encoder.layer.{i}.mlp.fc2.weight"))
- rename_keys.append((f"pretrained.blocks.{i}.mlp.fc2.bias", f"backbone.encoder.layer.{i}.mlp.fc2.bias"))
- rename_keys.append((f"pretrained.blocks.{i}.attn.proj.weight", f"backbone.encoder.layer.{i}.attention.output.dense.weight"))
- rename_keys.append((f"pretrained.blocks.{i}.attn.proj.bias", f"backbone.encoder.layer.{i}.attention.output.dense.bias"))
-
- # Head
- rename_keys.append(("pretrained.norm.weight", "backbone.layernorm.weight"))
- rename_keys.append(("pretrained.norm.bias", "backbone.layernorm.bias"))
-
- # activation postprocessing (readout projections + resize blocks)
- # Depth Anything does not use CLS token => readout_projects not required
-
- for i in range(4):
- rename_keys.append((f"depth_head.projects.{i}.weight", f"neck.reassemble_stage.layers.{i}.projection.weight"))
- rename_keys.append((f"depth_head.projects.{i}.bias", f"neck.reassemble_stage.layers.{i}.projection.bias"))
-
- if i != 2:
- rename_keys.append((f"depth_head.resize_layers.{i}.weight", f"neck.reassemble_stage.layers.{i}.resize.weight"))
- rename_keys.append((f"depth_head.resize_layers.{i}.bias", f"neck.reassemble_stage.layers.{i}.resize.bias"))
-
- # refinenet (tricky here)
- mapping = {1:3, 2:2, 3:1, 4:0}
-
- for i in range(1, 5):
- j = mapping[i]
- rename_keys.append((f"depth_head.scratch.refinenet{i}.out_conv.weight", f"neck.fusion_stage.layers.{j}.projection.weight"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.out_conv.bias", f"neck.fusion_stage.layers.{j}.projection.bias"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit1.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.weight"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit1.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.bias"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit1.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.weight"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit1.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.bias"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit2.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.weight"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit2.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.bias"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit2.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.weight"))
- rename_keys.append((f"depth_head.scratch.refinenet{i}.resConfUnit2.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.bias"))
-
- # scratch convolutions
- for i in range(4):
- rename_keys.append((f"depth_head.scratch.layer{i+1}_rn.weight", f"neck.convs.{i}.weight"))
-
- # head
- rename_keys.append(("depth_head.scratch.output_conv1.weight", "head.conv1.weight"))
- rename_keys.append(("depth_head.scratch.output_conv1.bias", "head.conv1.bias"))
- rename_keys.append(("depth_head.scratch.output_conv2.0.weight", "head.conv2.weight"))
- rename_keys.append(("depth_head.scratch.output_conv2.0.bias", "head.conv2.bias"))
- rename_keys.append(("depth_head.scratch.output_conv2.2.weight", "head.conv3.weight"))
- rename_keys.append(("depth_head.scratch.output_conv2.2.bias", "head.conv3.bias"))
-
- return rename_keys
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config):
- hidden_size = config.backbone_config.hidden_size
- for i in range(config.backbone_config.num_hidden_layers):
- # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"pretrained.blocks.{i}.attn.qkv.weight")
- in_proj_bias = state_dict.pop(f"pretrained.blocks.{i}.attn.qkv.bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[:hidden_size, :]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[:hidden_size]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- hidden_size : hidden_size * 2, :
- ]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
- hidden_size : hidden_size * 2
- ]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-hidden_size:, :]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-hidden_size:]
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-name_to_checkpoint = {
- "depth-anything-small": "depth_anything_vits14.pth",
- "depth-anything-base": "depth_anything_vitb14.pth",
- "depth-anything-large": "depth_anything_vitl14.pth",
-}
-
-
-@torch.no_grad()
-def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub, verify_logits):
- """
- Copy/paste/tweak model's weights to our DPT structure.
- """
-
- # define DPT configuration
- config = get_dpt_config(model_name)
-
- model_name_to_filename = {
- "depth-anything-small": "depth_anything_vits14.pth",
- "depth-anything-base": "depth_anything_vitb14.pth",
- "depth-anything-large": "depth_anything_vitl14.pth",
- }
-
- # load original state_dict
- filename = model_name_to_filename[model_name]
- filepath = hf_hub_download(
- repo_id="LiheYoung/Depth-Anything", filename=f"checkpoints/{filename}", repo_type="space"
- )
- state_dict = torch.load(filepath, map_location="cpu")
- # rename keys
- rename_keys = create_rename_keys(config)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- # read in qkv matrices
- read_in_q_k_v(state_dict, config)
-
- # load HuggingFace model
- model = DepthAnythingForDepthEstimation(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- processor = DPTImageProcessor(
- do_resize=True,
- size={"height": 518, "width": 518},
- ensure_multiple_of=14,
- keep_aspect_ratio=True,
- do_rescale=True,
- do_normalize=True,
- image_mean=[0.485, 0.456, 0.406],
- image_std=[0.229, 0.224, 0.225],
- )
-
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- image = Image.open(requests.get(url, stream=True).raw)
-
- pixel_values = processor(image, return_tensors="pt").pixel_values
-
- # Verify forward pass
- with torch.no_grad():
- outputs = model(pixel_values)
- predicted_depth = outputs.predicted_depth
-
- print("Shape of predicted depth:", predicted_depth.shape)
- print("First values:", predicted_depth[0, :3, :3])
-
- # assert logits
- if verify_logits:
- expected_shape = torch.Size([1, 518, 686])
- if model_name == "depth-anything-small":
- expected_slice = torch.tensor(
- [[8.8204, 8.6468, 8.6195], [8.3313, 8.6027, 8.7526], [8.6526, 8.6866, 8.7453]],
- )
- elif model_name == "depth-anything-base":
- expected_slice = torch.tensor(
- [[26.3997, 26.3004, 26.3928], [26.2260, 26.2092, 26.3427], [26.0719, 26.0483, 26.1254]],
- )
- elif model_name == "depth-anything-large":
- expected_slice = torch.tensor(
- [[87.9968, 87.7493, 88.2704], [87.1927, 87.6611, 87.3640], [86.7789, 86.9469, 86.7991]]
- )
- else:
- raise ValueError("Not supported")
-
- assert predicted_depth.shape == torch.Size(expected_shape)
- assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-6)
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model and processor to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- print("Pushing model and processor to hub...")
- model.push_to_hub(repo_id=f"LiheYoung/{model_name}-hf")
- processor.push_to_hub(repo_id=f"LiheYoung/{model_name}-hf")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="depth-anything-small",
- type=str,
- choices=name_to_checkpoint.keys(),
- help="Name of the model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- help="Whether to push the model to the hub after conversion.",
- )
- parser.add_argument(
- "--verify_logits",
- action="store_false",
- required=False,
- help="Whether to verify the logits after conversion.",
- )
-
- args = parser.parse_args()
- convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.verify_logits)
diff --git a/transformers/models/depth_anything/modeling_depth_anything.py b/transformers/models/depth_anything/modeling_depth_anything.py
deleted file mode 100644
index 788b0d911396f19be6a6e54b39dd52139ba937da..0000000000000000000000000000000000000000
--- a/transformers/models/depth_anything/modeling_depth_anything.py
+++ /dev/null
@@ -1,463 +0,0 @@
-# coding=utf-8
-# Copyright 2024 TikTok and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Depth Anything model."""
-
-
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from ...file_utils import (
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- replace_return_docstrings,
-)
-from ...modeling_outputs import DepthEstimatorOutput
-from ...modeling_utils import PreTrainedModel
-from ...utils import logging
-from ..auto import AutoBackbone
-from .configuration_depth_anything import DepthAnythingConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "DepthAnythingConfig"
-
-
-from ..deprecated._archive_maps import DEPTH_ANYTHING_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-DEPTH_ANYTHING_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`DepthAnythingConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DEPTH_ANYTHING_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`DPTImageProcessor.__call__`]
- for details.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class DepthAnythingReassembleLayer(nn.Module):
- def __init__(self, config, channels, factor):
- super().__init__()
- self.projection = nn.Conv2d(in_channels=config.reassemble_hidden_size, out_channels=channels, kernel_size=1)
-
- # up/down sampling depending on factor
- if factor > 1:
- self.resize = nn.ConvTranspose2d(channels, channels, kernel_size=factor, stride=factor, padding=0)
- elif factor == 1:
- self.resize = nn.Identity()
- elif factor < 1:
- # so should downsample
- self.resize = nn.Conv2d(channels, channels, kernel_size=3, stride=int(1 / factor), padding=1)
-
- # Copied from transformers.models.dpt.modeling_dpt.DPTReassembleLayer.forward
- def forward(self, hidden_state):
- hidden_state = self.projection(hidden_state)
- hidden_state = self.resize(hidden_state)
-
- return hidden_state
-
-
-class DepthAnythingReassembleStage(nn.Module):
- """
- This class reassembles the hidden states of the backbone into image-like feature representations at various
- resolutions.
-
- This happens in 3 stages:
- 1. Take the patch embeddings and reshape them to image-like feature representations.
- 2. Project the channel dimension of the hidden states according to `config.neck_hidden_sizes`.
- 3. Resizing the spatial dimensions (height, width).
-
- Args:
- config (`[DepthAnythingConfig]`):
- Model configuration class defining the model architecture.
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.config = config
- self.layers = nn.ModuleList()
- for channels, factor in zip(config.neck_hidden_sizes, config.reassemble_factors):
- self.layers.append(DepthAnythingReassembleLayer(config, channels=channels, factor=factor))
-
- def forward(self, hidden_states: List[torch.Tensor], patch_height=None, patch_width=None) -> List[torch.Tensor]:
- """
- Args:
- hidden_states (`List[torch.FloatTensor]`, each of shape `(batch_size, sequence_length + 1, hidden_size)`):
- List of hidden states from the backbone.
- """
- out = []
-
- for i, hidden_state in enumerate(hidden_states):
- # reshape to (batch_size, num_channels, height, width)
- hidden_state = hidden_state[:, 1:]
- batch_size, _, num_channels = hidden_state.shape
- hidden_state = hidden_state.reshape(batch_size, patch_height, patch_width, num_channels)
- hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
- hidden_state = self.layers[i](hidden_state)
- out.append(hidden_state)
-
- return out
-
-
-class DepthAnythingPreActResidualLayer(nn.Module):
- """
- ResidualConvUnit, pre-activate residual unit.
-
- Args:
- config (`[DepthAnythingConfig]`):
- Model configuration class defining the model architecture.
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.activation1 = nn.ReLU()
- self.convolution1 = nn.Conv2d(
- config.fusion_hidden_size,
- config.fusion_hidden_size,
- kernel_size=3,
- stride=1,
- padding=1,
- bias=True,
- )
-
- self.activation2 = nn.ReLU()
- self.convolution2 = nn.Conv2d(
- config.fusion_hidden_size,
- config.fusion_hidden_size,
- kernel_size=3,
- stride=1,
- padding=1,
- bias=True,
- )
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- residual = hidden_state
- hidden_state = self.activation1(hidden_state)
- hidden_state = self.convolution1(hidden_state)
- hidden_state = self.activation2(hidden_state)
- hidden_state = self.convolution2(hidden_state)
-
- return hidden_state + residual
-
-
-class DepthAnythingFeatureFusionLayer(nn.Module):
- """Feature fusion layer, merges feature maps from different stages.
-
- Args:
- config (`[DepthAnythingConfig]`):
- Model configuration class defining the model architecture.
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.projection = nn.Conv2d(config.fusion_hidden_size, config.fusion_hidden_size, kernel_size=1, bias=True)
-
- self.residual_layer1 = DepthAnythingPreActResidualLayer(config)
- self.residual_layer2 = DepthAnythingPreActResidualLayer(config)
-
- def forward(self, hidden_state, residual=None, size=None):
- if residual is not None:
- if hidden_state.shape != residual.shape:
- residual = nn.functional.interpolate(
- residual, size=(hidden_state.shape[2], hidden_state.shape[3]), mode="bilinear", align_corners=False
- )
- hidden_state = hidden_state + self.residual_layer1(residual)
-
- hidden_state = self.residual_layer2(hidden_state)
-
- modifier = {"scale_factor": 2} if size is None else {"size": size}
-
- hidden_state = nn.functional.interpolate(
- hidden_state,
- **modifier,
- mode="bilinear",
- align_corners=True,
- )
- hidden_state = self.projection(hidden_state)
-
- return hidden_state
-
-
-class DepthAnythingFeatureFusionStage(nn.Module):
- # Copied from transformers.models.dpt.modeling_dpt.DPTFeatureFusionStage.__init__ with DPT->DepthAnything
- def __init__(self, config):
- super().__init__()
- self.layers = nn.ModuleList()
- for _ in range(len(config.neck_hidden_sizes)):
- self.layers.append(DepthAnythingFeatureFusionLayer(config))
-
- def forward(self, hidden_states, size=None):
- # reversing the hidden_states, we start from the last
- hidden_states = hidden_states[::-1]
-
- fused_hidden_states = []
- # first layer only uses the last hidden_state
- size = hidden_states[1].shape[2:]
- fused_hidden_state = self.layers[0](hidden_states[0], size=size)
- fused_hidden_states.append(fused_hidden_state)
-
- # looping from the last layer to the second
- for idx, (hidden_state, layer) in enumerate(zip(hidden_states[1:], self.layers[1:])):
- size = hidden_states[1:][idx + 1].shape[2:] if idx != (len(hidden_states[1:]) - 1) else None
-
- fused_hidden_state = layer(fused_hidden_state, hidden_state, size=size)
-
- fused_hidden_states.append(fused_hidden_state)
-
- return fused_hidden_states
-
-
-# Copied from transformers.models.dpt.modeling_dpt.DPTPreTrainedModel with DPT->DepthAnything,dpt->depth_anything
-class DepthAnythingPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DepthAnythingConfig
- base_model_prefix = "depth_anything"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-class DepthAnythingNeck(nn.Module):
- """
- DepthAnythingNeck. A neck is a module that is normally used between the backbone and the head. It takes a list of tensors as
- input and produces another list of tensors as output. For DepthAnything, it includes 2 stages:
-
- * DepthAnythingReassembleStage
- * DepthAnythingFeatureFusionStage.
-
- Args:
- config (dict): config dict.
- """
-
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- self.reassemble_stage = DepthAnythingReassembleStage(config)
-
- self.convs = nn.ModuleList()
- for channel in config.neck_hidden_sizes:
- self.convs.append(nn.Conv2d(channel, config.fusion_hidden_size, kernel_size=3, padding=1, bias=False))
-
- # fusion
- self.fusion_stage = DepthAnythingFeatureFusionStage(config)
-
- def forward(self, hidden_states: List[torch.Tensor], patch_height=None, patch_width=None) -> List[torch.Tensor]:
- """
- Args:
- hidden_states (`List[torch.FloatTensor]`, each of shape `(batch_size, sequence_length, hidden_size)` or `(batch_size, hidden_size, height, width)`):
- List of hidden states from the backbone.
- """
- if not isinstance(hidden_states, (tuple, list)):
- raise ValueError("hidden_states should be a tuple or list of tensors")
-
- if len(hidden_states) != len(self.config.neck_hidden_sizes):
- raise ValueError("The number of hidden states should be equal to the number of neck hidden sizes.")
-
- # postprocess hidden states
- hidden_states = self.reassemble_stage(hidden_states, patch_height, patch_width)
-
- features = [self.convs[i](feature) for i, feature in enumerate(hidden_states)]
-
- # fusion blocks
- output = self.fusion_stage(features)
-
- return output
-
-
-class DepthAnythingDepthEstimationHead(nn.Module):
- """
- Output head consisting of 3 convolutional layers. It progressively halves the feature dimension and upsamples
- the predictions to the input resolution after the first convolutional layer (details can be found in the DPT paper's
- supplementary material).
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.head_in_index = config.head_in_index
- self.patch_size = config.patch_size
-
- features = config.fusion_hidden_size
- self.conv1 = nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1)
- self.conv2 = nn.Conv2d(features // 2, config.head_hidden_size, kernel_size=3, stride=1, padding=1)
- self.activation1 = nn.ReLU()
- self.conv3 = nn.Conv2d(config.head_hidden_size, 1, kernel_size=1, stride=1, padding=0)
- self.activation2 = nn.ReLU()
-
- def forward(self, hidden_states: List[torch.Tensor], patch_height, patch_width) -> torch.Tensor:
- hidden_states = hidden_states[self.head_in_index]
-
- predicted_depth = self.conv1(hidden_states)
- predicted_depth = nn.functional.interpolate(
- predicted_depth,
- (int(patch_height * self.patch_size), int(patch_width * self.patch_size)),
- mode="bilinear",
- align_corners=True,
- )
- predicted_depth = self.conv2(predicted_depth)
- predicted_depth = self.activation1(predicted_depth)
- predicted_depth = self.conv3(predicted_depth)
- predicted_depth = self.activation2(predicted_depth)
- predicted_depth = predicted_depth.squeeze(dim=1) # shape (batch_size, height, width)
-
- return predicted_depth
-
-
-@add_start_docstrings(
- """
- Depth Anything Model with a depth estimation head on top (consisting of 3 convolutional layers) e.g. for KITTI, NYUv2.
- """,
- DEPTH_ANYTHING_START_DOCSTRING,
-)
-class DepthAnythingForDepthEstimation(DepthAnythingPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.backbone = AutoBackbone.from_config(config.backbone_config)
- self.neck = DepthAnythingNeck(config)
- self.head = DepthAnythingDepthEstimationHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DEPTH_ANYTHING_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DepthEstimatorOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], DepthEstimatorOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
- Ground truth depth estimation maps for computing the loss.
-
- Returns:
-
- Examples:
- ```python
- >>> from transformers import AutoImageProcessor, AutoModelForDepthEstimation
- >>> import torch
- >>> import numpy as np
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-small-hf")
- >>> model = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-small-hf")
-
- >>> # prepare image for the model
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> with torch.no_grad():
- ... outputs = model(**inputs)
- ... predicted_depth = outputs.predicted_depth
-
- >>> # interpolate to original size
- >>> prediction = torch.nn.functional.interpolate(
- ... predicted_depth.unsqueeze(1),
- ... size=image.size[::-1],
- ... mode="bicubic",
- ... align_corners=False,
- ... )
-
- >>> # visualize the prediction
- >>> output = prediction.squeeze().cpu().numpy()
- >>> formatted = (output * 255 / np.max(output)).astype("uint8")
- >>> depth = Image.fromarray(formatted)
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-
- outputs = self.backbone.forward_with_filtered_kwargs(
- pixel_values, output_hidden_states=output_hidden_states, output_attentions=output_attentions
- )
- hidden_states = outputs.feature_maps
-
- _, _, height, width = pixel_values.shape
- patch_size = self.config.patch_size
- patch_height = height // patch_size
- patch_width = width // patch_size
-
- hidden_states = self.neck(hidden_states, patch_height, patch_width)
-
- predicted_depth = self.head(hidden_states, patch_height, patch_width)
-
- loss = None
- if labels is not None:
- raise NotImplementedError("Training is not implemented yet")
-
- if not return_dict:
- if output_hidden_states:
- output = (predicted_depth,) + outputs[1:]
- else:
- output = (predicted_depth,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return DepthEstimatorOutput(
- loss=loss,
- predicted_depth=predicted_depth,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/deta/__init__.py b/transformers/models/deta/__init__.py
deleted file mode 100644
index 2d25a6a71602b38a48b23de4ab227969217ae16e..0000000000000000000000000000000000000000
--- a/transformers/models/deta/__init__.py
+++ /dev/null
@@ -1,73 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {
- "configuration_deta": ["DETA_PRETRAINED_CONFIG_ARCHIVE_MAP", "DetaConfig"],
-}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["image_processing_deta"] = ["DetaImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_deta"] = [
- "DETA_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DetaForObjectDetection",
- "DetaModel",
- "DetaPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_deta import DETA_PRETRAINED_CONFIG_ARCHIVE_MAP, DetaConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .image_processing_deta import DetaImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_deta import (
- DETA_PRETRAINED_MODEL_ARCHIVE_LIST,
- DetaForObjectDetection,
- DetaModel,
- DetaPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/deta/__pycache__/__init__.cpython-310.pyc b/transformers/models/deta/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 88497d634dc7b6cdb459a0d4667f95cdc5acbcd0..0000000000000000000000000000000000000000
Binary files a/transformers/models/deta/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deta/__pycache__/configuration_deta.cpython-310.pyc b/transformers/models/deta/__pycache__/configuration_deta.cpython-310.pyc
deleted file mode 100644
index f1ae2833fb51d2df143d1060ce387ea8b12dec3d..0000000000000000000000000000000000000000
Binary files a/transformers/models/deta/__pycache__/configuration_deta.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deta/__pycache__/convert_deta_resnet_to_pytorch.cpython-310.pyc b/transformers/models/deta/__pycache__/convert_deta_resnet_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 6b9134738213d2aa8ce26d355a1fafcfa6d779a0..0000000000000000000000000000000000000000
Binary files a/transformers/models/deta/__pycache__/convert_deta_resnet_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deta/__pycache__/convert_deta_swin_to_pytorch.cpython-310.pyc b/transformers/models/deta/__pycache__/convert_deta_swin_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 35578d349ce13c6c6df93da6867ca9a22d828186..0000000000000000000000000000000000000000
Binary files a/transformers/models/deta/__pycache__/convert_deta_swin_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deta/__pycache__/image_processing_deta.cpython-310.pyc b/transformers/models/deta/__pycache__/image_processing_deta.cpython-310.pyc
deleted file mode 100644
index b483ade010aeb9465cc8e1cbf1332bea9706d8d9..0000000000000000000000000000000000000000
Binary files a/transformers/models/deta/__pycache__/image_processing_deta.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deta/__pycache__/modeling_deta.cpython-310.pyc b/transformers/models/deta/__pycache__/modeling_deta.cpython-310.pyc
deleted file mode 100644
index 3e1e949e39d09071462454f5a08fa36d6c780654..0000000000000000000000000000000000000000
Binary files a/transformers/models/deta/__pycache__/modeling_deta.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/deta/configuration_deta.py b/transformers/models/deta/configuration_deta.py
deleted file mode 100644
index 1604bc56e6396d1c4414a0cdf8b9032827539fe9..0000000000000000000000000000000000000000
--- a/transformers/models/deta/configuration_deta.py
+++ /dev/null
@@ -1,271 +0,0 @@
-# coding=utf-8
-# Copyright 2022 SenseTime and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DETA model configuration"""
-
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-from ..auto import CONFIG_MAPPING
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DETA_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DetaConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DetaModel`]. It is used to instantiate a DETA
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the DETA
- [SenseTime/deformable-detr](https://huggingface.co/SenseTime/deformable-detr) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- backbone_config (`PretrainedConfig` or `dict`, *optional*, defaults to `ResNetConfig()`):
- The configuration of the backbone model.
- backbone (`str`, *optional*):
- Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
- will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
- is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
- use_pretrained_backbone (`bool`, *optional*, `False`):
- Whether to use pretrained weights for the backbone.
- use_timm_backbone (`bool`, *optional*, `False`):
- Whether to load `backbone` from the timm library. If `False`, the backbone is loaded from the transformers
- library.
- backbone_kwargs (`dict`, *optional*):
- Keyword arguments to be passed to AutoBackbone when loading from a checkpoint
- e.g. `{'out_indices': (0, 1, 2, 3)}`. Cannot be specified if `backbone_config` is set.
- num_queries (`int`, *optional*, defaults to 900):
- Number of object queries, i.e. detection slots. This is the maximal number of objects [`DetaModel`] can
- detect in a single image. In case `two_stage` is set to `True`, we use `two_stage_num_proposals` instead.
- d_model (`int`, *optional*, defaults to 256):
- Dimension of the layers.
- encoder_layers (`int`, *optional*, defaults to 6):
- Number of encoder layers.
- decoder_layers (`int`, *optional*, defaults to 6):
- Number of decoder layers.
- encoder_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the Transformer encoder.
- decoder_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the Transformer decoder.
- decoder_ffn_dim (`int`, *optional*, defaults to 2048):
- Dimension of the "intermediate" (often named feed-forward) layer in decoder.
- encoder_ffn_dim (`int`, *optional*, defaults to 2048):
- Dimension of the "intermediate" (often named feed-forward) layer in decoder.
- activation_function (`str` or `function`, *optional*, defaults to `"relu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- activation_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for activations inside the fully connected layer.
- init_std (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- init_xavier_std (`float`, *optional*, defaults to 1):
- The scaling factor used for the Xavier initialization gain in the HM Attention map module.
- encoder_layerdrop (`float`, *optional*, defaults to 0.0):
- The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
- for more details.
- auxiliary_loss (`bool`, *optional*, defaults to `False`):
- Whether auxiliary decoding losses (loss at each decoder layer) are to be used.
- position_embedding_type (`str`, *optional*, defaults to `"sine"`):
- Type of position embeddings to be used on top of the image features. One of `"sine"` or `"learned"`.
- class_cost (`float`, *optional*, defaults to 1):
- Relative weight of the classification error in the Hungarian matching cost.
- bbox_cost (`float`, *optional*, defaults to 5):
- Relative weight of the L1 error of the bounding box coordinates in the Hungarian matching cost.
- giou_cost (`float`, *optional*, defaults to 2):
- Relative weight of the generalized IoU loss of the bounding box in the Hungarian matching cost.
- mask_loss_coefficient (`float`, *optional*, defaults to 1):
- Relative weight of the Focal loss in the panoptic segmentation loss.
- dice_loss_coefficient (`float`, *optional*, defaults to 1):
- Relative weight of the DICE/F-1 loss in the panoptic segmentation loss.
- bbox_loss_coefficient (`float`, *optional*, defaults to 5):
- Relative weight of the L1 bounding box loss in the object detection loss.
- giou_loss_coefficient (`float`, *optional*, defaults to 2):
- Relative weight of the generalized IoU loss in the object detection loss.
- eos_coefficient (`float`, *optional*, defaults to 0.1):
- Relative classification weight of the 'no-object' class in the object detection loss.
- num_feature_levels (`int`, *optional*, defaults to 5):
- The number of input feature levels.
- encoder_n_points (`int`, *optional*, defaults to 4):
- The number of sampled keys in each feature level for each attention head in the encoder.
- decoder_n_points (`int`, *optional*, defaults to 4):
- The number of sampled keys in each feature level for each attention head in the decoder.
- two_stage (`bool`, *optional*, defaults to `True`):
- Whether to apply a two-stage deformable DETR, where the region proposals are also generated by a variant of
- DETA, which are further fed into the decoder for iterative bounding box refinement.
- two_stage_num_proposals (`int`, *optional*, defaults to 300):
- The number of region proposals to be generated, in case `two_stage` is set to `True`.
- with_box_refine (`bool`, *optional*, defaults to `True`):
- Whether to apply iterative bounding box refinement, where each decoder layer refines the bounding boxes
- based on the predictions from the previous layer.
- focal_alpha (`float`, *optional*, defaults to 0.25):
- Alpha parameter in the focal loss.
- assign_first_stage (`bool`, *optional*, defaults to `True`):
- Whether to assign each prediction i to the highest overlapping ground truth object if the overlap is larger than a threshold 0.7.
- assign_second_stage (`bool`, *optional*, defaults to `True`):
- Whether to assign second assignment procedure in the second stage closely follows the first stage assignment procedure.
- disable_custom_kernels (`bool`, *optional*, defaults to `True`):
- Disable the use of custom CUDA and CPU kernels. This option is necessary for the ONNX export, as custom
- kernels are not supported by PyTorch ONNX export.
-
- Examples:
-
- ```python
- >>> from transformers import DetaConfig, DetaModel
-
- >>> # Initializing a DETA SenseTime/deformable-detr style configuration
- >>> configuration = DetaConfig()
-
- >>> # Initializing a model (with random weights) from the SenseTime/deformable-detr style configuration
- >>> model = DetaModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "deta"
- attribute_map = {
- "hidden_size": "d_model",
- "num_attention_heads": "encoder_attention_heads",
- }
-
- def __init__(
- self,
- backbone_config=None,
- backbone=None,
- use_pretrained_backbone=False,
- use_timm_backbone=False,
- backbone_kwargs=None,
- num_queries=900,
- max_position_embeddings=2048,
- encoder_layers=6,
- encoder_ffn_dim=2048,
- encoder_attention_heads=8,
- decoder_layers=6,
- decoder_ffn_dim=1024,
- decoder_attention_heads=8,
- encoder_layerdrop=0.0,
- is_encoder_decoder=True,
- activation_function="relu",
- d_model=256,
- dropout=0.1,
- attention_dropout=0.0,
- activation_dropout=0.0,
- init_std=0.02,
- init_xavier_std=1.0,
- return_intermediate=True,
- auxiliary_loss=False,
- position_embedding_type="sine",
- num_feature_levels=5,
- encoder_n_points=4,
- decoder_n_points=4,
- two_stage=True,
- two_stage_num_proposals=300,
- with_box_refine=True,
- assign_first_stage=True,
- assign_second_stage=True,
- class_cost=1,
- bbox_cost=5,
- giou_cost=2,
- mask_loss_coefficient=1,
- dice_loss_coefficient=1,
- bbox_loss_coefficient=5,
- giou_loss_coefficient=2,
- eos_coefficient=0.1,
- focal_alpha=0.25,
- disable_custom_kernels=True,
- **kwargs,
- ):
- if use_pretrained_backbone:
- raise ValueError("Pretrained backbones are not supported yet.")
-
- if backbone_config is not None and backbone is not None:
- raise ValueError("You can't specify both `backbone` and `backbone_config`.")
-
- if backbone_config is None and backbone is None:
- logger.info("`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.")
- backbone_config = CONFIG_MAPPING["resnet"](out_features=["stage2", "stage3", "stage4"])
- else:
- if isinstance(backbone_config, dict):
- backbone_model_type = backbone_config.pop("model_type")
- config_class = CONFIG_MAPPING[backbone_model_type]
- backbone_config = config_class.from_dict(backbone_config)
-
- if backbone_kwargs is not None and backbone_kwargs and backbone_config is not None:
- raise ValueError("You can't specify both `backbone_kwargs` and `backbone_config`.")
-
- self.backbone_config = backbone_config
- self.backbone = backbone
- self.use_pretrained_backbone = use_pretrained_backbone
- self.use_timm_backbone = use_timm_backbone
- self.backbone_kwargs = backbone_kwargs
- self.num_queries = num_queries
- self.max_position_embeddings = max_position_embeddings
- self.d_model = d_model
- self.encoder_ffn_dim = encoder_ffn_dim
- self.encoder_layers = encoder_layers
- self.encoder_attention_heads = encoder_attention_heads
- self.decoder_ffn_dim = decoder_ffn_dim
- self.decoder_layers = decoder_layers
- self.decoder_attention_heads = decoder_attention_heads
- self.dropout = dropout
- self.attention_dropout = attention_dropout
- self.activation_dropout = activation_dropout
- self.activation_function = activation_function
- self.init_std = init_std
- self.init_xavier_std = init_xavier_std
- self.encoder_layerdrop = encoder_layerdrop
- self.auxiliary_loss = auxiliary_loss
- self.position_embedding_type = position_embedding_type
- # deformable attributes
- self.num_feature_levels = num_feature_levels
- self.encoder_n_points = encoder_n_points
- self.decoder_n_points = decoder_n_points
- self.two_stage = two_stage
- self.two_stage_num_proposals = two_stage_num_proposals
- self.with_box_refine = with_box_refine
- self.assign_first_stage = assign_first_stage
- self.assign_second_stage = assign_second_stage
- if two_stage is True and with_box_refine is False:
- raise ValueError("If two_stage is True, with_box_refine must be True.")
- # Hungarian matcher
- self.class_cost = class_cost
- self.bbox_cost = bbox_cost
- self.giou_cost = giou_cost
- # Loss coefficients
- self.mask_loss_coefficient = mask_loss_coefficient
- self.dice_loss_coefficient = dice_loss_coefficient
- self.bbox_loss_coefficient = bbox_loss_coefficient
- self.giou_loss_coefficient = giou_loss_coefficient
- self.eos_coefficient = eos_coefficient
- self.focal_alpha = focal_alpha
- self.disable_custom_kernels = disable_custom_kernels
- super().__init__(is_encoder_decoder=is_encoder_decoder, **kwargs)
-
- @property
- def num_attention_heads(self) -> int:
- return self.encoder_attention_heads
-
- @property
- def hidden_size(self) -> int:
- return self.d_model
diff --git a/transformers/models/deta/convert_deta_resnet_to_pytorch.py b/transformers/models/deta/convert_deta_resnet_to_pytorch.py
deleted file mode 100644
index cc17568bd64133169b047a3d767bcbf1b2582b25..0000000000000000000000000000000000000000
--- a/transformers/models/deta/convert_deta_resnet_to_pytorch.py
+++ /dev/null
@@ -1,320 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DETA checkpoints from the original repository.
-
-URL: https://github.com/jozhang97/DETA/tree/master"""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
-from PIL import Image
-
-from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_deta_config():
- config = DetaConfig(
- num_queries=900,
- encoder_ffn_dim=2048,
- decoder_ffn_dim=2048,
- num_feature_levels=5,
- assign_first_stage=True,
- with_box_refine=True,
- two_stage=True,
- )
-
- # set labels
- config.num_labels = 91
- repo_id = "huggingface/label-files"
- filename = "coco-detection-id2label.json"
- id2label = json.load(open(cached_download(hf_hub_url(repo_id, filename, repo_type="dataset")), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
-
- return config
-
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys(config):
- rename_keys = []
-
- # stem
- # fmt: off
- rename_keys.append(("backbone.0.body.conv1.weight", "model.backbone.model.embedder.embedder.convolution.weight"))
- rename_keys.append(("backbone.0.body.bn1.weight", "model.backbone.model.embedder.embedder.normalization.weight"))
- rename_keys.append(("backbone.0.body.bn1.bias", "model.backbone.model.embedder.embedder.normalization.bias"))
- rename_keys.append(("backbone.0.body.bn1.running_mean", "model.backbone.model.embedder.embedder.normalization.running_mean"))
- rename_keys.append(("backbone.0.body.bn1.running_var", "model.backbone.model.embedder.embedder.normalization.running_var"))
- # stages
- for stage_idx in range(len(config.backbone_config.depths)):
- for layer_idx in range(config.backbone_config.depths[stage_idx]):
- # shortcut
- if layer_idx == 0:
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var",
- )
- )
- # 3 convs
- for i in range(3):
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var",
- f"model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var",
- )
- )
- # transformer encoder
- for i in range(config.encoder_layers):
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight", f"model.encoder.layers.{i}.self_attn.sampling_offsets.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias", f"model.encoder.layers.{i}.self_attn.sampling_offsets.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.attention_weights.weight", f"model.encoder.layers.{i}.self_attn.attention_weights.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.attention_weights.bias", f"model.encoder.layers.{i}.self_attn.attention_weights.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.value_proj.weight", f"model.encoder.layers.{i}.self_attn.value_proj.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.value_proj.bias", f"model.encoder.layers.{i}.self_attn.value_proj.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.output_proj.weight", f"model.encoder.layers.{i}.self_attn.output_proj.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.output_proj.bias", f"model.encoder.layers.{i}.self_attn.output_proj.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm1.weight", f"model.encoder.layers.{i}.self_attn_layer_norm.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm1.bias", f"model.encoder.layers.{i}.self_attn_layer_norm.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.weight", f"model.encoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.bias", f"model.encoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.weight", f"model.encoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.bias", f"model.encoder.layers.{i}.fc2.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.weight", f"model.encoder.layers.{i}.final_layer_norm.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.bias", f"model.encoder.layers.{i}.final_layer_norm.bias"))
-
- # transformer decoder
- for i in range(config.decoder_layers):
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight", f"model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias", f"model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.attention_weights.weight", f"model.decoder.layers.{i}.encoder_attn.attention_weights.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.attention_weights.bias", f"model.decoder.layers.{i}.encoder_attn.attention_weights.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.value_proj.weight", f"model.decoder.layers.{i}.encoder_attn.value_proj.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.value_proj.bias", f"model.decoder.layers.{i}.encoder_attn.value_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.output_proj.weight", f"model.decoder.layers.{i}.encoder_attn.output_proj.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.output_proj.bias", f"model.decoder.layers.{i}.encoder_attn.output_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm1.weight", f"model.decoder.layers.{i}.encoder_attn_layer_norm.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm1.bias", f"model.decoder.layers.{i}.encoder_attn_layer_norm.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.self_attn.out_proj.weight", f"model.decoder.layers.{i}.self_attn.out_proj.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.self_attn.out_proj.bias", f"model.decoder.layers.{i}.self_attn.out_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm2.weight", f"model.decoder.layers.{i}.self_attn_layer_norm.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm2.bias", f"model.decoder.layers.{i}.self_attn_layer_norm.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.weight", f"model.decoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.bias", f"model.decoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.weight", f"model.decoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.bias", f"model.decoder.layers.{i}.fc2.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.weight", f"model.decoder.layers.{i}.final_layer_norm.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.bias", f"model.decoder.layers.{i}.final_layer_norm.bias"))
-
- # fmt: on
-
- return rename_keys
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-def read_in_decoder_q_k_v(state_dict, config):
- # transformer decoder self-attention layers
- hidden_size = config.d_model
- for i in range(config.decoder_layers):
- # read in weights + bias of input projection layer of self-attention
- in_proj_weight = state_dict.pop(f"transformer.decoder.layers.{i}.self_attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"transformer.decoder.layers.{i}.self_attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"model.decoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:hidden_size, :]
- state_dict[f"model.decoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:hidden_size]
- state_dict[f"model.decoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[
- hidden_size : hidden_size * 2, :
- ]
- state_dict[f"model.decoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[hidden_size : hidden_size * 2]
- state_dict[f"model.decoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-hidden_size:, :]
- state_dict[f"model.decoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-hidden_size:]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
-
- return im
-
-
-@torch.no_grad()
-def convert_deta_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub):
- """
- Copy/paste/tweak model's weights to our DETA structure.
- """
-
- # load config
- config = get_deta_config()
-
- # load original state dict
- if model_name == "deta-resnet-50":
- filename = "adet_checkpoint0011.pth"
- elif model_name == "deta-resnet-50-24-epochs":
- filename = "adet_2x_checkpoint0023.pth"
- else:
- raise ValueError(f"Model name {model_name} not supported")
- checkpoint_path = hf_hub_download(repo_id="nielsr/deta-checkpoints", filename=filename)
- state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
-
- # rename keys
- rename_keys = create_rename_keys(config)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- read_in_decoder_q_k_v(state_dict, config)
-
- # fix some prefixes
- for key in state_dict.copy().keys():
- if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
- val = state_dict.pop(key)
- state_dict[key.replace("transformer.decoder", "model.decoder")] = val
- if "input_proj" in key:
- val = state_dict.pop(key)
- state_dict["model." + key] = val
- if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
- val = state_dict.pop(key)
- state_dict[key.replace("transformer", "model")] = val
-
- # finally, create HuggingFace model and load state dict
- model = DetaForObjectDetection(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- device = "cuda" if torch.cuda.is_available() else "cpu"
- model.to(device)
-
- # load image processor
- processor = DetaImageProcessor(format="coco_detection")
-
- # verify our conversion on image
- img = prepare_img()
- encoding = processor(images=img, return_tensors="pt")
- pixel_values = encoding["pixel_values"]
- outputs = model(pixel_values.to(device))
-
- # verify logits
- if model_name == "deta-resnet-50":
- expected_logits = torch.tensor(
- [[-7.3978, -2.5406, -4.1668], [-8.2684, -3.9933, -3.8096], [-7.0515, -3.7973, -5.8516]]
- )
- expected_boxes = torch.tensor([[0.5043, 0.4973, 0.9998], [0.2542, 0.5489, 0.4748], [0.5490, 0.2765, 0.0570]])
- elif model_name == "deta-resnet-50-24-epochs":
- expected_logits = torch.tensor(
- [[-7.1688, -2.4857, -4.8669], [-7.8630, -3.8154, -4.2674], [-7.2730, -4.1865, -5.5323]]
- )
- expected_boxes = torch.tensor([[0.5021, 0.4971, 0.9994], [0.2546, 0.5486, 0.4731], [0.1686, 0.1986, 0.2142]])
-
- assert torch.allclose(outputs.logits[0, :3, :3], expected_logits.to(device), atol=1e-4)
- assert torch.allclose(outputs.pred_boxes[0, :3, :3], expected_boxes.to(device), atol=1e-4)
- print("Everything ok!")
-
- if pytorch_dump_folder_path:
- # Save model and processor
- logger.info(f"Saving PyTorch model and processor to {pytorch_dump_folder_path}...")
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- # Push to hub
- if push_to_hub:
- print("Pushing model and processor to hub...")
- model.push_to_hub(f"jozhang97/{model_name}")
- processor.push_to_hub(f"jozhang97/{model_name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- "--model_name",
- type=str,
- default="deta-resnet-50",
- choices=["deta-resnet-50", "deta-resnet-50-24-epochs"],
- help="Name of the model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- help="Path to the folder to output PyTorch model.",
- )
- parser.add_argument(
- "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
- )
- args = parser.parse_args()
- convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/deta/convert_deta_swin_to_pytorch.py b/transformers/models/deta/convert_deta_swin_to_pytorch.py
deleted file mode 100644
index 911bc434e14265f9fe21dc8166b4d9eafb0d9cc0..0000000000000000000000000000000000000000
--- a/transformers/models/deta/convert_deta_swin_to_pytorch.py
+++ /dev/null
@@ -1,327 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DETA checkpoints from the original repository.
-
-URL: https://github.com/jozhang97/DETA/tree/master"""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
-from PIL import Image
-
-from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor, SwinConfig
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_deta_config(model_name):
- backbone_config = SwinConfig(
- embed_dim=192,
- depths=(2, 2, 18, 2),
- num_heads=(6, 12, 24, 48),
- window_size=12,
- out_features=["stage2", "stage3", "stage4"],
- )
-
- config = DetaConfig(
- backbone_config=backbone_config,
- num_queries=900,
- encoder_ffn_dim=2048,
- decoder_ffn_dim=2048,
- num_feature_levels=5,
- assign_first_stage=True,
- with_box_refine=True,
- two_stage=True,
- )
-
- # set labels
- repo_id = "huggingface/label-files"
- if "o365" in model_name:
- num_labels = 366
- filename = "object365-id2label.json"
- else:
- num_labels = 91
- filename = "coco-detection-id2label.json"
-
- config.num_labels = num_labels
- id2label = json.load(open(cached_download(hf_hub_url(repo_id, filename, repo_type="dataset")), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
-
- return config
-
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys(config):
- rename_keys = []
-
- # stem
- # fmt: off
- rename_keys.append(("backbone.0.body.patch_embed.proj.weight", "model.backbone.model.embeddings.patch_embeddings.projection.weight"))
- rename_keys.append(("backbone.0.body.patch_embed.proj.bias", "model.backbone.model.embeddings.patch_embeddings.projection.bias"))
- rename_keys.append(("backbone.0.body.patch_embed.norm.weight", "model.backbone.model.embeddings.norm.weight"))
- rename_keys.append(("backbone.0.body.patch_embed.norm.bias", "model.backbone.model.embeddings.norm.bias"))
- # stages
- for i in range(len(config.backbone_config.depths)):
- for j in range(config.backbone_config.depths[i]):
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.norm1.weight", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.norm1.bias", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_bias_table", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_index", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.attn.proj.weight", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.attn.proj.bias", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.norm2.weight", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.norm2.bias", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.weight", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.bias", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.weight", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.weight"))
- rename_keys.append((f"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.bias", f"model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.bias"))
-
- if i < 3:
- rename_keys.append((f"backbone.0.body.layers.{i}.downsample.reduction.weight", f"model.backbone.model.encoder.layers.{i}.downsample.reduction.weight"))
- rename_keys.append((f"backbone.0.body.layers.{i}.downsample.norm.weight", f"model.backbone.model.encoder.layers.{i}.downsample.norm.weight"))
- rename_keys.append((f"backbone.0.body.layers.{i}.downsample.norm.bias", f"model.backbone.model.encoder.layers.{i}.downsample.norm.bias"))
-
- rename_keys.append(("backbone.0.body.norm1.weight", "model.backbone.model.hidden_states_norms.stage2.weight"))
- rename_keys.append(("backbone.0.body.norm1.bias", "model.backbone.model.hidden_states_norms.stage2.bias"))
- rename_keys.append(("backbone.0.body.norm2.weight", "model.backbone.model.hidden_states_norms.stage3.weight"))
- rename_keys.append(("backbone.0.body.norm2.bias", "model.backbone.model.hidden_states_norms.stage3.bias"))
- rename_keys.append(("backbone.0.body.norm3.weight", "model.backbone.model.hidden_states_norms.stage4.weight"))
- rename_keys.append(("backbone.0.body.norm3.bias", "model.backbone.model.hidden_states_norms.stage4.bias"))
-
- # transformer encoder
- for i in range(config.encoder_layers):
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight", f"model.encoder.layers.{i}.self_attn.sampling_offsets.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias", f"model.encoder.layers.{i}.self_attn.sampling_offsets.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.attention_weights.weight", f"model.encoder.layers.{i}.self_attn.attention_weights.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.attention_weights.bias", f"model.encoder.layers.{i}.self_attn.attention_weights.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.value_proj.weight", f"model.encoder.layers.{i}.self_attn.value_proj.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.value_proj.bias", f"model.encoder.layers.{i}.self_attn.value_proj.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.output_proj.weight", f"model.encoder.layers.{i}.self_attn.output_proj.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.self_attn.output_proj.bias", f"model.encoder.layers.{i}.self_attn.output_proj.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm1.weight", f"model.encoder.layers.{i}.self_attn_layer_norm.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm1.bias", f"model.encoder.layers.{i}.self_attn_layer_norm.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.weight", f"model.encoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.bias", f"model.encoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.weight", f"model.encoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.bias", f"model.encoder.layers.{i}.fc2.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.weight", f"model.encoder.layers.{i}.final_layer_norm.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.bias", f"model.encoder.layers.{i}.final_layer_norm.bias"))
-
- # transformer decoder
- for i in range(config.decoder_layers):
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight", f"model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias", f"model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.attention_weights.weight", f"model.decoder.layers.{i}.encoder_attn.attention_weights.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.attention_weights.bias", f"model.decoder.layers.{i}.encoder_attn.attention_weights.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.value_proj.weight", f"model.decoder.layers.{i}.encoder_attn.value_proj.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.value_proj.bias", f"model.decoder.layers.{i}.encoder_attn.value_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.output_proj.weight", f"model.decoder.layers.{i}.encoder_attn.output_proj.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.cross_attn.output_proj.bias", f"model.decoder.layers.{i}.encoder_attn.output_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm1.weight", f"model.decoder.layers.{i}.encoder_attn_layer_norm.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm1.bias", f"model.decoder.layers.{i}.encoder_attn_layer_norm.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.self_attn.out_proj.weight", f"model.decoder.layers.{i}.self_attn.out_proj.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.self_attn.out_proj.bias", f"model.decoder.layers.{i}.self_attn.out_proj.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm2.weight", f"model.decoder.layers.{i}.self_attn_layer_norm.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm2.bias", f"model.decoder.layers.{i}.self_attn_layer_norm.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.weight", f"model.decoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.bias", f"model.decoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.weight", f"model.decoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.bias", f"model.decoder.layers.{i}.fc2.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.weight", f"model.decoder.layers.{i}.final_layer_norm.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.bias", f"model.decoder.layers.{i}.final_layer_norm.bias"))
-
- # fmt: on
-
- return rename_keys
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_swin_q_k_v(state_dict, backbone_config):
- num_features = [int(backbone_config.embed_dim * 2**i) for i in range(len(backbone_config.depths))]
- for i in range(len(backbone_config.depths)):
- dim = num_features[i]
- for j in range(backbone_config.depths[i]):
- # fmt: off
- # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight")
- in_proj_bias = state_dict.pop(f"backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.query.weight"] = in_proj_weight[:dim, :]
- state_dict[f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.query.bias"] = in_proj_bias[: dim]
- state_dict[f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.key.weight"] = in_proj_weight[
- dim : dim * 2, :
- ]
- state_dict[f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.key.bias"] = in_proj_bias[
- dim : dim * 2
- ]
- state_dict[f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.value.weight"] = in_proj_weight[
- -dim :, :
- ]
- state_dict[f"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.value.bias"] = in_proj_bias[-dim :]
- # fmt: on
-
-
-def read_in_decoder_q_k_v(state_dict, config):
- # transformer decoder self-attention layers
- hidden_size = config.d_model
- for i in range(config.decoder_layers):
- # read in weights + bias of input projection layer of self-attention
- in_proj_weight = state_dict.pop(f"transformer.decoder.layers.{i}.self_attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"transformer.decoder.layers.{i}.self_attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"model.decoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:hidden_size, :]
- state_dict[f"model.decoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:hidden_size]
- state_dict[f"model.decoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[
- hidden_size : hidden_size * 2, :
- ]
- state_dict[f"model.decoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[hidden_size : hidden_size * 2]
- state_dict[f"model.decoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-hidden_size:, :]
- state_dict[f"model.decoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-hidden_size:]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
-
- return im
-
-
-@torch.no_grad()
-def convert_deta_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub):
- """
- Copy/paste/tweak model's weights to our DETA structure.
- """
-
- # load config
- config = get_deta_config(model_name)
-
- # load original state dict
- if model_name == "deta-swin-large":
- checkpoint_path = hf_hub_download(repo_id="nielsr/deta-checkpoints", filename="adet_swin_ft.pth")
- elif model_name == "deta-swin-large-o365":
- checkpoint_path = hf_hub_download(repo_id="jozhang97/deta-swin-l-o365", filename="deta_swin_pt_o365.pth")
- else:
- raise ValueError(f"Model name {model_name} not supported")
-
- state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
-
- # original state dict
- for name, param in state_dict.items():
- print(name, param.shape)
-
- # rename keys
- rename_keys = create_rename_keys(config)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- read_in_swin_q_k_v(state_dict, config.backbone_config)
- read_in_decoder_q_k_v(state_dict, config)
-
- # fix some prefixes
- for key in state_dict.copy().keys():
- if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
- val = state_dict.pop(key)
- state_dict[key.replace("transformer.decoder", "model.decoder")] = val
- if "input_proj" in key:
- val = state_dict.pop(key)
- state_dict["model." + key] = val
- if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
- val = state_dict.pop(key)
- state_dict[key.replace("transformer", "model")] = val
-
- # finally, create HuggingFace model and load state dict
- model = DetaForObjectDetection(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- device = "cuda" if torch.cuda.is_available() else "cpu"
- model.to(device)
-
- # load image processor
- processor = DetaImageProcessor(format="coco_detection")
-
- # verify our conversion on image
- img = prepare_img()
- encoding = processor(images=img, return_tensors="pt")
- pixel_values = encoding["pixel_values"]
- outputs = model(pixel_values.to(device))
-
- # verify logits
- print("Logits:", outputs.logits[0, :3, :3])
- print("Boxes:", outputs.pred_boxes[0, :3, :3])
- if model_name == "deta-swin-large":
- expected_logits = torch.tensor(
- [[-7.6308, -2.8485, -5.3737], [-7.2037, -4.5505, -4.8027], [-7.2943, -4.2611, -4.6617]]
- )
- expected_boxes = torch.tensor([[0.4987, 0.4969, 0.9999], [0.2549, 0.5498, 0.4805], [0.5498, 0.2757, 0.0569]])
- elif model_name == "deta-swin-large-o365":
- expected_logits = torch.tensor(
- [[-8.0122, -3.5720, -4.9717], [-8.1547, -3.6886, -4.6389], [-7.6610, -3.6194, -5.0134]]
- )
- expected_boxes = torch.tensor([[0.2523, 0.5549, 0.4881], [0.7715, 0.4149, 0.4601], [0.5503, 0.2753, 0.0575]])
- assert torch.allclose(outputs.logits[0, :3, :3], expected_logits.to(device), atol=1e-4)
- assert torch.allclose(outputs.pred_boxes[0, :3, :3], expected_boxes.to(device), atol=1e-4)
- print("Everything ok!")
-
- if pytorch_dump_folder_path:
- # Save model and processor
- logger.info(f"Saving PyTorch model and processor to {pytorch_dump_folder_path}...")
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- # Push to hub
- if push_to_hub:
- print("Pushing model and processor to hub...")
- model.push_to_hub(f"jozhang97/{model_name}")
- processor.push_to_hub(f"jozhang97/{model_name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- "--model_name",
- type=str,
- default="deta-swin-large",
- choices=["deta-swin-large", "deta-swin-large-o365"],
- help="Name of the model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- help="Path to the folder to output PyTorch model.",
- )
- parser.add_argument(
- "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
- )
- args = parser.parse_args()
- convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/deta/image_processing_deta.py b/transformers/models/deta/image_processing_deta.py
deleted file mode 100644
index 45c5c6cb285a8f5851d23bc8aa38e764593c0846..0000000000000000000000000000000000000000
--- a/transformers/models/deta/image_processing_deta.py
+++ /dev/null
@@ -1,1174 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for Deformable DETR."""
-
-import pathlib
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
-
-import numpy as np
-
-from ...feature_extraction_utils import BatchFeature
-from ...image_processing_utils import BaseImageProcessor, get_size_dict
-from ...image_transforms import (
- PaddingMode,
- center_to_corners_format,
- corners_to_center_format,
- pad,
- rescale,
- resize,
- rgb_to_id,
- to_channel_dimension_format,
-)
-from ...image_utils import (
- IMAGENET_DEFAULT_MEAN,
- IMAGENET_DEFAULT_STD,
- AnnotationFormat,
- AnnotationType,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- get_image_size,
- infer_channel_dimension_format,
- is_batched,
- is_scaled_image,
- to_numpy_array,
- valid_images,
- validate_annotations,
- validate_preprocess_arguments,
-)
-from ...utils import (
- is_flax_available,
- is_jax_tensor,
- is_tf_available,
- is_tf_tensor,
- is_torch_available,
- is_torch_tensor,
- is_torchvision_available,
- is_vision_available,
- logging,
-)
-from ...utils.generic import TensorType
-
-
-if is_torch_available():
- import torch
-
-
-if is_torchvision_available():
- from torchvision.ops.boxes import batched_nms
-
-if is_vision_available():
- import PIL
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-SUPPORTED_ANNOTATION_FORMATS = (AnnotationFormat.COCO_DETECTION, AnnotationFormat.COCO_PANOPTIC)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_size_with_aspect_ratio
-def get_size_with_aspect_ratio(image_size, size, max_size=None) -> Tuple[int, int]:
- """
- Computes the output image size given the input image size and the desired output size.
-
- Args:
- image_size (`Tuple[int, int]`):
- The input image size.
- size (`int`):
- The desired output size.
- max_size (`int`, *optional*):
- The maximum allowed output size.
- """
- height, width = image_size
- if max_size is not None:
- min_original_size = float(min((height, width)))
- max_original_size = float(max((height, width)))
- if max_original_size / min_original_size * size > max_size:
- size = int(round(max_size * min_original_size / max_original_size))
-
- if (height <= width and height == size) or (width <= height and width == size):
- return height, width
-
- if width < height:
- ow = size
- oh = int(size * height / width)
- else:
- oh = size
- ow = int(size * width / height)
- return (oh, ow)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_resize_output_image_size
-def get_resize_output_image_size(
- input_image: np.ndarray,
- size: Union[int, Tuple[int, int], List[int]],
- max_size: Optional[int] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
-) -> Tuple[int, int]:
- """
- Computes the output image size given the input image size and the desired output size. If the desired output size
- is a tuple or list, the output image size is returned as is. If the desired output size is an integer, the output
- image size is computed by keeping the aspect ratio of the input image size.
-
- Args:
- input_image (`np.ndarray`):
- The image to resize.
- size (`int` or `Tuple[int, int]` or `List[int]`):
- The desired output size.
- max_size (`int`, *optional*):
- The maximum allowed output size.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred from the input image.
- """
- image_size = get_image_size(input_image, input_data_format)
- if isinstance(size, (list, tuple)):
- return size
-
- return get_size_with_aspect_ratio(image_size, size, max_size)
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_numpy_to_framework_fn
-def get_numpy_to_framework_fn(arr) -> Callable:
- """
- Returns a function that converts a numpy array to the framework of the input array.
-
- Args:
- arr (`np.ndarray`): The array to convert.
- """
- if isinstance(arr, np.ndarray):
- return np.array
- if is_tf_available() and is_tf_tensor(arr):
- import tensorflow as tf
-
- return tf.convert_to_tensor
- if is_torch_available() and is_torch_tensor(arr):
- import torch
-
- return torch.tensor
- if is_flax_available() and is_jax_tensor(arr):
- import jax.numpy as jnp
-
- return jnp.array
- raise ValueError(f"Cannot convert arrays of type {type(arr)}")
-
-
-# Copied from transformers.models.detr.image_processing_detr.safe_squeeze
-def safe_squeeze(arr: np.ndarray, axis: Optional[int] = None) -> np.ndarray:
- """
- Squeezes an array, but only if the axis specified has dim 1.
- """
- if axis is None:
- return arr.squeeze()
-
- try:
- return arr.squeeze(axis=axis)
- except ValueError:
- return arr
-
-
-# Copied from transformers.models.detr.image_processing_detr.normalize_annotation
-def normalize_annotation(annotation: Dict, image_size: Tuple[int, int]) -> Dict:
- image_height, image_width = image_size
- norm_annotation = {}
- for key, value in annotation.items():
- if key == "boxes":
- boxes = value
- boxes = corners_to_center_format(boxes)
- boxes /= np.asarray([image_width, image_height, image_width, image_height], dtype=np.float32)
- norm_annotation[key] = boxes
- else:
- norm_annotation[key] = value
- return norm_annotation
-
-
-# Copied from transformers.models.detr.image_processing_detr.max_across_indices
-def max_across_indices(values: Iterable[Any]) -> List[Any]:
- """
- Return the maximum value across all indices of an iterable of values.
- """
- return [max(values_i) for values_i in zip(*values)]
-
-
-# Copied from transformers.models.detr.image_processing_detr.get_max_height_width
-def get_max_height_width(
- images: List[np.ndarray], input_data_format: Optional[Union[str, ChannelDimension]] = None
-) -> List[int]:
- """
- Get the maximum height and width across all images in a batch.
- """
- if input_data_format is None:
- input_data_format = infer_channel_dimension_format(images[0])
-
- if input_data_format == ChannelDimension.FIRST:
- _, max_height, max_width = max_across_indices([img.shape for img in images])
- elif input_data_format == ChannelDimension.LAST:
- max_height, max_width, _ = max_across_indices([img.shape for img in images])
- else:
- raise ValueError(f"Invalid channel dimension format: {input_data_format}")
- return (max_height, max_width)
-
-
-# Copied from transformers.models.detr.image_processing_detr.make_pixel_mask
-def make_pixel_mask(
- image: np.ndarray, output_size: Tuple[int, int], input_data_format: Optional[Union[str, ChannelDimension]] = None
-) -> np.ndarray:
- """
- Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.
-
- Args:
- image (`np.ndarray`):
- Image to make the pixel mask for.
- output_size (`Tuple[int, int]`):
- Output size of the mask.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- mask = np.zeros(output_size, dtype=np.int64)
- mask[:input_height, :input_width] = 1
- return mask
-
-
-# Copied from transformers.models.detr.image_processing_detr.convert_coco_poly_to_mask
-def convert_coco_poly_to_mask(segmentations, height: int, width: int) -> np.ndarray:
- """
- Convert a COCO polygon annotation to a mask.
-
- Args:
- segmentations (`List[List[float]]`):
- List of polygons, each polygon represented by a list of x-y coordinates.
- height (`int`):
- Height of the mask.
- width (`int`):
- Width of the mask.
- """
- try:
- from pycocotools import mask as coco_mask
- except ImportError:
- raise ImportError("Pycocotools is not installed in your environment.")
-
- masks = []
- for polygons in segmentations:
- rles = coco_mask.frPyObjects(polygons, height, width)
- mask = coco_mask.decode(rles)
- if len(mask.shape) < 3:
- mask = mask[..., None]
- mask = np.asarray(mask, dtype=np.uint8)
- mask = np.any(mask, axis=2)
- masks.append(mask)
- if masks:
- masks = np.stack(masks, axis=0)
- else:
- masks = np.zeros((0, height, width), dtype=np.uint8)
-
- return masks
-
-
-# Copied from transformers.models.detr.image_processing_detr.prepare_coco_detection_annotation with DETR->DETA
-def prepare_coco_detection_annotation(
- image,
- target,
- return_segmentation_masks: bool = False,
- input_data_format: Optional[Union[ChannelDimension, str]] = None,
-):
- """
- Convert the target in COCO format into the format expected by DETA.
- """
- image_height, image_width = get_image_size(image, channel_dim=input_data_format)
-
- image_id = target["image_id"]
- image_id = np.asarray([image_id], dtype=np.int64)
-
- # Get all COCO annotations for the given image.
- annotations = target["annotations"]
- annotations = [obj for obj in annotations if "iscrowd" not in obj or obj["iscrowd"] == 0]
-
- classes = [obj["category_id"] for obj in annotations]
- classes = np.asarray(classes, dtype=np.int64)
-
- # for conversion to coco api
- area = np.asarray([obj["area"] for obj in annotations], dtype=np.float32)
- iscrowd = np.asarray([obj["iscrowd"] if "iscrowd" in obj else 0 for obj in annotations], dtype=np.int64)
-
- boxes = [obj["bbox"] for obj in annotations]
- # guard against no boxes via resizing
- boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4)
- boxes[:, 2:] += boxes[:, :2]
- boxes[:, 0::2] = boxes[:, 0::2].clip(min=0, max=image_width)
- boxes[:, 1::2] = boxes[:, 1::2].clip(min=0, max=image_height)
-
- keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
-
- new_target = {}
- new_target["image_id"] = image_id
- new_target["class_labels"] = classes[keep]
- new_target["boxes"] = boxes[keep]
- new_target["area"] = area[keep]
- new_target["iscrowd"] = iscrowd[keep]
- new_target["orig_size"] = np.asarray([int(image_height), int(image_width)], dtype=np.int64)
-
- if annotations and "keypoints" in annotations[0]:
- keypoints = [obj["keypoints"] for obj in annotations]
- # Converting the filtered keypoints list to a numpy array
- keypoints = np.asarray(keypoints, dtype=np.float32)
- # Apply the keep mask here to filter the relevant annotations
- keypoints = keypoints[keep]
- num_keypoints = keypoints.shape[0]
- keypoints = keypoints.reshape((-1, 3)) if num_keypoints else keypoints
- new_target["keypoints"] = keypoints
-
- if return_segmentation_masks:
- segmentation_masks = [obj["segmentation"] for obj in annotations]
- masks = convert_coco_poly_to_mask(segmentation_masks, image_height, image_width)
- new_target["masks"] = masks[keep]
-
- return new_target
-
-
-# Copied from transformers.models.detr.image_processing_detr.masks_to_boxes
-def masks_to_boxes(masks: np.ndarray) -> np.ndarray:
- """
- Compute the bounding boxes around the provided panoptic segmentation masks.
-
- Args:
- masks: masks in format `[number_masks, height, width]` where N is the number of masks
-
- Returns:
- boxes: bounding boxes in format `[number_masks, 4]` in xyxy format
- """
- if masks.size == 0:
- return np.zeros((0, 4))
-
- h, w = masks.shape[-2:]
- y = np.arange(0, h, dtype=np.float32)
- x = np.arange(0, w, dtype=np.float32)
- # see https://github.com/pytorch/pytorch/issues/50276
- y, x = np.meshgrid(y, x, indexing="ij")
-
- x_mask = masks * np.expand_dims(x, axis=0)
- x_max = x_mask.reshape(x_mask.shape[0], -1).max(-1)
- x = np.ma.array(x_mask, mask=~(np.array(masks, dtype=bool)))
- x_min = x.filled(fill_value=1e8)
- x_min = x_min.reshape(x_min.shape[0], -1).min(-1)
-
- y_mask = masks * np.expand_dims(y, axis=0)
- y_max = y_mask.reshape(x_mask.shape[0], -1).max(-1)
- y = np.ma.array(y_mask, mask=~(np.array(masks, dtype=bool)))
- y_min = y.filled(fill_value=1e8)
- y_min = y_min.reshape(y_min.shape[0], -1).min(-1)
-
- return np.stack([x_min, y_min, x_max, y_max], 1)
-
-
-# Copied from transformers.models.detr.image_processing_detr.prepare_coco_panoptic_annotation with DETR->DETA
-def prepare_coco_panoptic_annotation(
- image: np.ndarray,
- target: Dict,
- masks_path: Union[str, pathlib.Path],
- return_masks: bool = True,
- input_data_format: Union[ChannelDimension, str] = None,
-) -> Dict:
- """
- Prepare a coco panoptic annotation for DETA.
- """
- image_height, image_width = get_image_size(image, channel_dim=input_data_format)
- annotation_path = pathlib.Path(masks_path) / target["file_name"]
-
- new_target = {}
- new_target["image_id"] = np.asarray([target["image_id"] if "image_id" in target else target["id"]], dtype=np.int64)
- new_target["size"] = np.asarray([image_height, image_width], dtype=np.int64)
- new_target["orig_size"] = np.asarray([image_height, image_width], dtype=np.int64)
-
- if "segments_info" in target:
- masks = np.asarray(PIL.Image.open(annotation_path), dtype=np.uint32)
- masks = rgb_to_id(masks)
-
- ids = np.array([segment_info["id"] for segment_info in target["segments_info"]])
- masks = masks == ids[:, None, None]
- masks = masks.astype(np.uint8)
- if return_masks:
- new_target["masks"] = masks
- new_target["boxes"] = masks_to_boxes(masks)
- new_target["class_labels"] = np.array(
- [segment_info["category_id"] for segment_info in target["segments_info"]], dtype=np.int64
- )
- new_target["iscrowd"] = np.asarray(
- [segment_info["iscrowd"] for segment_info in target["segments_info"]], dtype=np.int64
- )
- new_target["area"] = np.asarray(
- [segment_info["area"] for segment_info in target["segments_info"]], dtype=np.float32
- )
-
- return new_target
-
-
-# Copied from transformers.models.detr.image_processing_detr.resize_annotation
-def resize_annotation(
- annotation: Dict[str, Any],
- orig_size: Tuple[int, int],
- target_size: Tuple[int, int],
- threshold: float = 0.5,
- resample: PILImageResampling = PILImageResampling.NEAREST,
-):
- """
- Resizes an annotation to a target size.
-
- Args:
- annotation (`Dict[str, Any]`):
- The annotation dictionary.
- orig_size (`Tuple[int, int]`):
- The original size of the input image.
- target_size (`Tuple[int, int]`):
- The target size of the image, as returned by the preprocessing `resize` step.
- threshold (`float`, *optional*, defaults to 0.5):
- The threshold used to binarize the segmentation masks.
- resample (`PILImageResampling`, defaults to `PILImageResampling.NEAREST`):
- The resampling filter to use when resizing the masks.
- """
- ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(target_size, orig_size))
- ratio_height, ratio_width = ratios
-
- new_annotation = {}
- new_annotation["size"] = target_size
-
- for key, value in annotation.items():
- if key == "boxes":
- boxes = value
- scaled_boxes = boxes * np.asarray([ratio_width, ratio_height, ratio_width, ratio_height], dtype=np.float32)
- new_annotation["boxes"] = scaled_boxes
- elif key == "area":
- area = value
- scaled_area = area * (ratio_width * ratio_height)
- new_annotation["area"] = scaled_area
- elif key == "masks":
- masks = value[:, None]
- masks = np.array([resize(mask, target_size, resample=resample) for mask in masks])
- masks = masks.astype(np.float32)
- masks = masks[:, 0] > threshold
- new_annotation["masks"] = masks
- elif key == "size":
- new_annotation["size"] = target_size
- else:
- new_annotation[key] = value
-
- return new_annotation
-
-
-class DetaImageProcessor(BaseImageProcessor):
- r"""
- Constructs a Deformable DETR image processor.
-
- Args:
- format (`str`, *optional*, defaults to `"coco_detection"`):
- Data format of the annotations. One of "coco_detection" or "coco_panoptic".
- do_resize (`bool`, *optional*, defaults to `True`):
- Controls whether to resize the image's (height, width) dimensions to the specified `size`. Can be
- overridden by the `do_resize` parameter in the `preprocess` method.
- size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 800, "longest_edge": 1333}`):
- Size of the image's (height, width) dimensions after resizing. Can be overridden by the `size` parameter in
- the `preprocess` method.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Controls whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the
- `do_rescale` parameter in the `preprocess` method.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the
- `preprocess` method.
- do_normalize:
- Controls whether to normalize the image. Can be overridden by the `do_normalize` parameter in the
- `preprocess` method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_DEFAULT_MEAN`):
- Mean values to use when normalizing the image. Can be a single value or a list of values, one for each
- channel. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_DEFAULT_STD`):
- Standard deviation values to use when normalizing the image. Can be a single value or a list of values, one
- for each channel. Can be overridden by the `image_std` parameter in the `preprocess` method.
- do_convert_annotations (`bool`, *optional*, defaults to `True`):
- Controls whether to convert the annotations to the format expected by the DETR model. Converts the
- bounding boxes to the format `(center_x, center_y, width, height)` and in the range `[0, 1]`.
- Can be overridden by the `do_convert_annotations` parameter in the `preprocess` method.
- do_pad (`bool`, *optional*, defaults to `True`):
- Controls whether to pad the image. Can be overridden by the `do_pad` parameter in the `preprocess`
- method. If `True` will pad the images in the batch to the largest height and width in the batch.
- Padding will be applied to the bottom and right of the image with zeros.
- """
-
- model_input_names = ["pixel_values", "pixel_mask"]
-
- def __init__(
- self,
- format: Union[str, AnnotationFormat] = AnnotationFormat.COCO_DETECTION,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- do_normalize: bool = True,
- image_mean: Union[float, List[float]] = None,
- image_std: Union[float, List[float]] = None,
- do_convert_annotations: bool = True,
- do_pad: bool = True,
- **kwargs,
- ) -> None:
- if "pad_and_return_pixel_mask" in kwargs:
- do_pad = kwargs.pop("pad_and_return_pixel_mask")
-
- size = size if size is not None else {"shortest_edge": 800, "longest_edge": 1333}
- size = get_size_dict(size, default_to_square=False)
-
- if do_convert_annotations is None:
- do_convert_annotations = do_normalize
-
- super().__init__(**kwargs)
- self.format = format
- self.do_resize = do_resize
- self.size = size
- self.resample = resample
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_normalize = do_normalize
- self.do_convert_annotations = do_convert_annotations
- self.image_mean = image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
- self.do_pad = do_pad
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_annotation with DETR->DETA
- def prepare_annotation(
- self,
- image: np.ndarray,
- target: Dict,
- format: Optional[AnnotationFormat] = None,
- return_segmentation_masks: bool = None,
- masks_path: Optional[Union[str, pathlib.Path]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> Dict:
- """
- Prepare an annotation for feeding into DETA model.
- """
- format = format if format is not None else self.format
-
- if format == AnnotationFormat.COCO_DETECTION:
- return_segmentation_masks = False if return_segmentation_masks is None else return_segmentation_masks
- target = prepare_coco_detection_annotation(
- image, target, return_segmentation_masks, input_data_format=input_data_format
- )
- elif format == AnnotationFormat.COCO_PANOPTIC:
- return_segmentation_masks = True if return_segmentation_masks is None else return_segmentation_masks
- target = prepare_coco_panoptic_annotation(
- image,
- target,
- masks_path=masks_path,
- return_masks=return_segmentation_masks,
- input_data_format=input_data_format,
- )
- else:
- raise ValueError(f"Format {format} is not supported.")
- return target
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare
- def prepare(self, image, target, return_segmentation_masks=None, masks_path=None):
- logger.warning_once(
- "The `prepare` method is deprecated and will be removed in a v4.33. "
- "Please use `prepare_annotation` instead. Note: the `prepare_annotation` method "
- "does not return the image anymore.",
- )
- target = self.prepare_annotation(image, target, return_segmentation_masks, masks_path, self.format)
- return image, target
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.convert_coco_poly_to_mask
- def convert_coco_poly_to_mask(self, *args, **kwargs):
- logger.warning_once("The `convert_coco_poly_to_mask` method is deprecated and will be removed in v4.33. ")
- return convert_coco_poly_to_mask(*args, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_coco_detection
- def prepare_coco_detection(self, *args, **kwargs):
- logger.warning_once("The `prepare_coco_detection` method is deprecated and will be removed in v4.33. ")
- return prepare_coco_detection_annotation(*args, **kwargs)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.prepare_coco_panoptic
- def prepare_coco_panoptic(self, *args, **kwargs):
- logger.warning_once("The `prepare_coco_panoptic` method is deprecated and will be removed in v4.33. ")
- return prepare_coco_panoptic_annotation(*args, **kwargs)
-
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize the image to the given size. Size can be `min_size` (scalar) or `(height, width)` tuple. If size is an
- int, smaller edge of the image will be matched to this number.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- The desired output size. Can contain keys `shortest_edge` and `longest_edge` or `height` and `width`.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image.
- data_format (`ChannelDimension`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred from the input
- image.
- """
- size = get_size_dict(size, default_to_square=False)
- if "shortest_edge" in size and "longest_edge" in size:
- size = get_resize_output_image_size(
- image, size["shortest_edge"], size["longest_edge"], input_data_format=input_data_format
- )
- elif "height" in size and "width" in size:
- size = (size["height"], size["width"])
- else:
- raise ValueError(
- "Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got"
- f" {size.keys()}."
- )
- image = resize(
- image, size=size, resample=resample, data_format=data_format, input_data_format=input_data_format
- )
- return image
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.resize_annotation
- def resize_annotation(
- self,
- annotation,
- orig_size,
- size,
- resample: PILImageResampling = PILImageResampling.NEAREST,
- ) -> Dict:
- """
- Resize the annotation to match the resized image. If size is an int, smaller edge of the mask will be matched
- to this number.
- """
- return resize_annotation(annotation, orig_size=orig_size, target_size=size, resample=resample)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.rescale
- def rescale(
- self,
- image: np.ndarray,
- rescale_factor: float,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.ndarray:
- """
- Rescale the image by the given factor. image = image * rescale_factor.
-
- Args:
- image (`np.ndarray`):
- Image to rescale.
- rescale_factor (`float`):
- The value to use for rescaling.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the input image. If unset, is inferred from the input image. Can be
- one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- """
- return rescale(image, rescale_factor, data_format=data_format, input_data_format=input_data_format)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.normalize_annotation
- def normalize_annotation(self, annotation: Dict, image_size: Tuple[int, int]) -> Dict:
- """
- Normalize the boxes in the annotation from `[top_left_x, top_left_y, bottom_right_x, bottom_right_y]` to
- `[center_x, center_y, width, height]` format and from absolute to relative pixel values.
- """
- return normalize_annotation(annotation, image_size=image_size)
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor._update_annotation_for_padded_image
- def _update_annotation_for_padded_image(
- self,
- annotation: Dict,
- input_image_size: Tuple[int, int],
- output_image_size: Tuple[int, int],
- padding,
- update_bboxes,
- ) -> Dict:
- """
- Update the annotation for a padded image.
- """
- new_annotation = {}
- new_annotation["size"] = output_image_size
-
- for key, value in annotation.items():
- if key == "masks":
- masks = value
- masks = pad(
- masks,
- padding,
- mode=PaddingMode.CONSTANT,
- constant_values=0,
- input_data_format=ChannelDimension.FIRST,
- )
- masks = safe_squeeze(masks, 1)
- new_annotation["masks"] = masks
- elif key == "boxes" and update_bboxes:
- boxes = value
- boxes *= np.asarray(
- [
- input_image_size[1] / output_image_size[1],
- input_image_size[0] / output_image_size[0],
- input_image_size[1] / output_image_size[1],
- input_image_size[0] / output_image_size[0],
- ]
- )
- new_annotation["boxes"] = boxes
- elif key == "size":
- new_annotation["size"] = output_image_size
- else:
- new_annotation[key] = value
- return new_annotation
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor._pad_image
- def _pad_image(
- self,
- image: np.ndarray,
- output_size: Tuple[int, int],
- annotation: Optional[Dict[str, Any]] = None,
- constant_values: Union[float, Iterable[float]] = 0,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- update_bboxes: bool = True,
- ) -> np.ndarray:
- """
- Pad an image with zeros to the given size.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- output_height, output_width = output_size
-
- pad_bottom = output_height - input_height
- pad_right = output_width - input_width
- padding = ((0, pad_bottom), (0, pad_right))
- padded_image = pad(
- image,
- padding,
- mode=PaddingMode.CONSTANT,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- )
- if annotation is not None:
- annotation = self._update_annotation_for_padded_image(
- annotation, (input_height, input_width), (output_height, output_width), padding, update_bboxes
- )
- return padded_image, annotation
-
- # Copied from transformers.models.detr.image_processing_detr.DetrImageProcessor.pad
- def pad(
- self,
- images: List[np.ndarray],
- annotations: Optional[Union[AnnotationType, List[AnnotationType]]] = None,
- constant_values: Union[float, Iterable[float]] = 0,
- return_pixel_mask: bool = True,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- update_bboxes: bool = True,
- ) -> BatchFeature:
- """
- Pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width
- in the batch and optionally returns their corresponding pixel mask.
-
- Args:
- images (List[`np.ndarray`]):
- Images to pad.
- annotations (`AnnotationType` or `List[AnnotationType]`, *optional*):
- Annotations to transform according to the padding that is applied to the images.
- constant_values (`float` or `Iterable[float]`, *optional*):
- The value to use for the padding if `mode` is `"constant"`.
- return_pixel_mask (`bool`, *optional*, defaults to `True`):
- Whether to return a pixel mask.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the image. If not provided, it will be the same as the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- update_bboxes (`bool`, *optional*, defaults to `True`):
- Whether to update the bounding boxes in the annotations to match the padded images. If the
- bounding boxes have not been converted to relative coordinates and `(centre_x, centre_y, width, height)`
- format, the bounding boxes will not be updated.
- """
- pad_size = get_max_height_width(images, input_data_format=input_data_format)
-
- annotation_list = annotations if annotations is not None else [None] * len(images)
- padded_images = []
- padded_annotations = []
- for image, annotation in zip(images, annotation_list):
- padded_image, padded_annotation = self._pad_image(
- image,
- pad_size,
- annotation,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- update_bboxes=update_bboxes,
- )
- padded_images.append(padded_image)
- padded_annotations.append(padded_annotation)
-
- data = {"pixel_values": padded_images}
-
- if return_pixel_mask:
- masks = [
- make_pixel_mask(image=image, output_size=pad_size, input_data_format=input_data_format)
- for image in images
- ]
- data["pixel_mask"] = masks
-
- encoded_inputs = BatchFeature(data=data, tensor_type=return_tensors)
-
- if annotations is not None:
- encoded_inputs["labels"] = [
- BatchFeature(annotation, tensor_type=return_tensors) for annotation in padded_annotations
- ]
-
- return encoded_inputs
-
- def preprocess(
- self,
- images: ImageInput,
- annotations: Optional[Union[List[Dict], List[List[Dict]]]] = None,
- return_segmentation_masks: bool = None,
- masks_path: Optional[Union[str, pathlib.Path]] = None,
- do_resize: Optional[bool] = None,
- size: Optional[Dict[str, int]] = None,
- resample=None, # PILImageResampling
- do_rescale: Optional[bool] = None,
- rescale_factor: Optional[Union[int, float]] = None,
- do_normalize: Optional[bool] = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- do_convert_annotations: Optional[bool] = None,
- do_pad: Optional[bool] = None,
- format: Optional[Union[str, AnnotationFormat]] = None,
- return_tensors: Optional[Union[TensorType, str]] = None,
- data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> BatchFeature:
- """
- Preprocess an image or a batch of images so that it can be used by the model.
-
- Args:
- images (`ImageInput`):
- Image or batch of images to preprocess. Expects a single or batch of images with pixel values ranging
- from 0 to 255. If passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- annotations (`List[Dict]` or `List[List[Dict]]`, *optional*):
- List of annotations associated with the image or batch of images. If annotation is for object
- detection, the annotations should be a dictionary with the following keys:
- - "image_id" (`int`): The image id.
- - "annotations" (`List[Dict]`): List of annotations for an image. Each annotation should be a
- dictionary. An image can have no annotations, in which case the list should be empty.
- If annotation is for segmentation, the annotations should be a dictionary with the following keys:
- - "image_id" (`int`): The image id.
- - "segments_info" (`List[Dict]`): List of segments for an image. Each segment should be a dictionary.
- An image can have no segments, in which case the list should be empty.
- - "file_name" (`str`): The file name of the image.
- return_segmentation_masks (`bool`, *optional*, defaults to self.return_segmentation_masks):
- Whether to return segmentation masks.
- masks_path (`str` or `pathlib.Path`, *optional*):
- Path to the directory containing the segmentation masks.
- do_resize (`bool`, *optional*, defaults to self.do_resize):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to self.size):
- Size of the image after resizing.
- resample (`PILImageResampling`, *optional*, defaults to self.resample):
- Resampling filter to use when resizing the image.
- do_rescale (`bool`, *optional*, defaults to self.do_rescale):
- Whether to rescale the image.
- rescale_factor (`float`, *optional*, defaults to self.rescale_factor):
- Rescale factor to use when rescaling the image.
- do_normalize (`bool`, *optional*, defaults to self.do_normalize):
- Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to self.image_mean):
- Mean to use when normalizing the image.
- image_std (`float` or `List[float]`, *optional*, defaults to self.image_std):
- Standard deviation to use when normalizing the image.
- do_convert_annotations (`bool`, *optional*, defaults to self.do_convert_annotations):
- Whether to convert the annotations to the format expected by the model. Converts the bounding
- boxes from the format `(top_left_x, top_left_y, width, height)` to `(center_x, center_y, width, height)`
- and in relative coordinates.
- do_pad (`bool`, *optional*, defaults to self.do_pad):
- Whether to pad the image. If `True` will pad the images in the batch to the largest image in the batch
- and create a pixel mask. Padding will be applied to the bottom and right of the image with zeros.
- format (`str` or `AnnotationFormat`, *optional*, defaults to self.format):
- Format of the annotations.
- return_tensors (`str` or `TensorType`, *optional*, defaults to self.return_tensors):
- Type of tensors to return. If `None`, will return the list of images.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - Unset: Use the channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- if "pad_and_return_pixel_mask" in kwargs:
- logger.warning_once(
- "The `pad_and_return_pixel_mask` argument is deprecated and will be removed in a future version, "
- "use `do_pad` instead.",
- )
- do_pad = kwargs.pop("pad_and_return_pixel_mask")
-
- do_resize = self.do_resize if do_resize is None else do_resize
- size = self.size if size is None else size
- size = get_size_dict(size=size, default_to_square=False)
- resample = self.resample if resample is None else resample
- do_rescale = self.do_rescale if do_rescale is None else do_rescale
- rescale_factor = self.rescale_factor if rescale_factor is None else rescale_factor
- do_normalize = self.do_normalize if do_normalize is None else do_normalize
- image_mean = self.image_mean if image_mean is None else image_mean
- image_std = self.image_std if image_std is None else image_std
- do_convert_annotations = (
- self.do_convert_annotations if do_convert_annotations is None else do_convert_annotations
- )
- do_pad = self.do_pad if do_pad is None else do_pad
- format = self.format if format is None else format
-
- # Here, the pad() method pads to the maximum of (width, height). It does not need to be validated.
-
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
-
- if not is_batched(images):
- images = [images]
- annotations = [annotations] if annotations is not None else None
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- if annotations is not None and len(images) != len(annotations):
- raise ValueError(
- f"The number of images ({len(images)}) and annotations ({len(annotations)}) do not match."
- )
-
- format = AnnotationFormat(format)
- if annotations is not None:
- validate_annotations(format, SUPPORTED_ANNOTATION_FORMATS, annotations)
-
- if (
- masks_path is not None
- and format == AnnotationFormat.COCO_PANOPTIC
- and not isinstance(masks_path, (pathlib.Path, str))
- ):
- raise ValueError(
- "The path to the directory containing the mask PNG files should be provided as a"
- f" `pathlib.Path` or string object, but is {type(masks_path)} instead."
- )
-
- # All transformations expect numpy arrays
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- # prepare (COCO annotations as a list of Dict -> DETR target as a single Dict per image)
- if annotations is not None:
- prepared_images = []
- prepared_annotations = []
- for image, target in zip(images, annotations):
- target = self.prepare_annotation(
- image,
- target,
- format,
- return_segmentation_masks=return_segmentation_masks,
- masks_path=masks_path,
- input_data_format=input_data_format,
- )
- prepared_images.append(image)
- prepared_annotations.append(target)
- images = prepared_images
- annotations = prepared_annotations
- del prepared_images, prepared_annotations
-
- # transformations
- if do_resize:
- if annotations is not None:
- resized_images, resized_annotations = [], []
- for image, target in zip(images, annotations):
- orig_size = get_image_size(image, input_data_format)
- resized_image = self.resize(
- image, size=size, resample=resample, input_data_format=input_data_format
- )
- resized_annotation = self.resize_annotation(
- target, orig_size, get_image_size(resized_image, input_data_format)
- )
- resized_images.append(resized_image)
- resized_annotations.append(resized_annotation)
- images = resized_images
- annotations = resized_annotations
- del resized_images, resized_annotations
- else:
- images = [
- self.resize(image, size=size, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_rescale:
- images = [self.rescale(image, rescale_factor, input_data_format=input_data_format) for image in images]
-
- if do_normalize:
- images = [
- self.normalize(image, image_mean, image_std, input_data_format=input_data_format) for image in images
- ]
-
- if do_convert_annotations and annotations is not None:
- annotations = [
- self.normalize_annotation(annotation, get_image_size(image, input_data_format))
- for annotation, image in zip(annotations, images)
- ]
-
- if do_pad:
- # Pads images and returns their mask: {'pixel_values': ..., 'pixel_mask': ...}
- encoded_inputs = self.pad(
- images,
- annotations=annotations,
- return_pixel_mask=True,
- data_format=data_format,
- input_data_format=input_data_format,
- return_tensors=return_tensors,
- update_bboxes=do_convert_annotations,
- )
- else:
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
- for image in images
- ]
- encoded_inputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)
- if annotations is not None:
- encoded_inputs["labels"] = [
- BatchFeature(annotation, tensor_type=return_tensors) for annotation in annotations
- ]
-
- return encoded_inputs
-
- def post_process_object_detection(
- self,
- outputs,
- threshold: float = 0.5,
- target_sizes: Union[TensorType, List[Tuple]] = None,
- nms_threshold: float = 0.7,
- ):
- """
- Converts the output of [`DetaForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y,
- bottom_right_x, bottom_right_y) format. Only supports PyTorch.
-
- Args:
- outputs ([`DetrObjectDetectionOutput`]):
- Raw outputs of the model.
- threshold (`float`, *optional*, defaults to 0.5):
- Score threshold to keep object detection predictions.
- target_sizes (`torch.Tensor` or `List[Tuple[int, int]]`, *optional*):
- Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
- (height, width) of each image in the batch. If left to None, predictions will not be resized.
- nms_threshold (`float`, *optional*, defaults to 0.7):
- NMS threshold.
-
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
- in the batch as predicted by the model.
- """
- out_logits, out_bbox = outputs.logits, outputs.pred_boxes
- batch_size, num_queries, num_labels = out_logits.shape
-
- if target_sizes is not None:
- if len(out_logits) != len(target_sizes):
- raise ValueError(
- "Make sure that you pass in as many target sizes as the batch dimension of the logits"
- )
-
- prob = out_logits.sigmoid()
-
- all_scores = prob.view(batch_size, num_queries * num_labels).to(out_logits.device)
- all_indexes = torch.arange(num_queries * num_labels)[None].repeat(batch_size, 1).to(out_logits.device)
- all_boxes = torch.div(all_indexes, out_logits.shape[2], rounding_mode="floor")
- all_labels = all_indexes % out_logits.shape[2]
-
- boxes = center_to_corners_format(out_bbox)
- boxes = torch.gather(boxes, 1, all_boxes.unsqueeze(-1).repeat(1, 1, 4))
-
- # and from relative [0, 1] to absolute [0, height] coordinates
- if target_sizes is not None:
- if isinstance(target_sizes, List):
- img_h = torch.Tensor([i[0] for i in target_sizes])
- img_w = torch.Tensor([i[1] for i in target_sizes])
- else:
- img_h, img_w = target_sizes.unbind(1)
-
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device)
- boxes = boxes * scale_fct[:, None, :]
-
- results = []
- for b in range(batch_size):
- box = boxes[b]
- score = all_scores[b]
- lbls = all_labels[b]
-
- pre_topk = score.topk(min(10000, num_queries * num_labels)).indices
- box = box[pre_topk]
- score = score[pre_topk]
- lbls = lbls[pre_topk]
-
- # apply NMS
- keep_inds = batched_nms(box, score, lbls, nms_threshold)[:100]
- score = score[keep_inds]
- lbls = lbls[keep_inds]
- box = box[keep_inds]
-
- results.append(
- {
- "scores": score[score > threshold],
- "labels": lbls[score > threshold],
- "boxes": box[score > threshold],
- }
- )
-
- return results
diff --git a/transformers/models/deta/modeling_deta.py b/transformers/models/deta/modeling_deta.py
deleted file mode 100644
index ce0a5e79aa4eb1df665c0c54cd32792bdea93786..0000000000000000000000000000000000000000
--- a/transformers/models/deta/modeling_deta.py
+++ /dev/null
@@ -1,2877 +0,0 @@
-# coding=utf-8
-# Copyright 2022 SenseTime and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DETA model."""
-
-
-import copy
-import math
-import os
-import warnings
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor, nn
-from torch.autograd import Function
-from torch.autograd.function import once_differentiable
-
-from ...activations import ACT2FN
-from ...file_utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_scipy_available,
- is_torch_cuda_available,
- is_vision_available,
- replace_return_docstrings,
-)
-from ...modeling_attn_mask_utils import _prepare_4d_attention_mask
-from ...modeling_outputs import BaseModelOutput
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import meshgrid
-from ...utils import is_accelerate_available, is_ninja_available, is_torchvision_available, logging, requires_backends
-from ...utils.backbone_utils import load_backbone
-from .configuration_deta import DetaConfig
-
-
-logger = logging.get_logger(__name__)
-
-MultiScaleDeformableAttention = None
-
-
-# Copied from models.deformable_detr.load_cuda_kernels
-def load_cuda_kernels():
- from torch.utils.cpp_extension import load
-
- global MultiScaleDeformableAttention
-
- root = Path(__file__).resolve().parent.parent.parent / "kernels" / "deta"
- src_files = [
- root / filename
- for filename in [
- "vision.cpp",
- os.path.join("cpu", "ms_deform_attn_cpu.cpp"),
- os.path.join("cuda", "ms_deform_attn_cuda.cu"),
- ]
- ]
-
- load(
- "MultiScaleDeformableAttention",
- src_files,
- with_cuda=True,
- extra_include_paths=[str(root)],
- extra_cflags=["-DWITH_CUDA=1"],
- extra_cuda_cflags=[
- "-DCUDA_HAS_FP16=1",
- "-D__CUDA_NO_HALF_OPERATORS__",
- "-D__CUDA_NO_HALF_CONVERSIONS__",
- "-D__CUDA_NO_HALF2_OPERATORS__",
- ],
- )
-
-
-# Copied from transformers.models.deformable_detr.modeling_deformable_detr.MultiScaleDeformableAttentionFunction
-class MultiScaleDeformableAttentionFunction(Function):
- @staticmethod
- def forward(
- context,
- value,
- value_spatial_shapes,
- value_level_start_index,
- sampling_locations,
- attention_weights,
- im2col_step,
- ):
- context.im2col_step = im2col_step
- output = MultiScaleDeformableAttention.ms_deform_attn_forward(
- value,
- value_spatial_shapes,
- value_level_start_index,
- sampling_locations,
- attention_weights,
- context.im2col_step,
- )
- context.save_for_backward(
- value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights
- )
- return output
-
- @staticmethod
- @once_differentiable
- def backward(context, grad_output):
- (
- value,
- value_spatial_shapes,
- value_level_start_index,
- sampling_locations,
- attention_weights,
- ) = context.saved_tensors
- grad_value, grad_sampling_loc, grad_attn_weight = MultiScaleDeformableAttention.ms_deform_attn_backward(
- value,
- value_spatial_shapes,
- value_level_start_index,
- sampling_locations,
- attention_weights,
- grad_output,
- context.im2col_step,
- )
-
- return grad_value, None, None, grad_sampling_loc, grad_attn_weight, None
-
-
-if is_accelerate_available():
- from accelerate import PartialState
- from accelerate.utils import reduce
-
-if is_vision_available():
- from transformers.image_transforms import center_to_corners_format
-
-if is_torchvision_available():
- from torchvision.ops.boxes import batched_nms
-
-if is_scipy_available():
- from scipy.optimize import linear_sum_assignment
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "DetaConfig"
-_CHECKPOINT_FOR_DOC = "jozhang97/deta-swin-large-o365"
-
-
-from ..deprecated._archive_maps import DETA_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-# Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrDecoderOutput with DeformableDetr->Deta
-class DetaDecoderOutput(ModelOutput):
- """
- Base class for outputs of the DetaDecoder. This class adds two attributes to
- BaseModelOutputWithCrossAttentions, namely:
- - a stacked tensor of intermediate decoder hidden states (i.e. the output of each decoder layer)
- - a stacked tensor of intermediate reference points.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
- Stacked intermediate hidden states (output of each layer of the decoder).
- intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, hidden_size)`):
- Stacked intermediate reference points (reference points of each layer of the decoder).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
- the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- """
-
- last_hidden_state: torch.FloatTensor = None
- intermediate_hidden_states: torch.FloatTensor = None
- intermediate_reference_points: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class DetaModelOutput(ModelOutput):
- """
- Base class for outputs of the Deformable DETR encoder-decoder model.
-
- Args:
- init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Initial reference points sent through the Transformer decoder.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
- Stacked intermediate hidden states (output of each layer of the decoder).
- intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
- Stacked intermediate reference points (reference points of each layer of the decoder).
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, num_queries, hidden_size)`. Hidden-states of the decoder at the output of each layer
- plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, num_queries,
- num_queries)`. Attentions weights of the decoder, after the attention softmax, used to compute the weighted
- average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_queries, num_heads, 4, 4)`.
- Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
- weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_queries, num_heads, 4, 4)`.
- Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
- self-attention heads.
- enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
- Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
- picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
- foreground and background).
- enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
- Logits of predicted bounding boxes coordinates in the first stage.
- output_proposals (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.two_stage=True`):
- Logits of proposal bounding boxes coordinates in the gen_encoder_output_proposals.
- """
-
- init_reference_points: torch.FloatTensor = None
- last_hidden_state: torch.FloatTensor = None
- intermediate_hidden_states: torch.FloatTensor = None
- intermediate_reference_points: torch.FloatTensor = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- enc_outputs_class: Optional[torch.FloatTensor] = None
- enc_outputs_coord_logits: Optional[torch.FloatTensor] = None
- output_proposals: Optional[torch.FloatTensor] = None
-
-
-@dataclass
-class DetaObjectDetectionOutput(ModelOutput):
- """
- Output type of [`DetaForObjectDetection`].
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
- Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
- bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
- scale-invariant IoU loss.
- loss_dict (`Dict`, *optional*):
- A dictionary containing the individual losses. Useful for logging.
- logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
- Classification logits (including no-object) for all queries.
- pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
- values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
- possible padding). You can use [`~DetaProcessor.post_process_object_detection`] to retrieve the
- unnormalized bounding boxes.
- auxiliary_outputs (`list[Dict]`, *optional*):
- Optional, only returned when auxilary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
- and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
- `pred_boxes`) for each decoder layer.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, num_queries, hidden_size)`. Hidden-states of the decoder at the output of each layer
- plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, num_queries,
- num_queries)`. Attentions weights of the decoder, after the attention softmax, used to compute the weighted
- average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_queries, num_heads, 4, 4)`.
- Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
- weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_heads, 4,
- 4)`. Attentions weights of the encoder, after the attention softmax, used to compute the weighted average
- in the self-attention heads.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
- Stacked intermediate hidden states (output of each layer of the decoder).
- intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
- Stacked intermediate reference points (reference points of each layer of the decoder).
- init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Initial reference points sent through the Transformer decoder.
- enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
- Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
- picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
- foreground and background).
- enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
- Logits of predicted bounding boxes coordinates in the first stage.
- output_proposals (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.two_stage=True`):
- Logits of proposal bounding boxes coordinates in the gen_encoder_output_proposals.
- """
-
- loss: Optional[torch.FloatTensor] = None
- loss_dict: Optional[Dict] = None
- logits: torch.FloatTensor = None
- pred_boxes: torch.FloatTensor = None
- auxiliary_outputs: Optional[List[Dict]] = None
- init_reference_points: Optional[torch.FloatTensor] = None
- last_hidden_state: Optional[torch.FloatTensor] = None
- intermediate_hidden_states: Optional[torch.FloatTensor] = None
- intermediate_reference_points: Optional[torch.FloatTensor] = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- enc_outputs_class: Optional = None
- enc_outputs_coord_logits: Optional = None
- output_proposals: Optional[torch.FloatTensor] = None
-
-
-def _get_clones(module, N):
- return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
-
-
-def inverse_sigmoid(x, eps=1e-5):
- x = x.clamp(min=0, max=1)
- x1 = x.clamp(min=eps)
- x2 = (1 - x).clamp(min=eps)
- return torch.log(x1 / x2)
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrFrozenBatchNorm2d with Detr->Deta
-class DetaFrozenBatchNorm2d(nn.Module):
- """
- BatchNorm2d where the batch statistics and the affine parameters are fixed.
-
- Copy-paste from torchvision.misc.ops with added eps before rqsrt, without which any other models than
- torchvision.models.resnet[18,34,50,101] produce nans.
- """
-
- def __init__(self, n):
- super().__init__()
- self.register_buffer("weight", torch.ones(n))
- self.register_buffer("bias", torch.zeros(n))
- self.register_buffer("running_mean", torch.zeros(n))
- self.register_buffer("running_var", torch.ones(n))
-
- def _load_from_state_dict(
- self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
- ):
- num_batches_tracked_key = prefix + "num_batches_tracked"
- if num_batches_tracked_key in state_dict:
- del state_dict[num_batches_tracked_key]
-
- super()._load_from_state_dict(
- state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
- )
-
- def forward(self, x):
- # move reshapes to the beginning
- # to make it user-friendly
- weight = self.weight.reshape(1, -1, 1, 1)
- bias = self.bias.reshape(1, -1, 1, 1)
- running_var = self.running_var.reshape(1, -1, 1, 1)
- running_mean = self.running_mean.reshape(1, -1, 1, 1)
- epsilon = 1e-5
- scale = weight * (running_var + epsilon).rsqrt()
- bias = bias - running_mean * scale
- return x * scale + bias
-
-
-# Copied from transformers.models.detr.modeling_detr.replace_batch_norm with Detr->Deta
-def replace_batch_norm(model):
- r"""
- Recursively replace all `torch.nn.BatchNorm2d` with `DetaFrozenBatchNorm2d`.
-
- Args:
- model (torch.nn.Module):
- input model
- """
- for name, module in model.named_children():
- if isinstance(module, nn.BatchNorm2d):
- new_module = DetaFrozenBatchNorm2d(module.num_features)
-
- if not module.weight.device == torch.device("meta"):
- new_module.weight.data.copy_(module.weight)
- new_module.bias.data.copy_(module.bias)
- new_module.running_mean.data.copy_(module.running_mean)
- new_module.running_var.data.copy_(module.running_var)
-
- model._modules[name] = new_module
-
- if len(list(module.children())) > 0:
- replace_batch_norm(module)
-
-
-class DetaBackboneWithPositionalEncodings(nn.Module):
- """
- Backbone model with positional embeddings.
-
- nn.BatchNorm2d layers are replaced by DetaFrozenBatchNorm2d as defined above.
- """
-
- def __init__(self, config):
- super().__init__()
-
- backbone = load_backbone(config)
- with torch.no_grad():
- replace_batch_norm(backbone)
- self.model = backbone
- self.intermediate_channel_sizes = self.model.channels
-
- # TODO fix this
- if config.backbone_config.model_type == "resnet":
- for name, parameter in self.model.named_parameters():
- if "stages.1" not in name and "stages.2" not in name and "stages.3" not in name:
- parameter.requires_grad_(False)
-
- self.position_embedding = build_position_encoding(config)
-
- def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
- """
- Outputs feature maps of latter stages C_3 through C_5 in ResNet if `config.num_feature_levels > 1`, otherwise
- outputs feature maps of C_5.
- """
- # first, send pixel_values through the backbone to get list of feature maps
- features = self.model(pixel_values).feature_maps
-
- # next, create position embeddings
- out = []
- pos = []
- for feature_map in features:
- # downsample pixel_mask to match shape of corresponding feature_map
- mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
- position_embeddings = self.position_embedding(feature_map, mask).to(feature_map.dtype)
- out.append((feature_map, mask))
- pos.append(position_embeddings)
-
- return out, pos
-
-
-# Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrSinePositionEmbedding with DeformableDetr->Deta
-class DetaSinePositionEmbedding(nn.Module):
- """
- This is a more standard version of the position embedding, very similar to the one used by the Attention is all you
- need paper, generalized to work on images.
- """
-
- def __init__(self, embedding_dim=64, temperature=10000, normalize=False, scale=None):
- super().__init__()
- self.embedding_dim = embedding_dim
- self.temperature = temperature
- self.normalize = normalize
- if scale is not None and normalize is False:
- raise ValueError("normalize should be True if scale is passed")
- if scale is None:
- scale = 2 * math.pi
- self.scale = scale
-
- def forward(self, pixel_values, pixel_mask):
- if pixel_mask is None:
- raise ValueError("No pixel mask provided")
- y_embed = pixel_mask.cumsum(1, dtype=torch.float32)
- x_embed = pixel_mask.cumsum(2, dtype=torch.float32)
- if self.normalize:
- eps = 1e-6
- y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale
- x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale
-
- dim_t = torch.arange(self.embedding_dim, dtype=torch.int64, device=pixel_values.device).float()
- dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.embedding_dim)
-
- pos_x = x_embed[:, :, :, None] / dim_t
- pos_y = y_embed[:, :, :, None] / dim_t
- pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
- return pos
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrLearnedPositionEmbedding
-class DetaLearnedPositionEmbedding(nn.Module):
- """
- This module learns positional embeddings up to a fixed maximum size.
- """
-
- def __init__(self, embedding_dim=256):
- super().__init__()
- self.row_embeddings = nn.Embedding(50, embedding_dim)
- self.column_embeddings = nn.Embedding(50, embedding_dim)
-
- def forward(self, pixel_values, pixel_mask=None):
- height, width = pixel_values.shape[-2:]
- width_values = torch.arange(width, device=pixel_values.device)
- height_values = torch.arange(height, device=pixel_values.device)
- x_emb = self.column_embeddings(width_values)
- y_emb = self.row_embeddings(height_values)
- pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1)
- pos = pos.permute(2, 0, 1)
- pos = pos.unsqueeze(0)
- pos = pos.repeat(pixel_values.shape[0], 1, 1, 1)
- return pos
-
-
-# Copied from transformers.models.detr.modeling_detr.build_position_encoding with Detr->Deta
-def build_position_encoding(config):
- n_steps = config.d_model // 2
- if config.position_embedding_type == "sine":
- # TODO find a better way of exposing other arguments
- position_embedding = DetaSinePositionEmbedding(n_steps, normalize=True)
- elif config.position_embedding_type == "learned":
- position_embedding = DetaLearnedPositionEmbedding(n_steps)
- else:
- raise ValueError(f"Not supported {config.position_embedding_type}")
-
- return position_embedding
-
-
-# Copied from transformers.models.deformable_detr.modeling_deformable_detr.multi_scale_deformable_attention
-def multi_scale_deformable_attention(
- value: Tensor, value_spatial_shapes: Tensor, sampling_locations: Tensor, attention_weights: Tensor
-) -> Tensor:
- batch_size, _, num_heads, hidden_dim = value.shape
- _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
- value_list = value.split([height.item() * width.item() for height, width in value_spatial_shapes], dim=1)
- sampling_grids = 2 * sampling_locations - 1
- sampling_value_list = []
- for level_id, (height, width) in enumerate(value_spatial_shapes):
- # batch_size, height*width, num_heads, hidden_dim
- # -> batch_size, height*width, num_heads*hidden_dim
- # -> batch_size, num_heads*hidden_dim, height*width
- # -> batch_size*num_heads, hidden_dim, height, width
- value_l_ = (
- value_list[level_id].flatten(2).transpose(1, 2).reshape(batch_size * num_heads, hidden_dim, height, width)
- )
- # batch_size, num_queries, num_heads, num_points, 2
- # -> batch_size, num_heads, num_queries, num_points, 2
- # -> batch_size*num_heads, num_queries, num_points, 2
- sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1)
- # batch_size*num_heads, hidden_dim, num_queries, num_points
- sampling_value_l_ = nn.functional.grid_sample(
- value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False
- )
- sampling_value_list.append(sampling_value_l_)
- # (batch_size, num_queries, num_heads, num_levels, num_points)
- # -> (batch_size, num_heads, num_queries, num_levels, num_points)
- # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points)
- attention_weights = attention_weights.transpose(1, 2).reshape(
- batch_size * num_heads, 1, num_queries, num_levels * num_points
- )
- output = (
- (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights)
- .sum(-1)
- .view(batch_size, num_heads * hidden_dim, num_queries)
- )
- return output.transpose(1, 2).contiguous()
-
-
-# Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrMultiscaleDeformableAttention with DeformableDetr->Deta
-class DetaMultiscaleDeformableAttention(nn.Module):
- """
- Multiscale deformable attention as proposed in Deformable DETR.
- """
-
- def __init__(self, config: DetaConfig, num_heads: int, n_points: int):
- super().__init__()
-
- kernel_loaded = MultiScaleDeformableAttention is not None
- if is_torch_cuda_available() and is_ninja_available() and not kernel_loaded:
- try:
- load_cuda_kernels()
- except Exception as e:
- logger.warning(f"Could not load the custom kernel for multi-scale deformable attention: {e}")
-
- if config.d_model % num_heads != 0:
- raise ValueError(
- f"embed_dim (d_model) must be divisible by num_heads, but got {config.d_model} and {num_heads}"
- )
- dim_per_head = config.d_model // num_heads
- # check if dim_per_head is power of 2
- if not ((dim_per_head & (dim_per_head - 1) == 0) and dim_per_head != 0):
- warnings.warn(
- "You'd better set embed_dim (d_model) in DetaMultiscaleDeformableAttention to make the"
- " dimension of each attention head a power of 2 which is more efficient in the authors' CUDA"
- " implementation."
- )
-
- self.im2col_step = 64
-
- self.d_model = config.d_model
- self.n_levels = config.num_feature_levels
- self.n_heads = num_heads
- self.n_points = n_points
-
- self.sampling_offsets = nn.Linear(config.d_model, num_heads * self.n_levels * n_points * 2)
- self.attention_weights = nn.Linear(config.d_model, num_heads * self.n_levels * n_points)
- self.value_proj = nn.Linear(config.d_model, config.d_model)
- self.output_proj = nn.Linear(config.d_model, config.d_model)
-
- self.disable_custom_kernels = config.disable_custom_kernels
-
- self._reset_parameters()
-
- def _reset_parameters(self):
- nn.init.constant_(self.sampling_offsets.weight.data, 0.0)
- default_dtype = torch.get_default_dtype()
- thetas = torch.arange(self.n_heads, dtype=torch.int64).to(default_dtype) * (2.0 * math.pi / self.n_heads)
- grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
- grid_init = (
- (grid_init / grid_init.abs().max(-1, keepdim=True)[0])
- .view(self.n_heads, 1, 1, 2)
- .repeat(1, self.n_levels, self.n_points, 1)
- )
- for i in range(self.n_points):
- grid_init[:, :, i, :] *= i + 1
- with torch.no_grad():
- self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
- nn.init.constant_(self.attention_weights.weight.data, 0.0)
- nn.init.constant_(self.attention_weights.bias.data, 0.0)
- nn.init.xavier_uniform_(self.value_proj.weight.data)
- nn.init.constant_(self.value_proj.bias.data, 0.0)
- nn.init.xavier_uniform_(self.output_proj.weight.data)
- nn.init.constant_(self.output_proj.bias.data, 0.0)
-
- def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Optional[Tensor]):
- return tensor if position_embeddings is None else tensor + position_embeddings
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- position_embeddings: Optional[torch.Tensor] = None,
- reference_points=None,
- spatial_shapes=None,
- level_start_index=None,
- output_attentions: bool = False,
- ):
- # add position embeddings to the hidden states before projecting to queries and keys
- if position_embeddings is not None:
- hidden_states = self.with_pos_embed(hidden_states, position_embeddings)
-
- batch_size, num_queries, _ = hidden_states.shape
- batch_size, sequence_length, _ = encoder_hidden_states.shape
- if (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() != sequence_length:
- raise ValueError(
- "Make sure to align the spatial shapes with the sequence length of the encoder hidden states"
- )
-
- value = self.value_proj(encoder_hidden_states)
- if attention_mask is not None:
- # we invert the attention_mask
- value = value.masked_fill(~attention_mask[..., None], float(0))
- value = value.view(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads)
- sampling_offsets = self.sampling_offsets(hidden_states).view(
- batch_size, num_queries, self.n_heads, self.n_levels, self.n_points, 2
- )
- attention_weights = self.attention_weights(hidden_states).view(
- batch_size, num_queries, self.n_heads, self.n_levels * self.n_points
- )
- attention_weights = F.softmax(attention_weights, -1).view(
- batch_size, num_queries, self.n_heads, self.n_levels, self.n_points
- )
- # batch_size, num_queries, n_heads, n_levels, n_points, 2
- num_coordinates = reference_points.shape[-1]
- if num_coordinates == 2:
- offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1)
- sampling_locations = (
- reference_points[:, :, None, :, None, :]
- + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
- )
- elif num_coordinates == 4:
- sampling_locations = (
- reference_points[:, :, None, :, None, :2]
- + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
- )
- else:
- raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}")
-
- if self.disable_custom_kernels:
- # PyTorch implementation
- output = multi_scale_deformable_attention(value, spatial_shapes, sampling_locations, attention_weights)
- else:
- try:
- # custom kernel
- output = MultiScaleDeformableAttentionFunction.apply(
- value,
- spatial_shapes,
- level_start_index,
- sampling_locations,
- attention_weights,
- self.im2col_step,
- )
- except Exception:
- # PyTorch implementation
- output = multi_scale_deformable_attention(value, spatial_shapes, sampling_locations, attention_weights)
- output = self.output_proj(output)
-
- return output, attention_weights
-
-
-# Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrMultiheadAttention with DeformableDetr->Deta,Deformable DETR->DETA
-class DetaMultiheadAttention(nn.Module):
- """
- Multi-headed attention from 'Attention Is All You Need' paper.
-
- Here, we add position embeddings to the queries and keys (as explained in the Deformable DETR paper).
- """
-
- def __init__(
- self,
- embed_dim: int,
- num_heads: int,
- dropout: float = 0.0,
- bias: bool = True,
- ):
- super().__init__()
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.dropout = dropout
- self.head_dim = embed_dim // num_heads
- if self.head_dim * num_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {num_heads})."
- )
- self.scaling = self.head_dim**-0.5
-
- self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
-
- def _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int):
- return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def with_pos_embed(self, tensor: torch.Tensor, position_embeddings: Optional[Tensor]):
- return tensor if position_embeddings is None else tensor + position_embeddings
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_embeddings: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel"""
-
- batch_size, target_len, embed_dim = hidden_states.size()
- # add position embeddings to the hidden states before projecting to queries and keys
- if position_embeddings is not None:
- hidden_states_original = hidden_states
- hidden_states = self.with_pos_embed(hidden_states, position_embeddings)
-
- # get queries, keys and values
- query_states = self.q_proj(hidden_states) * self.scaling
- key_states = self._shape(self.k_proj(hidden_states), -1, batch_size)
- value_states = self._shape(self.v_proj(hidden_states_original), -1, batch_size)
-
- proj_shape = (batch_size * self.num_heads, -1, self.head_dim)
- query_states = self._shape(query_states, target_len, batch_size).view(*proj_shape)
- key_states = key_states.view(*proj_shape)
- value_states = value_states.view(*proj_shape)
-
- source_len = key_states.size(1)
-
- attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-
- if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len):
- raise ValueError(
- f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- # expand attention_mask
- if attention_mask is not None:
- # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len]
- attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
-
- if attention_mask is not None:
- if attention_mask.size() != (batch_size, 1, target_len, source_len):
- raise ValueError(
- f"Attention mask should be of size {(batch_size, 1, target_len, source_len)}, but is"
- f" {attention_mask.size()}"
- )
- attn_weights = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attention_mask
- attn_weights = attn_weights.view(batch_size * self.num_heads, target_len, source_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if output_attentions:
- # this operation is a bit awkward, but it's required to
- # make sure that attn_weights keeps its gradient.
- # In order to do so, attn_weights have to reshaped
- # twice and have to be reused in the following
- attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len)
- attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-
- attn_output = torch.bmm(attn_probs, value_states)
-
- if attn_output.size() != (batch_size * self.num_heads, target_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.head_dim)
- attn_output = attn_output.transpose(1, 2)
- attn_output = attn_output.reshape(batch_size, target_len, embed_dim)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped
-
-
-class DetaEncoderLayer(nn.Module):
- def __init__(self, config: DetaConfig):
- super().__init__()
- self.embed_dim = config.d_model
- self.self_attn = DetaMultiscaleDeformableAttention(
- config,
- num_heads=config.encoder_attention_heads,
- n_points=config.encoder_n_points,
- )
- self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
- self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
- self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: torch.Tensor,
- position_embeddings: torch.Tensor = None,
- reference_points=None,
- spatial_shapes=None,
- level_start_index=None,
- output_attentions: bool = False,
- ):
- """
- Args:
- hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Input to the layer.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
- Attention mask.
- position_embeddings (`torch.FloatTensor`, *optional*):
- Position embeddings, to be added to `hidden_states`.
- reference_points (`torch.FloatTensor`, *optional*):
- Reference points.
- spatial_shapes (`torch.LongTensor`, *optional*):
- Spatial shapes of the backbone feature maps.
- level_start_index (`torch.LongTensor`, *optional*):
- Level start index.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- residual = hidden_states
-
- # Apply Multi-scale Deformable Attention Module on the multi-scale feature maps.
- hidden_states, attn_weights = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- encoder_hidden_states=hidden_states,
- encoder_attention_mask=attention_mask,
- position_embeddings=position_embeddings,
- reference_points=reference_points,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- residual = hidden_states
- hidden_states = self.activation_fn(self.fc1(hidden_states))
- hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
-
- hidden_states = self.fc2(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
- hidden_states = residual + hidden_states
- hidden_states = self.final_layer_norm(hidden_states)
-
- if self.training:
- if torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any():
- clamp_value = torch.finfo(hidden_states.dtype).max - 1000
- hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
-
-class DetaDecoderLayer(nn.Module):
- def __init__(self, config: DetaConfig):
- super().__init__()
- self.embed_dim = config.d_model
-
- # self-attention
- self.self_attn = DetaMultiheadAttention(
- embed_dim=self.embed_dim,
- num_heads=config.decoder_attention_heads,
- dropout=config.attention_dropout,
- )
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
-
- self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- # cross-attention
- self.encoder_attn = DetaMultiscaleDeformableAttention(
- config,
- num_heads=config.decoder_attention_heads,
- n_points=config.decoder_n_points,
- )
- self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- # feedforward neural networks
- self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
- self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- position_embeddings: Optional[torch.Tensor] = None,
- reference_points=None,
- spatial_shapes=None,
- level_start_index=None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- ):
- """
- Args:
- hidden_states (`torch.FloatTensor`):
- Input to the layer of shape `(batch, seq_len, embed_dim)`.
- position_embeddings (`torch.FloatTensor`, *optional*):
- Position embeddings that are added to the queries and keys in the self-attention layer.
- reference_points (`torch.FloatTensor`, *optional*):
- Reference points.
- spatial_shapes (`torch.LongTensor`, *optional*):
- Spatial shapes.
- level_start_index (`torch.LongTensor`, *optional*):
- Level start index.
- encoder_hidden_states (`torch.FloatTensor`):
- cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
- encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
- `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
- values.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- residual = hidden_states
-
- # Self Attention
- hidden_states, self_attn_weights = self.self_attn(
- hidden_states=hidden_states,
- position_embeddings=position_embeddings,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- second_residual = hidden_states
-
- # Cross-Attention
- cross_attn_weights = None
- hidden_states, cross_attn_weights = self.encoder_attn(
- hidden_states=hidden_states,
- attention_mask=encoder_attention_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- position_embeddings=position_embeddings,
- reference_points=reference_points,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = second_residual + hidden_states
-
- hidden_states = self.encoder_attn_layer_norm(hidden_states)
-
- # Fully Connected
- residual = hidden_states
- hidden_states = self.activation_fn(self.fc1(hidden_states))
- hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
- hidden_states = self.fc2(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.final_layer_norm(hidden_states)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights, cross_attn_weights)
-
- return outputs
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrClassificationHead
-class DetaClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, input_dim: int, inner_dim: int, num_classes: int, pooler_dropout: float):
- super().__init__()
- self.dense = nn.Linear(input_dim, inner_dim)
- self.dropout = nn.Dropout(p=pooler_dropout)
- self.out_proj = nn.Linear(inner_dim, num_classes)
-
- def forward(self, hidden_states: torch.Tensor):
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.dense(hidden_states)
- hidden_states = torch.tanh(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.out_proj(hidden_states)
- return hidden_states
-
-
-class DetaPreTrainedModel(PreTrainedModel):
- config_class = DetaConfig
- base_model_prefix = "model"
- main_input_name = "pixel_values"
- _no_split_modules = [r"DetaBackboneWithPositionalEncodings", r"DetaEncoderLayer", r"DetaDecoderLayer"]
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- std = self.config.init_std
-
- if isinstance(module, DetaLearnedPositionEmbedding):
- nn.init.uniform_(module.row_embeddings.weight)
- nn.init.uniform_(module.column_embeddings.weight)
- elif isinstance(module, DetaMultiscaleDeformableAttention):
- module._reset_parameters()
- elif isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- if hasattr(module, "reference_points") and not self.config.two_stage:
- nn.init.xavier_uniform_(module.reference_points.weight.data, gain=1.0)
- nn.init.constant_(module.reference_points.bias.data, 0.0)
- if hasattr(module, "level_embed"):
- nn.init.normal_(module.level_embed)
-
-
-DETA_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`DetaConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DETA_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Padding will be ignored by default should you provide it.
-
- Pixel values can be obtained using [`AutoImageProcessor`]. See [`AutoImageProcessor.__call__`] for details.
-
- pixel_mask (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
- Mask to avoid performing attention on padding pixel values. Mask values selected in `[0, 1]`:
-
- - 1 for pixels that are real (i.e. **not masked**),
- - 0 for pixels that are padding (i.e. **masked**).
-
- [What are attention masks?](../glossary#attention-mask)
-
- decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*):
- Not used by default. Can be used to mask object queries.
- encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*):
- Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
- `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of
- hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you
- can choose to directly pass a flattened representation of an image.
- decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
- Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an
- embedded representation.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class DetaEncoder(DetaPreTrainedModel):
- """
- Transformer encoder consisting of *config.encoder_layers* deformable attention layers. Each layer is a
- [`DetaEncoderLayer`].
-
- The encoder updates the flattened multi-scale feature maps through multiple deformable attention layers.
-
- Args:
- config: DetaConfig
- """
-
- def __init__(self, config: DetaConfig):
- super().__init__(config)
-
- self.dropout = config.dropout
- self.layers = nn.ModuleList([DetaEncoderLayer(config) for _ in range(config.encoder_layers)])
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @staticmethod
- def get_reference_points(spatial_shapes, valid_ratios, device):
- """
- Get reference points for each feature map. Used in decoder.
-
- Args:
- spatial_shapes (`torch.LongTensor` of shape `(num_feature_levels, 2)`):
- Spatial shapes of each feature map.
- valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`):
- Valid ratios of each feature map.
- device (`torch.device`):
- Device on which to create the tensors.
- Returns:
- `torch.FloatTensor` of shape `(batch_size, num_queries, num_feature_levels, 2)`
- """
- reference_points_list = []
- for level, (height, width) in enumerate(spatial_shapes):
- ref_y, ref_x = meshgrid(
- torch.linspace(0.5, height - 0.5, height, dtype=torch.float32, device=device),
- torch.linspace(0.5, width - 0.5, width, dtype=torch.float32, device=device),
- indexing="ij",
- )
- # TODO: valid_ratios could be useless here. check https://github.com/fundamentalvision/Deformable-DETR/issues/36
- ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, level, 1] * height)
- ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, level, 0] * width)
- ref = torch.stack((ref_x, ref_y), -1)
- reference_points_list.append(ref)
- reference_points = torch.cat(reference_points_list, 1)
- reference_points = reference_points[:, :, None] * valid_ratios[:, None]
- return reference_points
-
- def forward(
- self,
- inputs_embeds=None,
- attention_mask=None,
- position_embeddings=None,
- spatial_shapes=None,
- level_start_index=None,
- valid_ratios=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- ):
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Flattened feature map (output of the backbone + projection layer) that is passed to the encoder.
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`:
- - 1 for pixel features that are real (i.e. **not masked**),
- - 0 for pixel features that are padding (i.e. **masked**).
- [What are attention masks?](../glossary#attention-mask)
- position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Position embeddings that are added to the queries and keys in each self-attention layer.
- spatial_shapes (`torch.LongTensor` of shape `(num_feature_levels, 2)`):
- Spatial shapes of each feature map.
- level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`):
- Starting index of each feature map.
- valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`):
- Ratio of valid area in each feature level.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- hidden_states = inputs_embeds
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
- reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=inputs_embeds.device)
-
- encoder_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
- for i, encoder_layer in enumerate(self.layers):
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
- layer_outputs = encoder_layer(
- hidden_states,
- attention_mask,
- position_embeddings=position_embeddings,
- reference_points=reference_points,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
- )
-
-
-class DetaDecoder(DetaPreTrainedModel):
- """
- Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DetaDecoderLayer`].
-
- The decoder updates the query embeddings through multiple self-attention and cross-attention layers.
-
- Some tweaks for Deformable DETR:
-
- - `position_embeddings`, `reference_points`, `spatial_shapes` and `valid_ratios` are added to the forward pass.
- - it also returns a stack of intermediate outputs and reference points from all decoding layers.
-
- Args:
- config: DetaConfig
- """
-
- def __init__(self, config: DetaConfig):
- super().__init__(config)
-
- self.dropout = config.dropout
- self.layers = nn.ModuleList([DetaDecoderLayer(config) for _ in range(config.decoder_layers)])
- self.gradient_checkpointing = False
-
- # hack implementation for iterative bounding box refinement and two-stage Deformable DETR
- self.bbox_embed = None
- self.class_embed = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def forward(
- self,
- inputs_embeds=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- position_embeddings=None,
- reference_points=None,
- spatial_shapes=None,
- level_start_index=None,
- valid_ratios=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- ):
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
- The query embeddings that are passed into the decoder.
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
- of the decoder.
- encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing cross-attention on padding pixel_values of the encoder. Mask values selected
- in `[0, 1]`:
- - 1 for pixels that are real (i.e. **not masked**),
- - 0 for pixels that are padding (i.e. **masked**).
- position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
- Position embeddings that are added to the queries and keys in each self-attention layer.
- reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)` is `as_two_stage` else `(batch_size, num_queries, 2)` or , *optional*):
- Reference point in range `[0, 1]`, top-left (0,0), bottom-right (1, 1), including padding area.
- spatial_shapes (`torch.FloatTensor` of shape `(num_feature_levels, 2)`):
- Spatial shapes of the feature maps.
- level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`, *optional*):
- Indexes for the start of each feature level. In range `[0, sequence_length]`.
- valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`, *optional*):
- Ratio of valid area in each feature level.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if inputs_embeds is not None:
- hidden_states = inputs_embeds
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
- intermediate = ()
- intermediate_reference_points = ()
-
- for idx, decoder_layer in enumerate(self.layers):
- if reference_points.shape[-1] == 4:
- reference_points_input = (
- reference_points[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[:, None]
- )
- else:
- if reference_points.shape[-1] != 2:
- raise ValueError("Reference points' last dimension must be of size 2")
- reference_points_input = reference_points[:, :, None] * valid_ratios[:, None]
-
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- decoder_layer.__call__,
- hidden_states,
- position_embeddings,
- reference_points_input,
- spatial_shapes,
- level_start_index,
- encoder_hidden_states,
- encoder_attention_mask,
- output_attentions,
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- position_embeddings=position_embeddings,
- encoder_hidden_states=encoder_hidden_states,
- reference_points=reference_points_input,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- # hack implementation for iterative bounding box refinement
- if self.bbox_embed is not None:
- tmp = self.bbox_embed[idx](hidden_states)
- if reference_points.shape[-1] == 4:
- new_reference_points = tmp + inverse_sigmoid(reference_points)
- new_reference_points = new_reference_points.sigmoid()
- else:
- if reference_points.shape[-1] != 2:
- raise ValueError(
- f"Reference points' last dimension must be of size 2, but is {reference_points.shape[-1]}"
- )
- new_reference_points = tmp
- new_reference_points[..., :2] = tmp[..., :2] + inverse_sigmoid(reference_points)
- new_reference_points = new_reference_points.sigmoid()
- reference_points = new_reference_points.detach()
-
- intermediate += (hidden_states,)
- intermediate_reference_points += (reference_points,)
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- if encoder_hidden_states is not None:
- all_cross_attentions += (layer_outputs[2],)
-
- # Keep batch_size as first dimension
- intermediate = torch.stack(intermediate, dim=1)
- intermediate_reference_points = torch.stack(intermediate_reference_points, dim=1)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- intermediate,
- intermediate_reference_points,
- all_hidden_states,
- all_self_attns,
- all_cross_attentions,
- ]
- if v is not None
- )
- return DetaDecoderOutput(
- last_hidden_state=hidden_states,
- intermediate_hidden_states=intermediate,
- intermediate_reference_points=intermediate_reference_points,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- cross_attentions=all_cross_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The bare DETA Model (consisting of a backbone and encoder-decoder Transformer) outputting raw hidden-states without
- any specific head on top.
- """,
- DETA_START_DOCSTRING,
-)
-class DetaModel(DetaPreTrainedModel):
- def __init__(self, config: DetaConfig):
- super().__init__(config)
-
- if config.two_stage:
- requires_backends(self, ["torchvision"])
-
- # Create backbone with positional encoding
- self.backbone = DetaBackboneWithPositionalEncodings(config)
- intermediate_channel_sizes = self.backbone.intermediate_channel_sizes
-
- # Create input projection layers
- if config.num_feature_levels > 1:
- num_backbone_outs = len(intermediate_channel_sizes)
- input_proj_list = []
- for _ in range(num_backbone_outs):
- in_channels = intermediate_channel_sizes[_]
- input_proj_list.append(
- nn.Sequential(
- nn.Conv2d(in_channels, config.d_model, kernel_size=1),
- nn.GroupNorm(32, config.d_model),
- )
- )
- for _ in range(config.num_feature_levels - num_backbone_outs):
- input_proj_list.append(
- nn.Sequential(
- nn.Conv2d(in_channels, config.d_model, kernel_size=3, stride=2, padding=1),
- nn.GroupNorm(32, config.d_model),
- )
- )
- in_channels = config.d_model
- self.input_proj = nn.ModuleList(input_proj_list)
- else:
- self.input_proj = nn.ModuleList(
- [
- nn.Sequential(
- nn.Conv2d(intermediate_channel_sizes[-1], config.d_model, kernel_size=1),
- nn.GroupNorm(32, config.d_model),
- )
- ]
- )
-
- if not config.two_stage:
- self.query_position_embeddings = nn.Embedding(config.num_queries, config.d_model * 2)
-
- self.encoder = DetaEncoder(config)
- self.decoder = DetaDecoder(config)
-
- self.level_embed = nn.Parameter(torch.Tensor(config.num_feature_levels, config.d_model))
-
- if config.two_stage:
- self.enc_output = nn.Linear(config.d_model, config.d_model)
- self.enc_output_norm = nn.LayerNorm(config.d_model)
- self.pos_trans = nn.Linear(config.d_model * 2, config.d_model * 2)
- self.pos_trans_norm = nn.LayerNorm(config.d_model * 2)
- self.pix_trans = nn.Linear(config.d_model, config.d_model)
- self.pix_trans_norm = nn.LayerNorm(config.d_model)
- else:
- self.reference_points = nn.Linear(config.d_model, 2)
-
- self.assign_first_stage = config.assign_first_stage
- self.two_stage_num_proposals = config.two_stage_num_proposals
-
- self.post_init()
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrModel.get_encoder
- def get_encoder(self):
- return self.encoder
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrModel.get_decoder
- def get_decoder(self):
- return self.decoder
-
- def freeze_backbone(self):
- for name, param in self.backbone.model.named_parameters():
- param.requires_grad_(False)
-
- def unfreeze_backbone(self):
- for name, param in self.backbone.model.named_parameters():
- param.requires_grad_(True)
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrModel.get_valid_ratio
- def get_valid_ratio(self, mask, dtype=torch.float32):
- """Get the valid ratio of all feature maps."""
-
- _, height, width = mask.shape
- valid_height = torch.sum(mask[:, :, 0], 1)
- valid_width = torch.sum(mask[:, 0, :], 1)
- valid_ratio_height = valid_height.to(dtype) / height
- valid_ratio_width = valid_width.to(dtype) / width
- valid_ratio = torch.stack([valid_ratio_width, valid_ratio_height], -1)
- return valid_ratio
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrModel.get_proposal_pos_embed
- def get_proposal_pos_embed(self, proposals):
- """Get the position embedding of the proposals."""
-
- num_pos_feats = self.config.d_model // 2
- temperature = 10000
- scale = 2 * math.pi
-
- dim_t = torch.arange(num_pos_feats, dtype=torch.int64, device=proposals.device).float()
- dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats)
- # batch_size, num_queries, 4
- proposals = proposals.sigmoid() * scale
- # batch_size, num_queries, 4, 128
- pos = proposals[:, :, :, None] / dim_t
- # batch_size, num_queries, 4, 64, 2 -> batch_size, num_queries, 512
- pos = torch.stack((pos[:, :, :, 0::2].sin(), pos[:, :, :, 1::2].cos()), dim=4).flatten(2)
- return pos
-
- def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes):
- """Generate the encoder output proposals from encoded enc_output.
-
- Args:
- enc_output (Tensor[batch_size, sequence_length, hidden_size]): Output of the encoder.
- padding_mask (Tensor[batch_size, sequence_length]): Padding mask for `enc_output`.
- spatial_shapes (Tensor[num_feature_levels, 2]): Spatial shapes of the feature maps.
-
- Returns:
- `tuple(torch.FloatTensor)`: A tuple of feature map and bbox prediction.
- - object_query (Tensor[batch_size, sequence_length, hidden_size]): Object query features. Later used to
- directly predict a bounding box. (without the need of a decoder)
- - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals, after an inverse
- sigmoid.
- """
- batch_size = enc_output.shape[0]
- proposals = []
- _cur = 0
- level_ids = []
- for level, (height, width) in enumerate(spatial_shapes):
- mask_flatten_ = padding_mask[:, _cur : (_cur + height * width)].view(batch_size, height, width, 1)
- valid_height = torch.sum(~mask_flatten_[:, :, 0, 0], 1)
- valid_width = torch.sum(~mask_flatten_[:, 0, :, 0], 1)
-
- grid_y, grid_x = meshgrid(
- torch.linspace(0, height - 1, height, dtype=torch.float32, device=enc_output.device),
- torch.linspace(0, width - 1, width, dtype=torch.float32, device=enc_output.device),
- indexing="ij",
- )
- grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1)
-
- scale = torch.cat([valid_width.unsqueeze(-1), valid_height.unsqueeze(-1)], 1).view(batch_size, 1, 1, 2)
- grid = (grid.unsqueeze(0).expand(batch_size, -1, -1, -1) + 0.5) / scale
- width_heigth = torch.ones_like(grid) * 0.05 * (2.0**level)
- proposal = torch.cat((grid, width_heigth), -1).view(batch_size, -1, 4)
- proposals.append(proposal)
- _cur += height * width
- level_ids.append(grid.new_ones(height * width, dtype=torch.long) * level)
- output_proposals = torch.cat(proposals, 1)
- output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True)
- output_proposals = torch.log(output_proposals / (1 - output_proposals)) # inverse sigmoid
- output_proposals = output_proposals.masked_fill(padding_mask.unsqueeze(-1), float("inf"))
- output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf"))
-
- # assign each pixel as an object query
- object_query = enc_output
- object_query = object_query.masked_fill(padding_mask.unsqueeze(-1), float(0))
- object_query = object_query.masked_fill(~output_proposals_valid, float(0))
- object_query = self.enc_output_norm(self.enc_output(object_query))
- level_ids = torch.cat(level_ids)
- return object_query, output_proposals, level_ids
-
- @add_start_docstrings_to_model_forward(DETA_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DetaModelOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.FloatTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], DetaModelOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, DetaModel
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("jozhang97/deta-swin-large-o365")
- >>> model = DetaModel.from_pretrained("jozhang97/deta-swin-large-o365", two_stage=False)
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> outputs = model(**inputs)
-
- >>> last_hidden_states = outputs.last_hidden_state
- >>> list(last_hidden_states.shape)
- [1, 900, 256]
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- batch_size, num_channels, height, width = pixel_values.shape
- device = pixel_values.device
-
- if pixel_mask is None:
- pixel_mask = torch.ones(((batch_size, height, width)), dtype=torch.long, device=device)
-
- # Extract multi-scale feature maps of same resolution `config.d_model` (cf Figure 4 in paper)
- # First, sent pixel_values + pixel_mask through Backbone to obtain the features
- # which is a list of tuples
- features, position_embeddings_list = self.backbone(pixel_values, pixel_mask)
-
- # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
- sources = []
- masks = []
- for level, (source, mask) in enumerate(features):
- sources.append(self.input_proj[level](source))
- masks.append(mask)
- if mask is None:
- raise ValueError("No attention mask was provided")
-
- # Lowest resolution feature maps are obtained via 3x3 stride 2 convolutions on the final stage
- if self.config.num_feature_levels > len(sources):
- _len_sources = len(sources)
- for level in range(_len_sources, self.config.num_feature_levels):
- if level == _len_sources:
- source = self.input_proj[level](features[-1][0])
- else:
- source = self.input_proj[level](sources[-1])
- mask = nn.functional.interpolate(pixel_mask[None].float(), size=source.shape[-2:]).to(torch.bool)[0]
- pos_l = self.backbone.position_embedding(source, mask).to(source.dtype)
- sources.append(source)
- masks.append(mask)
- position_embeddings_list.append(pos_l)
-
- # Create queries
- query_embeds = None
- if not self.config.two_stage:
- query_embeds = self.query_position_embeddings.weight
-
- # Prepare encoder inputs (by flattening)
- spatial_shapes = [(source.shape[2:]) for source in sources]
- source_flatten = [source.flatten(2).transpose(1, 2) for source in sources]
- mask_flatten = [mask.flatten(1) for mask in masks]
-
- lvl_pos_embed_flatten = []
- for level, pos_embed in enumerate(position_embeddings_list):
- pos_embed = pos_embed.flatten(2).transpose(1, 2)
- lvl_pos_embed = pos_embed + self.level_embed[level].view(1, 1, -1)
- lvl_pos_embed_flatten.append(lvl_pos_embed)
-
- source_flatten = torch.cat(source_flatten, 1)
- mask_flatten = torch.cat(mask_flatten, 1)
- lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
- spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=source_flatten.device)
- level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
- valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
- valid_ratios = valid_ratios.float()
-
- # Fourth, sent source_flatten + mask_flatten + lvl_pos_embed_flatten (backbone + proj layer output) through encoder
- # Also provide spatial_shapes, level_start_index and valid_ratios
- if encoder_outputs is None:
- encoder_outputs = self.encoder(
- inputs_embeds=source_flatten,
- attention_mask=mask_flatten,
- position_embeddings=lvl_pos_embed_flatten,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- valid_ratios=valid_ratios,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
- elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
- encoder_outputs = BaseModelOutput(
- last_hidden_state=encoder_outputs[0],
- hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
- attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
- )
-
- # Fifth, prepare decoder inputs
- batch_size, _, num_channels = encoder_outputs[0].shape
- enc_outputs_class = None
- enc_outputs_coord_logits = None
- output_proposals = None
- if self.config.two_stage:
- object_query_embedding, output_proposals, level_ids = self.gen_encoder_output_proposals(
- encoder_outputs[0], ~mask_flatten, spatial_shapes
- )
-
- # hack implementation for two-stage DETA
- # apply a detection head to each pixel (A.4 in paper)
- # linear projection for bounding box binary classification (i.e. foreground and background)
- enc_outputs_class = self.decoder.class_embed[-1](object_query_embedding)
- # 3-layer FFN to predict bounding boxes coordinates (bbox regression branch)
- delta_bbox = self.decoder.bbox_embed[-1](object_query_embedding)
- enc_outputs_coord_logits = delta_bbox + output_proposals
-
- # only keep top scoring `config.two_stage_num_proposals` proposals
- topk = self.two_stage_num_proposals
- proposal_logit = enc_outputs_class[..., 0]
-
- if self.assign_first_stage:
- proposal_boxes = center_to_corners_format(enc_outputs_coord_logits.sigmoid().float()).clamp(0, 1)
- topk_proposals = []
- for b in range(batch_size):
- prop_boxes_b = proposal_boxes[b]
- prop_logits_b = proposal_logit[b]
-
- # pre-nms per-level topk
- pre_nms_topk = 1000
- pre_nms_inds = []
- for lvl in range(len(spatial_shapes)):
- lvl_mask = level_ids == lvl
- pre_nms_inds.append(torch.topk(prop_logits_b.sigmoid() * lvl_mask, pre_nms_topk)[1])
- pre_nms_inds = torch.cat(pre_nms_inds)
-
- # nms on topk indices
- post_nms_inds = batched_nms(
- prop_boxes_b[pre_nms_inds], prop_logits_b[pre_nms_inds], level_ids[pre_nms_inds], 0.9
- )
- keep_inds = pre_nms_inds[post_nms_inds]
-
- if len(keep_inds) < self.two_stage_num_proposals:
- print(
- f"[WARNING] nms proposals ({len(keep_inds)}) < {self.two_stage_num_proposals}, running"
- " naive topk"
- )
- keep_inds = torch.topk(proposal_logit[b], topk)[1]
-
- # keep top Q/L indices for L levels
- q_per_l = topk // len(spatial_shapes)
- is_level_ordered = (
- level_ids[keep_inds][None]
- == torch.arange(len(spatial_shapes), device=level_ids.device)[:, None]
- )
- keep_inds_mask = is_level_ordered & (is_level_ordered.cumsum(1) <= q_per_l) # LS
- keep_inds_mask = keep_inds_mask.any(0) # S
-
- # pad to Q indices (might let ones filtered from pre-nms sneak by... unlikely because we pick high conf anyways)
- if keep_inds_mask.sum() < topk:
- num_to_add = topk - keep_inds_mask.sum()
- pad_inds = (~keep_inds_mask).nonzero()[:num_to_add]
- keep_inds_mask[pad_inds] = True
-
- keep_inds_topk = keep_inds[keep_inds_mask]
- topk_proposals.append(keep_inds_topk)
- topk_proposals = torch.stack(topk_proposals)
- else:
- topk_proposals = torch.topk(enc_outputs_class[..., 0], topk, dim=1)[1]
-
- topk_coords_logits = torch.gather(
- enc_outputs_coord_logits, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4)
- )
- topk_coords_logits = topk_coords_logits.detach()
- reference_points = topk_coords_logits.sigmoid()
- init_reference_points = reference_points
- pos_trans_out = self.pos_trans_norm(self.pos_trans(self.get_proposal_pos_embed(topk_coords_logits)))
- query_embed, target = torch.split(pos_trans_out, num_channels, dim=2)
-
- topk_feats = torch.stack(
- [object_query_embedding[b][topk_proposals[b]] for b in range(batch_size)]
- ).detach()
- target = target + self.pix_trans_norm(self.pix_trans(topk_feats))
- else:
- query_embed, target = torch.split(query_embeds, num_channels, dim=1)
- query_embed = query_embed.unsqueeze(0).expand(batch_size, -1, -1)
- target = target.unsqueeze(0).expand(batch_size, -1, -1)
- reference_points = self.reference_points(query_embed).sigmoid()
- init_reference_points = reference_points
-
- decoder_outputs = self.decoder(
- inputs_embeds=target,
- position_embeddings=query_embed,
- encoder_hidden_states=encoder_outputs[0],
- encoder_attention_mask=mask_flatten,
- reference_points=reference_points,
- spatial_shapes=spatial_shapes,
- level_start_index=level_start_index,
- valid_ratios=valid_ratios,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- enc_outputs = tuple(value for value in [enc_outputs_class, enc_outputs_coord_logits] if value is not None)
- tuple_outputs = (init_reference_points,) + decoder_outputs + encoder_outputs + enc_outputs
-
- return tuple_outputs
-
- return DetaModelOutput(
- init_reference_points=init_reference_points,
- last_hidden_state=decoder_outputs.last_hidden_state,
- intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
- intermediate_reference_points=decoder_outputs.intermediate_reference_points,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- enc_outputs_class=enc_outputs_class,
- enc_outputs_coord_logits=enc_outputs_coord_logits,
- output_proposals=output_proposals,
- )
-
-
-@add_start_docstrings(
- """
- DETA Model (consisting of a backbone and encoder-decoder Transformer) with object detection heads on top, for tasks
- such as COCO detection.
- """,
- DETA_START_DOCSTRING,
-)
-class DetaForObjectDetection(DetaPreTrainedModel):
- # When using clones, all layers > 0 will be clones, but layer 0 *is* required
- _tied_weights_keys = [r"bbox_embed\.\d+"]
- # We can't initialize the model on meta device as some weights are modified during the initialization
- _no_split_modules = None
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrForObjectDetection.__init__ with DeformableDetr->Deta
- def __init__(self, config: DetaConfig):
- super().__init__(config)
-
- # Deformable DETR encoder-decoder model
- self.model = DetaModel(config)
-
- # Detection heads on top
- self.class_embed = nn.Linear(config.d_model, config.num_labels)
- self.bbox_embed = DetaMLPPredictionHead(
- input_dim=config.d_model, hidden_dim=config.d_model, output_dim=4, num_layers=3
- )
-
- prior_prob = 0.01
- bias_value = -math.log((1 - prior_prob) / prior_prob)
- self.class_embed.bias.data = torch.ones(config.num_labels) * bias_value
- nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
- nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)
-
- # if two-stage, the last class_embed and bbox_embed is for region proposal generation
- num_pred = (config.decoder_layers + 1) if config.two_stage else config.decoder_layers
- if config.with_box_refine:
- self.class_embed = _get_clones(self.class_embed, num_pred)
- self.bbox_embed = _get_clones(self.bbox_embed, num_pred)
- nn.init.constant_(self.bbox_embed[0].layers[-1].bias.data[2:], -2.0)
- # hack implementation for iterative bounding box refinement
- self.model.decoder.bbox_embed = self.bbox_embed
- else:
- nn.init.constant_(self.bbox_embed.layers[-1].bias.data[2:], -2.0)
- self.class_embed = nn.ModuleList([self.class_embed for _ in range(num_pred)])
- self.bbox_embed = nn.ModuleList([self.bbox_embed for _ in range(num_pred)])
- self.model.decoder.bbox_embed = None
- if config.two_stage:
- # hack implementation for two-stage
- self.model.decoder.class_embed = self.class_embed
- for box_embed in self.bbox_embed:
- nn.init.constant_(box_embed.layers[-1].bias.data[2:], 0.0)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @torch.jit.unused
- def _set_aux_loss(self, outputs_class, outputs_coord):
- # this is a workaround to make torchscript happy, as torchscript
- # doesn't support dictionary with non-homogeneous values, such
- # as a dict having both a Tensor and a list.
- aux_loss = [
- {"logits": logits, "pred_boxes": pred_boxes}
- for logits, pred_boxes in zip(outputs_class.transpose(0, 1)[:-1], outputs_coord.transpose(0, 1)[:-1])
- ]
- return aux_loss
-
- @add_start_docstrings_to_model_forward(DETA_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DetaObjectDetectionOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.FloatTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[List[dict]] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], DetaObjectDetectionOutput]:
- r"""
- labels (`List[Dict]` of len `(batch_size,)`, *optional*):
- Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
- following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
- respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
- in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, DetaForObjectDetection
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("jozhang97/deta-swin-large")
- >>> model = DetaForObjectDetection.from_pretrained("jozhang97/deta-swin-large")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
- >>> target_sizes = torch.tensor([image.size[::-1]])
- >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[
- ... 0
- ... ]
- >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
- ... box = [round(i, 2) for i in box.tolist()]
- ... print(
- ... f"Detected {model.config.id2label[label.item()]} with confidence "
- ... f"{round(score.item(), 3)} at location {box}"
- ... )
- Detected cat with confidence 0.802 at location [9.87, 54.36, 316.93, 473.44]
- Detected cat with confidence 0.795 at location [346.62, 24.35, 639.62, 373.2]
- Detected remote with confidence 0.725 at location [40.41, 73.36, 175.77, 117.29]
- Detected remote with confidence 0.638 at location [333.34, 76.81, 370.22, 187.94]
- Detected couch with confidence 0.584 at location [0.03, 0.99, 640.02, 474.93]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # First, sent images through DETR base model to obtain encoder + decoder outputs
- outputs = self.model(
- pixel_values,
- pixel_mask=pixel_mask,
- decoder_attention_mask=decoder_attention_mask,
- encoder_outputs=encoder_outputs,
- inputs_embeds=inputs_embeds,
- decoder_inputs_embeds=decoder_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs.intermediate_hidden_states if return_dict else outputs[2]
- init_reference = outputs.init_reference_points if return_dict else outputs[0]
- inter_references = outputs.intermediate_reference_points if return_dict else outputs[3]
-
- # class logits + predicted bounding boxes
- outputs_classes = []
- outputs_coords = []
-
- for level in range(hidden_states.shape[1]):
- if level == 0:
- reference = init_reference
- else:
- reference = inter_references[:, level - 1]
- reference = inverse_sigmoid(reference)
- outputs_class = self.class_embed[level](hidden_states[:, level])
- delta_bbox = self.bbox_embed[level](hidden_states[:, level])
- if reference.shape[-1] == 4:
- outputs_coord_logits = delta_bbox + reference
- elif reference.shape[-1] == 2:
- delta_bbox[..., :2] += reference
- outputs_coord_logits = delta_bbox
- else:
- raise ValueError(f"reference.shape[-1] should be 4 or 2, but got {reference.shape[-1]}")
- outputs_coord = outputs_coord_logits.sigmoid()
- outputs_classes.append(outputs_class)
- outputs_coords.append(outputs_coord)
- # Keep batch_size as first dimension
- outputs_class = torch.stack(outputs_classes, dim=1)
- outputs_coord = torch.stack(outputs_coords, dim=1)
-
- logits = outputs_class[:, -1]
- pred_boxes = outputs_coord[:, -1]
-
- loss, loss_dict, auxiliary_outputs = None, None, None
- if labels is not None:
- # First: create the matcher
- matcher = DetaHungarianMatcher(
- class_cost=self.config.class_cost, bbox_cost=self.config.bbox_cost, giou_cost=self.config.giou_cost
- )
- # Second: create the criterion
- losses = ["labels", "boxes", "cardinality"]
- criterion = DetaLoss(
- matcher=matcher,
- num_classes=self.config.num_labels,
- focal_alpha=self.config.focal_alpha,
- losses=losses,
- num_queries=self.config.num_queries,
- assign_first_stage=self.config.assign_first_stage,
- assign_second_stage=self.config.assign_second_stage,
- )
- criterion.to(logits.device)
- # Third: compute the losses, based on outputs and labels
- outputs_loss = {}
- outputs_loss["logits"] = logits
- outputs_loss["pred_boxes"] = pred_boxes
- outputs_loss["init_reference"] = init_reference
- if self.config.auxiliary_loss:
- auxiliary_outputs = self._set_aux_loss(outputs_class, outputs_coord)
- outputs_loss["auxiliary_outputs"] = auxiliary_outputs
- if self.config.two_stage:
- enc_outputs_coord = outputs.enc_outputs_coord_logits.sigmoid()
- outputs_loss["enc_outputs"] = {
- "logits": outputs.enc_outputs_class,
- "pred_boxes": enc_outputs_coord,
- "anchors": outputs.output_proposals.sigmoid(),
- }
-
- loss_dict = criterion(outputs_loss, labels)
- # Fourth: compute total loss, as a weighted sum of the various losses
- weight_dict = {"loss_ce": 1, "loss_bbox": self.config.bbox_loss_coefficient}
- weight_dict["loss_giou"] = self.config.giou_loss_coefficient
- if self.config.auxiliary_loss:
- aux_weight_dict = {}
- for i in range(self.config.decoder_layers - 1):
- aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
- aux_weight_dict.update({k + "_enc": v for k, v in weight_dict.items()})
- weight_dict.update(aux_weight_dict)
- loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
-
- if not return_dict:
- if auxiliary_outputs is not None:
- output = (logits, pred_boxes) + auxiliary_outputs + outputs
- else:
- output = (logits, pred_boxes) + outputs
- tuple_outputs = ((loss, loss_dict) + output) if loss is not None else output
-
- return tuple_outputs
-
- dict_outputs = DetaObjectDetectionOutput(
- loss=loss,
- loss_dict=loss_dict,
- logits=logits,
- pred_boxes=pred_boxes,
- auxiliary_outputs=auxiliary_outputs,
- last_hidden_state=outputs.last_hidden_state,
- decoder_hidden_states=outputs.decoder_hidden_states,
- decoder_attentions=outputs.decoder_attentions,
- cross_attentions=outputs.cross_attentions,
- encoder_last_hidden_state=outputs.encoder_last_hidden_state,
- encoder_hidden_states=outputs.encoder_hidden_states,
- encoder_attentions=outputs.encoder_attentions,
- intermediate_hidden_states=outputs.intermediate_hidden_states,
- intermediate_reference_points=outputs.intermediate_reference_points,
- init_reference_points=outputs.init_reference_points,
- enc_outputs_class=outputs.enc_outputs_class,
- enc_outputs_coord_logits=outputs.enc_outputs_coord_logits,
- output_proposals=outputs.output_proposals,
- )
-
- return dict_outputs
-
-
-# Copied from transformers.models.detr.modeling_detr.dice_loss
-def dice_loss(inputs, targets, num_boxes):
- """
- Compute the DICE loss, similar to generalized IOU for masks
-
- Args:
- inputs: A float tensor of arbitrary shape.
- The predictions for each example.
- targets: A float tensor with the same shape as inputs. Stores the binary
- classification label for each element in inputs (0 for the negative class and 1 for the positive
- class).
- """
- inputs = inputs.sigmoid()
- inputs = inputs.flatten(1)
- numerator = 2 * (inputs * targets).sum(1)
- denominator = inputs.sum(-1) + targets.sum(-1)
- loss = 1 - (numerator + 1) / (denominator + 1)
- return loss.sum() / num_boxes
-
-
-# Copied from transformers.models.detr.modeling_detr.sigmoid_focal_loss
-def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2):
- """
- Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
-
- Args:
- inputs (`torch.FloatTensor` of arbitrary shape):
- The predictions for each example.
- targets (`torch.FloatTensor` with the same shape as `inputs`)
- A tensor storing the binary classification label for each element in the `inputs` (0 for the negative class
- and 1 for the positive class).
- alpha (`float`, *optional*, defaults to `0.25`):
- Optional weighting factor in the range (0,1) to balance positive vs. negative examples.
- gamma (`int`, *optional*, defaults to `2`):
- Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples.
-
- Returns:
- Loss tensor
- """
- prob = inputs.sigmoid()
- ce_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
- # add modulating factor
- p_t = prob * targets + (1 - prob) * (1 - targets)
- loss = ce_loss * ((1 - p_t) ** gamma)
-
- if alpha >= 0:
- alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
- loss = alpha_t * loss
-
- return loss.mean(1).sum() / num_boxes
-
-
-class DetaLoss(nn.Module):
- """
- This class computes the losses for `DetaForObjectDetection`. The process happens in two steps: 1) we compute
- hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair of matched
- ground-truth / prediction (supervised class and box).
-
- Args:
- matcher (`DetaHungarianMatcher`):
- Module able to compute a matching between targets and proposals.
- num_classes (`int`):
- Number of object categories, omitting the special no-object category.
- focal_alpha (`float`):
- Alpha parameter in focal loss.
- losses (`List[str]`):
- List of all the losses to be applied. See `get_loss` for a list of all available losses.
- """
-
- def __init__(
- self,
- matcher,
- num_classes,
- focal_alpha,
- losses,
- num_queries,
- assign_first_stage=False,
- assign_second_stage=False,
- ):
- super().__init__()
- self.matcher = matcher
- self.num_classes = num_classes
- self.focal_alpha = focal_alpha
- self.losses = losses
- self.assign_first_stage = assign_first_stage
- self.assign_second_stage = assign_second_stage
-
- if self.assign_first_stage:
- self.stg1_assigner = DetaStage1Assigner()
- if self.assign_second_stage:
- self.stg2_assigner = DetaStage2Assigner(num_queries)
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss.loss_labels
- def loss_labels(self, outputs, targets, indices, num_boxes):
- """
- Classification loss (Binary focal loss) targets dicts must contain the key "class_labels" containing a tensor
- of dim [nb_target_boxes]
- """
- if "logits" not in outputs:
- raise KeyError("No logits were found in the outputs")
- source_logits = outputs["logits"]
-
- idx = self._get_source_permutation_idx(indices)
- target_classes_o = torch.cat([t["class_labels"][J] for t, (_, J) in zip(targets, indices)])
- target_classes = torch.full(
- source_logits.shape[:2], self.num_classes, dtype=torch.int64, device=source_logits.device
- )
- target_classes[idx] = target_classes_o
-
- target_classes_onehot = torch.zeros(
- [source_logits.shape[0], source_logits.shape[1], source_logits.shape[2] + 1],
- dtype=source_logits.dtype,
- layout=source_logits.layout,
- device=source_logits.device,
- )
- target_classes_onehot.scatter_(2, target_classes.unsqueeze(-1), 1)
-
- target_classes_onehot = target_classes_onehot[:, :, :-1]
- loss_ce = (
- sigmoid_focal_loss(source_logits, target_classes_onehot, num_boxes, alpha=self.focal_alpha, gamma=2)
- * source_logits.shape[1]
- )
- losses = {"loss_ce": loss_ce}
-
- return losses
-
- @torch.no_grad()
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss.loss_cardinality
- def loss_cardinality(self, outputs, targets, indices, num_boxes):
- """
- Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes.
-
- This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients.
- """
- logits = outputs["logits"]
- device = logits.device
- target_lengths = torch.as_tensor([len(v["class_labels"]) for v in targets], device=device)
- # Count the number of predictions that are NOT "no-object" (which is the last class)
- card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1)
- card_err = nn.functional.l1_loss(card_pred.float(), target_lengths.float())
- losses = {"cardinality_error": card_err}
- return losses
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss.loss_boxes
- def loss_boxes(self, outputs, targets, indices, num_boxes):
- """
- Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss.
-
- Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes
- are expected in format (center_x, center_y, w, h), normalized by the image size.
- """
- if "pred_boxes" not in outputs:
- raise KeyError("No predicted boxes found in outputs")
- idx = self._get_source_permutation_idx(indices)
- source_boxes = outputs["pred_boxes"][idx]
- target_boxes = torch.cat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0)
-
- loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none")
-
- losses = {}
- losses["loss_bbox"] = loss_bbox.sum() / num_boxes
-
- loss_giou = 1 - torch.diag(
- generalized_box_iou(center_to_corners_format(source_boxes), center_to_corners_format(target_boxes))
- )
- losses["loss_giou"] = loss_giou.sum() / num_boxes
- return losses
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss._get_source_permutation_idx
- def _get_source_permutation_idx(self, indices):
- # permute predictions following indices
- batch_idx = torch.cat([torch.full_like(source, i) for i, (source, _) in enumerate(indices)])
- source_idx = torch.cat([source for (source, _) in indices])
- return batch_idx, source_idx
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss._get_target_permutation_idx
- def _get_target_permutation_idx(self, indices):
- # permute targets following indices
- batch_idx = torch.cat([torch.full_like(target, i) for i, (_, target) in enumerate(indices)])
- target_idx = torch.cat([target for (_, target) in indices])
- return batch_idx, target_idx
-
- # Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrLoss.get_loss
- def get_loss(self, loss, outputs, targets, indices, num_boxes):
- loss_map = {
- "labels": self.loss_labels,
- "cardinality": self.loss_cardinality,
- "boxes": self.loss_boxes,
- }
- if loss not in loss_map:
- raise ValueError(f"Loss {loss} not supported")
- return loss_map[loss](outputs, targets, indices, num_boxes)
-
- def forward(self, outputs, targets):
- """
- This performs the loss computation.
-
- Args:
- outputs (`dict`, *optional*):
- Dictionary of tensors, see the output specification of the model for the format.
- targets (`List[dict]`, *optional*):
- List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the
- losses applied, see each loss' doc.
- """
- outputs_without_aux = {k: v for k, v in outputs.items() if k not in ("auxiliary_outputs", "enc_outputs")}
-
- # Retrieve the matching between the outputs of the last layer and the targets
- if self.assign_second_stage:
- indices = self.stg2_assigner(outputs_without_aux, targets)
- else:
- indices = self.matcher(outputs_without_aux, targets)
-
- # Compute the average number of target boxes accross all nodes, for normalization purposes
- num_boxes = sum(len(t["class_labels"]) for t in targets)
- num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
- # Check that we have initialized the distributed state
- world_size = 1
- if is_accelerate_available():
- if PartialState._shared_state != {}:
- num_boxes = reduce(num_boxes)
- world_size = PartialState().num_processes
- num_boxes = torch.clamp(num_boxes / world_size, min=1).item()
-
- # Compute all the requested losses
- losses = {}
- for loss in self.losses:
- losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))
-
- # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
- if "auxiliary_outputs" in outputs:
- for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]):
- if not self.assign_second_stage:
- indices = self.matcher(auxiliary_outputs, targets)
- for loss in self.losses:
- l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes)
- l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
- losses.update(l_dict)
-
- if "enc_outputs" in outputs:
- enc_outputs = outputs["enc_outputs"]
- bin_targets = copy.deepcopy(targets)
- for bt in bin_targets:
- bt["class_labels"] = torch.zeros_like(bt["class_labels"])
- if self.assign_first_stage:
- indices = self.stg1_assigner(enc_outputs, bin_targets)
- else:
- indices = self.matcher(enc_outputs, bin_targets)
- for loss in self.losses:
- l_dict = self.get_loss(loss, enc_outputs, bin_targets, indices, num_boxes)
- l_dict = {k + "_enc": v for k, v in l_dict.items()}
- losses.update(l_dict)
-
- return losses
-
-
-# Copied from transformers.models.detr.modeling_detr.DetrMLPPredictionHead
-class DetaMLPPredictionHead(nn.Module):
- """
- Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates,
- height and width of a bounding box w.r.t. an image.
-
- Copied from https://github.com/facebookresearch/detr/blob/master/models/detr.py
-
- """
-
- def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
- super().__init__()
- self.num_layers = num_layers
- h = [hidden_dim] * (num_layers - 1)
- self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
-
- def forward(self, x):
- for i, layer in enumerate(self.layers):
- x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
- return x
-
-
-# Copied from transformers.models.deformable_detr.modeling_deformable_detr.DeformableDetrHungarianMatcher with DeformableDetr->Deta
-class DetaHungarianMatcher(nn.Module):
- """
- This class computes an assignment between the targets and the predictions of the network.
-
- For efficiency reasons, the targets don't include the no_object. Because of this, in general, there are more
- predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, while the others are
- un-matched (and thus treated as non-objects).
-
- Args:
- class_cost:
- The relative weight of the classification error in the matching cost.
- bbox_cost:
- The relative weight of the L1 error of the bounding box coordinates in the matching cost.
- giou_cost:
- The relative weight of the giou loss of the bounding box in the matching cost.
- """
-
- def __init__(self, class_cost: float = 1, bbox_cost: float = 1, giou_cost: float = 1):
- super().__init__()
- requires_backends(self, ["scipy"])
-
- self.class_cost = class_cost
- self.bbox_cost = bbox_cost
- self.giou_cost = giou_cost
- if class_cost == 0 and bbox_cost == 0 and giou_cost == 0:
- raise ValueError("All costs of the Matcher can't be 0")
-
- @torch.no_grad()
- def forward(self, outputs, targets):
- """
- Args:
- outputs (`dict`):
- A dictionary that contains at least these entries:
- * "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
- * "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates.
- targets (`List[dict]`):
- A list of targets (len(targets) = batch_size), where each target is a dict containing:
- * "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of
- ground-truth
- objects in the target) containing the class labels
- * "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates.
-
- Returns:
- `List[Tuple]`: A list of size `batch_size`, containing tuples of (index_i, index_j) where:
- - index_i is the indices of the selected predictions (in order)
- - index_j is the indices of the corresponding selected targets (in order)
- For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
- """
- batch_size, num_queries = outputs["logits"].shape[:2]
-
- # We flatten to compute the cost matrices in a batch
- out_prob = outputs["logits"].flatten(0, 1).sigmoid() # [batch_size * num_queries, num_classes]
- out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4]
-
- # Also concat the target labels and boxes
- target_ids = torch.cat([v["class_labels"] for v in targets])
- target_bbox = torch.cat([v["boxes"] for v in targets])
-
- # Compute the classification cost.
- alpha = 0.25
- gamma = 2.0
- neg_cost_class = (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log())
- pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log())
- class_cost = pos_cost_class[:, target_ids] - neg_cost_class[:, target_ids]
-
- # Compute the L1 cost between boxes
- bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)
-
- # Compute the giou cost between boxes
- giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))
-
- # Final cost matrix
- cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
- cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()
-
- sizes = [len(v["boxes"]) for v in targets]
- indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]
- return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
-
-
-# Copied from transformers.models.detr.modeling_detr._upcast
-def _upcast(t: Tensor) -> Tensor:
- # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
- if t.is_floating_point():
- return t if t.dtype in (torch.float32, torch.float64) else t.float()
- else:
- return t if t.dtype in (torch.int32, torch.int64) else t.int()
-
-
-# Copied from transformers.models.detr.modeling_detr.box_area
-def box_area(boxes: Tensor) -> Tensor:
- """
- Computes the area of a set of bounding boxes, which are specified by its (x1, y1, x2, y2) coordinates.
-
- Args:
- boxes (`torch.FloatTensor` of shape `(number_of_boxes, 4)`):
- Boxes for which the area will be computed. They are expected to be in (x1, y1, x2, y2) format with `0 <= x1
- < x2` and `0 <= y1 < y2`.
-
- Returns:
- `torch.FloatTensor`: a tensor containing the area for each box.
- """
- boxes = _upcast(boxes)
- return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
-
-
-# Copied from transformers.models.detr.modeling_detr.box_iou
-def box_iou(boxes1, boxes2):
- area1 = box_area(boxes1)
- area2 = box_area(boxes2)
-
- left_top = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
- right_bottom = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
-
- width_height = (right_bottom - left_top).clamp(min=0) # [N,M,2]
- inter = width_height[:, :, 0] * width_height[:, :, 1] # [N,M]
-
- union = area1[:, None] + area2 - inter
-
- iou = inter / union
- return iou, union
-
-
-# Copied from transformers.models.detr.modeling_detr.generalized_box_iou
-def generalized_box_iou(boxes1, boxes2):
- """
- Generalized IoU from https://giou.stanford.edu/. The boxes should be in [x0, y0, x1, y1] (corner) format.
-
- Returns:
- `torch.FloatTensor`: a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2)
- """
- # degenerate boxes gives inf / nan results
- # so do an early check
- if not (boxes1[:, 2:] >= boxes1[:, :2]).all():
- raise ValueError(f"boxes1 must be in [x0, y0, x1, y1] (corner) format, but got {boxes1}")
- if not (boxes2[:, 2:] >= boxes2[:, :2]).all():
- raise ValueError(f"boxes2 must be in [x0, y0, x1, y1] (corner) format, but got {boxes2}")
- iou, union = box_iou(boxes1, boxes2)
-
- top_left = torch.min(boxes1[:, None, :2], boxes2[:, :2])
- bottom_right = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
-
- width_height = (bottom_right - top_left).clamp(min=0) # [N,M,2]
- area = width_height[:, :, 0] * width_height[:, :, 1]
-
- return iou - (area - union) / area
-
-
-# from https://github.com/facebookresearch/detectron2/blob/cbbc1ce26473cb2a5cc8f58e8ada9ae14cb41052/detectron2/layers/wrappers.py#L100
-def nonzero_tuple(x):
- """
- A 'as_tuple=True' version of torch.nonzero to support torchscript. because of
- https://github.com/pytorch/pytorch/issues/38718
- """
- if torch.jit.is_scripting():
- if x.dim() == 0:
- return x.unsqueeze(0).nonzero().unbind(1)
- return x.nonzero().unbind(1)
- else:
- return x.nonzero(as_tuple=True)
-
-
-# from https://github.com/facebookresearch/detectron2/blob/9921a2caa585d4fa66c4b534b6fab6e74d89b582/detectron2/modeling/matcher.py#L9
-class DetaMatcher(object):
- """
- This class assigns to each predicted "element" (e.g., a box) a ground-truth element. Each predicted element will
- have exactly zero or one matches; each ground-truth element may be matched to zero or more predicted elements.
-
- The matching is determined by the MxN match_quality_matrix, that characterizes how well each (ground-truth,
- prediction)-pair match each other. For example, if the elements are boxes, this matrix may contain box
- intersection-over-union overlap values.
-
- The matcher returns (a) a vector of length N containing the index of the ground-truth element m in [0, M) that
- matches to prediction n in [0, N). (b) a vector of length N containing the labels for each prediction.
- """
-
- def __init__(self, thresholds: List[float], labels: List[int], allow_low_quality_matches: bool = False):
- """
- Args:
- thresholds (`list[float]`):
- A list of thresholds used to stratify predictions into levels.
- labels (`list[int`):
- A list of values to label predictions belonging at each level. A label can be one of {-1, 0, 1}
- signifying {ignore, negative class, positive class}, respectively.
- allow_low_quality_matches (`bool`, *optional*, defaults to `False`):
- If `True`, produce additional matches for predictions with maximum match quality lower than
- high_threshold. See `set_low_quality_matches_` for more details.
-
- For example,
- thresholds = [0.3, 0.5] labels = [0, -1, 1] All predictions with iou < 0.3 will be marked with 0 and
- thus will be considered as false positives while training. All predictions with 0.3 <= iou < 0.5 will
- be marked with -1 and thus will be ignored. All predictions with 0.5 <= iou will be marked with 1 and
- thus will be considered as true positives.
- """
- # Add -inf and +inf to first and last position in thresholds
- thresholds = thresholds[:]
- if thresholds[0] < 0:
- raise ValueError("Thresholds should be positive")
- thresholds.insert(0, -float("inf"))
- thresholds.append(float("inf"))
- # Currently torchscript does not support all + generator
- if not all(low <= high for (low, high) in zip(thresholds[:-1], thresholds[1:])):
- raise ValueError("Thresholds should be sorted.")
- if not all(l in [-1, 0, 1] for l in labels):
- raise ValueError("All labels should be either -1, 0 or 1")
- if len(labels) != len(thresholds) - 1:
- raise ValueError("Number of labels should be equal to number of thresholds - 1")
- self.thresholds = thresholds
- self.labels = labels
- self.allow_low_quality_matches = allow_low_quality_matches
-
- def __call__(self, match_quality_matrix):
- """
- Args:
- match_quality_matrix (Tensor[float]): an MxN tensor, containing the
- pairwise quality between M ground-truth elements and N predicted elements. All elements must be >= 0
- (due to the us of `torch.nonzero` for selecting indices in `set_low_quality_matches_`).
-
- Returns:
- matches (Tensor[int64]): a vector of length N, where matches[i] is a matched
- ground-truth index in [0, M)
- match_labels (Tensor[int8]): a vector of length N, where pred_labels[i] indicates
- whether a prediction is a true or false positive or ignored
- """
- assert match_quality_matrix.dim() == 2
- if match_quality_matrix.numel() == 0:
- default_matches = match_quality_matrix.new_full((match_quality_matrix.size(1),), 0, dtype=torch.int64)
- # When no gt boxes exist, we define IOU = 0 and therefore set labels
- # to `self.labels[0]`, which usually defaults to background class 0
- # To choose to ignore instead, can make labels=[-1,0,-1,1] + set appropriate thresholds
- default_match_labels = match_quality_matrix.new_full(
- (match_quality_matrix.size(1),), self.labels[0], dtype=torch.int8
- )
- return default_matches, default_match_labels
-
- assert torch.all(match_quality_matrix >= 0)
-
- # match_quality_matrix is M (gt) x N (predicted)
- # Max over gt elements (dim 0) to find best gt candidate for each prediction
- matched_vals, matches = match_quality_matrix.max(dim=0)
-
- match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)
-
- for l, low, high in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]):
- low_high = (matched_vals >= low) & (matched_vals < high)
- match_labels[low_high] = l
-
- if self.allow_low_quality_matches:
- self.set_low_quality_matches_(match_labels, match_quality_matrix)
-
- return matches, match_labels
-
- def set_low_quality_matches_(self, match_labels, match_quality_matrix):
- """
- Produce additional matches for predictions that have only low-quality matches. Specifically, for each
- ground-truth G find the set of predictions that have maximum overlap with it (including ties); for each
- prediction in that set, if it is unmatched, then match it to the ground-truth G.
-
- This function implements the RPN assignment case (i) in Sec. 3.1.2 of :paper:`Faster R-CNN`.
- """
- # For each gt, find the prediction with which it has highest quality
- highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
- # Find the highest quality match available, even if it is low, including ties.
- # Note that the matches qualities must be positive due to the use of
- # `torch.nonzero`.
- _, pred_inds_with_highest_quality = nonzero_tuple(match_quality_matrix == highest_quality_foreach_gt[:, None])
- # If an anchor was labeled positive only due to a low-quality match
- # with gt_A, but it has larger overlap with gt_B, it's matched index will still be gt_B.
- # This follows the implementation in Detectron, and is found to have no significant impact.
- match_labels[pred_inds_with_highest_quality] = 1
-
-
-# from https://github.com/facebookresearch/detectron2/blob/cbbc1ce26473cb2a5cc8f58e8ada9ae14cb41052/detectron2/modeling/sampling.py#L9
-def subsample_labels(labels: torch.Tensor, num_samples: int, positive_fraction: float, bg_label: int):
- """
- Return `num_samples` (or fewer, if not enough found) random samples from `labels` which is a mixture of positives &
- negatives. It will try to return as many positives as possible without exceeding `positive_fraction * num_samples`,
- and then try to fill the remaining slots with negatives.
-
- Args:
- labels (Tensor): (N, ) label vector with values:
- * -1: ignore
- * bg_label: background ("negative") class
- * otherwise: one or more foreground ("positive") classes
- num_samples (int): The total number of labels with value >= 0 to return.
- Values that are not sampled will be filled with -1 (ignore).
- positive_fraction (float): The number of subsampled labels with values > 0
- is `min(num_positives, int(positive_fraction * num_samples))`. The number of negatives sampled is
- `min(num_negatives, num_samples - num_positives_sampled)`. In order words, if there are not enough
- positives, the sample is filled with negatives. If there are also not enough negatives, then as many
- elements are sampled as is possible.
- bg_label (int): label index of background ("negative") class.
-
- Returns:
- pos_idx, neg_idx (Tensor):
- 1D vector of indices. The total length of both is `num_samples` or fewer.
- """
- positive = nonzero_tuple((labels != -1) & (labels != bg_label))[0]
- negative = nonzero_tuple(labels == bg_label)[0]
-
- num_pos = int(num_samples * positive_fraction)
- # protect against not enough positive examples
- num_pos = min(positive.numel(), num_pos)
- num_neg = num_samples - num_pos
- # protect against not enough negative examples
- num_neg = min(negative.numel(), num_neg)
-
- # randomly select positive and negative examples
- perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
- perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]
-
- pos_idx = positive[perm1]
- neg_idx = negative[perm2]
- return pos_idx, neg_idx
-
-
-def sample_topk_per_gt(pr_inds, gt_inds, iou, k):
- if len(gt_inds) == 0:
- return pr_inds, gt_inds
- # find topk matches for each gt
- gt_inds2, counts = gt_inds.unique(return_counts=True)
- scores, pr_inds2 = iou[gt_inds2].topk(k, dim=1)
- gt_inds2 = gt_inds2[:, None].repeat(1, k)
-
- # filter to as many matches that gt has
- pr_inds3 = torch.cat([pr[:c] for c, pr in zip(counts, pr_inds2)])
- gt_inds3 = torch.cat([gt[:c] for c, gt in zip(counts, gt_inds2)])
- return pr_inds3, gt_inds3
-
-
-# modified from https://github.com/facebookresearch/detectron2/blob/cbbc1ce26473cb2a5cc8f58e8ada9ae14cb41052/detectron2/modeling/roi_heads/roi_heads.py#L123
-class DetaStage2Assigner(nn.Module):
- def __init__(self, num_queries, max_k=4):
- super().__init__()
- self.positive_fraction = 0.25
- self.bg_label = 400 # number > 91 to filter out later
- self.batch_size_per_image = num_queries
- self.proposal_matcher = DetaMatcher(thresholds=[0.6], labels=[0, 1], allow_low_quality_matches=True)
- self.k = max_k
-
- def _sample_proposals(self, matched_idxs: torch.Tensor, matched_labels: torch.Tensor, gt_classes: torch.Tensor):
- """
- Based on the matching between N proposals and M groundtruth, sample the proposals and set their classification
- labels.
-
- Args:
- matched_idxs (Tensor): a vector of length N, each is the best-matched
- gt index in [0, M) for each proposal.
- matched_labels (Tensor): a vector of length N, the matcher's label
- (one of cfg.MODEL.ROI_HEADS.IOU_LABELS) for each proposal.
- gt_classes (Tensor): a vector of length M.
-
- Returns:
- Tensor: a vector of indices of sampled proposals. Each is in [0, N). Tensor: a vector of the same length,
- the classification label for
- each sampled proposal. Each sample is labeled as either a category in [0, num_classes) or the
- background (num_classes).
- """
- has_gt = gt_classes.numel() > 0
- # Get the corresponding GT for each proposal
- if has_gt:
- gt_classes = gt_classes[matched_idxs]
- # Label unmatched proposals (0 label from matcher) as background (label=num_classes)
- gt_classes[matched_labels == 0] = self.bg_label
- # Label ignore proposals (-1 label)
- gt_classes[matched_labels == -1] = -1
- else:
- gt_classes = torch.zeros_like(matched_idxs) + self.bg_label
-
- sampled_fg_idxs, sampled_bg_idxs = subsample_labels(
- gt_classes, self.batch_size_per_image, self.positive_fraction, self.bg_label
- )
-
- sampled_idxs = torch.cat([sampled_fg_idxs, sampled_bg_idxs], dim=0)
- return sampled_idxs, gt_classes[sampled_idxs]
-
- def forward(self, outputs, targets, return_cost_matrix=False):
- # COCO categories are from 1 to 90. They set num_classes=91 and apply sigmoid.
-
- bs = len(targets)
- indices = []
- ious = []
- for b in range(bs):
- iou, _ = box_iou(
- center_to_corners_format(targets[b]["boxes"]),
- center_to_corners_format(outputs["init_reference"][b].detach()),
- )
- matched_idxs, matched_labels = self.proposal_matcher(
- iou
- ) # proposal_id -> highest_iou_gt_id, proposal_id -> [1 if iou > 0.6, 0 ow]
- (
- sampled_idxs,
- sampled_gt_classes,
- ) = self._sample_proposals( # list of sampled proposal_ids, sampled_id -> [0, num_classes)+[bg_label]
- matched_idxs, matched_labels, targets[b]["class_labels"]
- )
- pos_pr_inds = sampled_idxs[sampled_gt_classes != self.bg_label]
- pos_gt_inds = matched_idxs[pos_pr_inds]
- pos_pr_inds, pos_gt_inds = self.postprocess_indices(pos_pr_inds, pos_gt_inds, iou)
- indices.append((pos_pr_inds, pos_gt_inds))
- ious.append(iou)
- if return_cost_matrix:
- return indices, ious
- return indices
-
- def postprocess_indices(self, pr_inds, gt_inds, iou):
- return sample_topk_per_gt(pr_inds, gt_inds, iou, self.k)
-
-
-# modified from https://github.com/facebookresearch/detectron2/blob/cbbc1ce26473cb2a5cc8f58e8ada9ae14cb41052/detectron2/modeling/proposal_generator/rpn.py#L181
-class DetaStage1Assigner(nn.Module):
- def __init__(self, t_low=0.3, t_high=0.7, max_k=4):
- super().__init__()
- self.positive_fraction = 0.5
- self.batch_size_per_image = 256
- self.k = max_k
- self.t_low = t_low
- self.t_high = t_high
- self.anchor_matcher = DetaMatcher(
- thresholds=[t_low, t_high], labels=[0, -1, 1], allow_low_quality_matches=True
- )
-
- def _subsample_labels(self, label):
- """
- Randomly sample a subset of positive and negative examples, and overwrite the label vector to the ignore value
- (-1) for all elements that are not included in the sample.
-
- Args:
- labels (Tensor): a vector of -1, 0, 1. Will be modified in-place and returned.
- """
- pos_idx, neg_idx = subsample_labels(label, self.batch_size_per_image, self.positive_fraction, 0)
- # Fill with the ignore label (-1), then set positive and negative labels
- label.fill_(-1)
- label.scatter_(0, pos_idx, 1)
- label.scatter_(0, neg_idx, 0)
- return label
-
- def forward(self, outputs, targets):
- bs = len(targets)
- indices = []
- for b in range(bs):
- anchors = outputs["anchors"][b]
- if len(targets[b]["boxes"]) == 0:
- indices.append(
- (
- torch.tensor([], dtype=torch.long, device=anchors.device),
- torch.tensor([], dtype=torch.long, device=anchors.device),
- )
- )
- continue
- iou, _ = box_iou(
- center_to_corners_format(targets[b]["boxes"]),
- center_to_corners_format(anchors),
- )
- matched_idxs, matched_labels = self.anchor_matcher(
- iou
- ) # proposal_id -> highest_iou_gt_id, proposal_id -> [1 if iou > 0.7, 0 if iou < 0.3, -1 ow]
- matched_labels = self._subsample_labels(matched_labels)
-
- all_pr_inds = torch.arange(len(anchors), device=matched_labels.device)
- pos_pr_inds = all_pr_inds[matched_labels == 1]
- pos_gt_inds = matched_idxs[pos_pr_inds]
- pos_pr_inds, pos_gt_inds = self.postprocess_indices(pos_pr_inds, pos_gt_inds, iou)
- pos_pr_inds, pos_gt_inds = pos_pr_inds.to(anchors.device), pos_gt_inds.to(anchors.device)
- indices.append((pos_pr_inds, pos_gt_inds))
- return indices
-
- def postprocess_indices(self, pr_inds, gt_inds, iou):
- return sample_topk_per_gt(pr_inds, gt_inds, iou, self.k)
diff --git a/transformers/models/detr/__init__.py b/transformers/models/detr/__init__.py
deleted file mode 100644
index 9cbaca9a54581fbe51cbf4bd88adac1660297152..0000000000000000000000000000000000000000
--- a/transformers/models/detr/__init__.py
+++ /dev/null
@@ -1,75 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {"configuration_detr": ["DETR_PRETRAINED_CONFIG_ARCHIVE_MAP", "DetrConfig", "DetrOnnxConfig"]}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_detr"] = ["DetrFeatureExtractor"]
- _import_structure["image_processing_detr"] = ["DetrImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_detr"] = [
- "DETR_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DetrForObjectDetection",
- "DetrForSegmentation",
- "DetrModel",
- "DetrPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_detr import DETR_PRETRAINED_CONFIG_ARCHIVE_MAP, DetrConfig, DetrOnnxConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_detr import DetrFeatureExtractor
- from .image_processing_detr import DetrImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_detr import (
- DETR_PRETRAINED_MODEL_ARCHIVE_LIST,
- DetrForObjectDetection,
- DetrForSegmentation,
- DetrModel,
- DetrPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/detr/__pycache__/__init__.cpython-310.pyc b/transformers/models/detr/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 6e9c792a53942a6261805aed42af2301dae071b2..0000000000000000000000000000000000000000
Binary files a/transformers/models/detr/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/detr/__pycache__/configuration_detr.cpython-310.pyc b/transformers/models/detr/__pycache__/configuration_detr.cpython-310.pyc
deleted file mode 100644
index df549a623e0d82a539826d04705457e50ff22599..0000000000000000000000000000000000000000
Binary files a/transformers/models/detr/__pycache__/configuration_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/detr/__pycache__/convert_detr_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/detr/__pycache__/convert_detr_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index fb2021cfc8124d0bb855f63d74a19efb24c2400e..0000000000000000000000000000000000000000
Binary files a/transformers/models/detr/__pycache__/convert_detr_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/detr/__pycache__/convert_detr_to_pytorch.cpython-310.pyc b/transformers/models/detr/__pycache__/convert_detr_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 46653a40e678bd1c1d6b2c64fc92242772016919..0000000000000000000000000000000000000000
Binary files a/transformers/models/detr/__pycache__/convert_detr_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/detr/__pycache__/feature_extraction_detr.cpython-310.pyc b/transformers/models/detr/__pycache__/feature_extraction_detr.cpython-310.pyc
deleted file mode 100644
index 6f7609eba93df7591a5ae80fbd20748d70578800..0000000000000000000000000000000000000000
Binary files a/transformers/models/detr/__pycache__/feature_extraction_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/detr/__pycache__/image_processing_detr.cpython-310.pyc b/transformers/models/detr/__pycache__/image_processing_detr.cpython-310.pyc
deleted file mode 100644
index e8067a90c326f10c71e9be64e65f3b32258d838f..0000000000000000000000000000000000000000
Binary files a/transformers/models/detr/__pycache__/image_processing_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/detr/__pycache__/modeling_detr.cpython-310.pyc b/transformers/models/detr/__pycache__/modeling_detr.cpython-310.pyc
deleted file mode 100644
index 5e575294222c05ac56ba8feb3657f3ad90ead390..0000000000000000000000000000000000000000
Binary files a/transformers/models/detr/__pycache__/modeling_detr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/detr/configuration_detr.py b/transformers/models/detr/configuration_detr.py
deleted file mode 100644
index 9b9b5afacd0b7f9dde6814920be350c69b8ffa88..0000000000000000000000000000000000000000
--- a/transformers/models/detr/configuration_detr.py
+++ /dev/null
@@ -1,284 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Facebook AI Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DETR model configuration"""
-
-from collections import OrderedDict
-from typing import Mapping
-
-from packaging import version
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-from ..auto import CONFIG_MAPPING
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DETR_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DetrConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DetrModel`]. It is used to instantiate a DETR
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the DETR
- [facebook/detr-resnet-50](https://huggingface.co/facebook/detr-resnet-50) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- use_timm_backbone (`bool`, *optional*, defaults to `True`):
- Whether or not to use the `timm` library for the backbone. If set to `False`, will use the [`AutoBackbone`]
- API.
- backbone_config (`PretrainedConfig` or `dict`, *optional*):
- The configuration of the backbone model. Only used in case `use_timm_backbone` is set to `False` in which
- case it will default to `ResNetConfig()`.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- num_queries (`int`, *optional*, defaults to 100):
- Number of object queries, i.e. detection slots. This is the maximal number of objects [`DetrModel`] can
- detect in a single image. For COCO, we recommend 100 queries.
- d_model (`int`, *optional*, defaults to 256):
- Dimension of the layers.
- encoder_layers (`int`, *optional*, defaults to 6):
- Number of encoder layers.
- decoder_layers (`int`, *optional*, defaults to 6):
- Number of decoder layers.
- encoder_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the Transformer encoder.
- decoder_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the Transformer decoder.
- decoder_ffn_dim (`int`, *optional*, defaults to 2048):
- Dimension of the "intermediate" (often named feed-forward) layer in decoder.
- encoder_ffn_dim (`int`, *optional*, defaults to 2048):
- Dimension of the "intermediate" (often named feed-forward) layer in decoder.
- activation_function (`str` or `function`, *optional*, defaults to `"relu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- activation_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for activations inside the fully connected layer.
- init_std (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- init_xavier_std (`float`, *optional*, defaults to 1):
- The scaling factor used for the Xavier initialization gain in the HM Attention map module.
- encoder_layerdrop (`float`, *optional*, defaults to 0.0):
- The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
- for more details.
- decoder_layerdrop (`float`, *optional*, defaults to 0.0):
- The LayerDrop probability for the decoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
- for more details.
- auxiliary_loss (`bool`, *optional*, defaults to `False`):
- Whether auxiliary decoding losses (loss at each decoder layer) are to be used.
- position_embedding_type (`str`, *optional*, defaults to `"sine"`):
- Type of position embeddings to be used on top of the image features. One of `"sine"` or `"learned"`.
- backbone (`str`, *optional*, defaults to `"resnet50"`):
- Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
- will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
- is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
- use_pretrained_backbone (`bool`, *optional*, `True`):
- Whether to use pretrained weights for the backbone.
- backbone_kwargs (`dict`, *optional*):
- Keyword arguments to be passed to AutoBackbone when loading from a checkpoint
- e.g. `{'out_indices': (0, 1, 2, 3)}`. Cannot be specified if `backbone_config` is set.
- dilation (`bool`, *optional*, defaults to `False`):
- Whether to replace stride with dilation in the last convolutional block (DC5). Only supported when
- `use_timm_backbone` = `True`.
- class_cost (`float`, *optional*, defaults to 1):
- Relative weight of the classification error in the Hungarian matching cost.
- bbox_cost (`float`, *optional*, defaults to 5):
- Relative weight of the L1 error of the bounding box coordinates in the Hungarian matching cost.
- giou_cost (`float`, *optional*, defaults to 2):
- Relative weight of the generalized IoU loss of the bounding box in the Hungarian matching cost.
- mask_loss_coefficient (`float`, *optional*, defaults to 1):
- Relative weight of the Focal loss in the panoptic segmentation loss.
- dice_loss_coefficient (`float`, *optional*, defaults to 1):
- Relative weight of the DICE/F-1 loss in the panoptic segmentation loss.
- bbox_loss_coefficient (`float`, *optional*, defaults to 5):
- Relative weight of the L1 bounding box loss in the object detection loss.
- giou_loss_coefficient (`float`, *optional*, defaults to 2):
- Relative weight of the generalized IoU loss in the object detection loss.
- eos_coefficient (`float`, *optional*, defaults to 0.1):
- Relative classification weight of the 'no-object' class in the object detection loss.
-
- Examples:
-
- ```python
- >>> from transformers import DetrConfig, DetrModel
-
- >>> # Initializing a DETR facebook/detr-resnet-50 style configuration
- >>> configuration = DetrConfig()
-
- >>> # Initializing a model (with random weights) from the facebook/detr-resnet-50 style configuration
- >>> model = DetrModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "detr"
- keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {
- "hidden_size": "d_model",
- "num_attention_heads": "encoder_attention_heads",
- }
-
- def __init__(
- self,
- use_timm_backbone=True,
- backbone_config=None,
- num_channels=3,
- num_queries=100,
- encoder_layers=6,
- encoder_ffn_dim=2048,
- encoder_attention_heads=8,
- decoder_layers=6,
- decoder_ffn_dim=2048,
- decoder_attention_heads=8,
- encoder_layerdrop=0.0,
- decoder_layerdrop=0.0,
- is_encoder_decoder=True,
- activation_function="relu",
- d_model=256,
- dropout=0.1,
- attention_dropout=0.0,
- activation_dropout=0.0,
- init_std=0.02,
- init_xavier_std=1.0,
- auxiliary_loss=False,
- position_embedding_type="sine",
- backbone="resnet50",
- use_pretrained_backbone=True,
- backbone_kwargs=None,
- dilation=False,
- class_cost=1,
- bbox_cost=5,
- giou_cost=2,
- mask_loss_coefficient=1,
- dice_loss_coefficient=1,
- bbox_loss_coefficient=5,
- giou_loss_coefficient=2,
- eos_coefficient=0.1,
- **kwargs,
- ):
- if not use_timm_backbone and use_pretrained_backbone:
- raise ValueError(
- "Loading pretrained backbone weights from the transformers library is not supported yet. `use_timm_backbone` must be set to `True` when `use_pretrained_backbone=True`"
- )
-
- if backbone_config is not None and backbone is not None:
- raise ValueError("You can't specify both `backbone` and `backbone_config`.")
-
- if backbone_config is not None and use_timm_backbone:
- raise ValueError("You can't specify both `backbone_config` and `use_timm_backbone`.")
-
- if backbone_kwargs is not None and backbone_kwargs and backbone_config is not None:
- raise ValueError("You can't specify both `backbone_kwargs` and `backbone_config`.")
-
- if not use_timm_backbone:
- if backbone_config is None:
- logger.info("`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.")
- backbone_config = CONFIG_MAPPING["resnet"](out_features=["stage4"])
- elif isinstance(backbone_config, dict):
- backbone_model_type = backbone_config.get("model_type")
- config_class = CONFIG_MAPPING[backbone_model_type]
- backbone_config = config_class.from_dict(backbone_config)
- # set timm attributes to None
- dilation, backbone, use_pretrained_backbone = None, None, None
-
- self.use_timm_backbone = use_timm_backbone
- self.backbone_config = backbone_config
- self.num_channels = num_channels
- self.num_queries = num_queries
- self.d_model = d_model
- self.encoder_ffn_dim = encoder_ffn_dim
- self.encoder_layers = encoder_layers
- self.encoder_attention_heads = encoder_attention_heads
- self.decoder_ffn_dim = decoder_ffn_dim
- self.decoder_layers = decoder_layers
- self.decoder_attention_heads = decoder_attention_heads
- self.dropout = dropout
- self.attention_dropout = attention_dropout
- self.activation_dropout = activation_dropout
- self.activation_function = activation_function
- self.init_std = init_std
- self.init_xavier_std = init_xavier_std
- self.encoder_layerdrop = encoder_layerdrop
- self.decoder_layerdrop = decoder_layerdrop
- self.num_hidden_layers = encoder_layers
- self.auxiliary_loss = auxiliary_loss
- self.position_embedding_type = position_embedding_type
- self.backbone = backbone
- self.use_pretrained_backbone = use_pretrained_backbone
- self.backbone_kwargs = backbone_kwargs
- self.dilation = dilation
- # Hungarian matcher
- self.class_cost = class_cost
- self.bbox_cost = bbox_cost
- self.giou_cost = giou_cost
- # Loss coefficients
- self.mask_loss_coefficient = mask_loss_coefficient
- self.dice_loss_coefficient = dice_loss_coefficient
- self.bbox_loss_coefficient = bbox_loss_coefficient
- self.giou_loss_coefficient = giou_loss_coefficient
- self.eos_coefficient = eos_coefficient
- super().__init__(is_encoder_decoder=is_encoder_decoder, **kwargs)
-
- @property
- def num_attention_heads(self) -> int:
- return self.encoder_attention_heads
-
- @property
- def hidden_size(self) -> int:
- return self.d_model
-
- @classmethod
- def from_backbone_config(cls, backbone_config: PretrainedConfig, **kwargs):
- """Instantiate a [`DetrConfig`] (or a derived class) from a pre-trained backbone model configuration.
-
- Args:
- backbone_config ([`PretrainedConfig`]):
- The backbone configuration.
- Returns:
- [`DetrConfig`]: An instance of a configuration object
- """
- return cls(backbone_config=backbone_config, **kwargs)
-
-
-class DetrOnnxConfig(OnnxConfig):
- torch_onnx_minimum_version = version.parse("1.11")
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- return OrderedDict(
- [
- ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
- ("pixel_mask", {0: "batch"}),
- ]
- )
-
- @property
- def atol_for_validation(self) -> float:
- return 1e-5
-
- @property
- def default_onnx_opset(self) -> int:
- return 12
diff --git a/transformers/models/detr/convert_detr_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/detr/convert_detr_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index 72de2be8701a9cf97a4e152be38da54bf87ac3d9..0000000000000000000000000000000000000000
--- a/transformers/models/detr/convert_detr_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,278 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DETR checkpoints with timm backbone."""
-
-
-import argparse
-import json
-from collections import OrderedDict
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-rename_keys = []
-for i in range(6):
- # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.self_attn.out_proj.weight", f"encoder.layers.{i}.self_attn.out_proj.weight")
- )
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.self_attn.out_proj.bias", f"encoder.layers.{i}.self_attn.out_proj.bias")
- )
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.weight", f"encoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.bias", f"encoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.weight", f"encoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.bias", f"encoder.layers.{i}.fc2.bias"))
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.norm1.weight", f"encoder.layers.{i}.self_attn_layer_norm.weight")
- )
- rename_keys.append((f"transformer.encoder.layers.{i}.norm1.bias", f"encoder.layers.{i}.self_attn_layer_norm.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.weight", f"encoder.layers.{i}.final_layer_norm.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.bias", f"encoder.layers.{i}.final_layer_norm.bias"))
- # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.self_attn.out_proj.weight", f"decoder.layers.{i}.self_attn.out_proj.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.self_attn.out_proj.bias", f"decoder.layers.{i}.self_attn.out_proj.bias")
- )
- rename_keys.append(
- (
- f"transformer.decoder.layers.{i}.multihead_attn.out_proj.weight",
- f"decoder.layers.{i}.encoder_attn.out_proj.weight",
- )
- )
- rename_keys.append(
- (
- f"transformer.decoder.layers.{i}.multihead_attn.out_proj.bias",
- f"decoder.layers.{i}.encoder_attn.out_proj.bias",
- )
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.weight", f"decoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.bias", f"decoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.weight", f"decoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.bias", f"decoder.layers.{i}.fc2.bias"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm1.weight", f"decoder.layers.{i}.self_attn_layer_norm.weight")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.norm1.bias", f"decoder.layers.{i}.self_attn_layer_norm.bias"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm2.weight", f"decoder.layers.{i}.encoder_attn_layer_norm.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm2.bias", f"decoder.layers.{i}.encoder_attn_layer_norm.bias")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.weight", f"decoder.layers.{i}.final_layer_norm.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.bias", f"decoder.layers.{i}.final_layer_norm.bias"))
-
-# convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
-rename_keys.extend(
- [
- ("input_proj.weight", "input_projection.weight"),
- ("input_proj.bias", "input_projection.bias"),
- ("query_embed.weight", "query_position_embeddings.weight"),
- ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
- ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
- ("class_embed.weight", "class_labels_classifier.weight"),
- ("class_embed.bias", "class_labels_classifier.bias"),
- ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
- ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
- ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
- ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
- ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
- ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
- ]
-)
-
-
-def rename_key(state_dict, old, new):
- val = state_dict.pop(old)
- state_dict[new] = val
-
-
-def rename_backbone_keys(state_dict):
- new_state_dict = OrderedDict()
- for key, value in state_dict.items():
- if "backbone.0.body" in key:
- new_key = key.replace("backbone.0.body", "backbone.conv_encoder.model")
- new_state_dict[new_key] = value
- else:
- new_state_dict[key] = value
-
- return new_state_dict
-
-
-def read_in_q_k_v(state_dict, is_panoptic=False):
- prefix = ""
- if is_panoptic:
- prefix = "detr."
-
- # first: transformer encoder
- for i in range(6):
- # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"encoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
- state_dict[f"encoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
- state_dict[f"encoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
- state_dict[f"encoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
- state_dict[f"encoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
- state_dict[f"encoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
- # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
- for i in range(6):
- # read in weights + bias of input projection layer of self-attention
- in_proj_weight = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"decoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
- state_dict[f"decoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
- state_dict[f"decoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
- state_dict[f"decoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
- state_dict[f"decoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
- state_dict[f"decoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
- # read in weights + bias of input projection layer of cross-attention
- in_proj_weight_cross_attn = state_dict.pop(
- f"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"
- )
- in_proj_bias_cross_attn = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias")
- # next, add query, keys and values (in that order) of cross-attention to the state dict
- state_dict[f"decoder.layers.{i}.encoder_attn.q_proj.weight"] = in_proj_weight_cross_attn[:256, :]
- state_dict[f"decoder.layers.{i}.encoder_attn.q_proj.bias"] = in_proj_bias_cross_attn[:256]
- state_dict[f"decoder.layers.{i}.encoder_attn.k_proj.weight"] = in_proj_weight_cross_attn[256:512, :]
- state_dict[f"decoder.layers.{i}.encoder_attn.k_proj.bias"] = in_proj_bias_cross_attn[256:512]
- state_dict[f"decoder.layers.{i}.encoder_attn.v_proj.weight"] = in_proj_weight_cross_attn[-256:, :]
- state_dict[f"decoder.layers.{i}.encoder_attn.v_proj.bias"] = in_proj_bias_cross_attn[-256:]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
-
- return im
-
-
-@torch.no_grad()
-def convert_detr_checkpoint(model_name, pytorch_dump_folder_path):
- """
- Copy/paste/tweak model's weights to our DETR structure.
- """
-
- # load default config
- config = DetrConfig()
- # set backbone and dilation attributes
- if "resnet101" in model_name:
- config.backbone = "resnet101"
- if "dc5" in model_name:
- config.dilation = True
- is_panoptic = "panoptic" in model_name
- if is_panoptic:
- config.num_labels = 250
- else:
- config.num_labels = 91
- repo_id = "huggingface/label-files"
- filename = "coco-detection-id2label.json"
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
-
- # load image processor
- format = "coco_panoptic" if is_panoptic else "coco_detection"
- image_processor = DetrImageProcessor(format=format)
-
- # prepare image
- img = prepare_img()
- encoding = image_processor(images=img, return_tensors="pt")
- pixel_values = encoding["pixel_values"]
-
- logger.info(f"Converting model {model_name}...")
-
- # load original model from torch hub
- detr = torch.hub.load("facebookresearch/detr", model_name, pretrained=True).eval()
- state_dict = detr.state_dict()
- # rename keys
- for src, dest in rename_keys:
- if is_panoptic:
- src = "detr." + src
- rename_key(state_dict, src, dest)
- state_dict = rename_backbone_keys(state_dict)
- # query, key and value matrices need special treatment
- read_in_q_k_v(state_dict, is_panoptic=is_panoptic)
- # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
- prefix = "detr.model." if is_panoptic else "model."
- for key in state_dict.copy().keys():
- if is_panoptic:
- if (
- key.startswith("detr")
- and not key.startswith("class_labels_classifier")
- and not key.startswith("bbox_predictor")
- ):
- val = state_dict.pop(key)
- state_dict["detr.model" + key[4:]] = val
- elif "class_labels_classifier" in key or "bbox_predictor" in key:
- val = state_dict.pop(key)
- state_dict["detr." + key] = val
- elif key.startswith("bbox_attention") or key.startswith("mask_head"):
- continue
- else:
- val = state_dict.pop(key)
- state_dict[prefix + key] = val
- else:
- if not key.startswith("class_labels_classifier") and not key.startswith("bbox_predictor"):
- val = state_dict.pop(key)
- state_dict[prefix + key] = val
- # finally, create HuggingFace model and load state dict
- model = DetrForSegmentation(config) if is_panoptic else DetrForObjectDetection(config)
- model.load_state_dict(state_dict)
- model.eval()
- # verify our conversion
- original_outputs = detr(pixel_values)
- outputs = model(pixel_values)
- assert torch.allclose(outputs.logits, original_outputs["pred_logits"], atol=1e-4)
- assert torch.allclose(outputs.pred_boxes, original_outputs["pred_boxes"], atol=1e-4)
- if is_panoptic:
- assert torch.allclose(outputs.pred_masks, original_outputs["pred_masks"], atol=1e-4)
-
- # Save model and image processor
- logger.info(f"Saving PyTorch model and image processor to {pytorch_dump_folder_path}...")
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- model.save_pretrained(pytorch_dump_folder_path)
- image_processor.save_pretrained(pytorch_dump_folder_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- "--model_name", default="detr_resnet50", type=str, help="Name of the DETR model you'd like to convert."
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
- )
- args = parser.parse_args()
- convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path)
diff --git a/transformers/models/detr/convert_detr_to_pytorch.py b/transformers/models/detr/convert_detr_to_pytorch.py
deleted file mode 100644
index a52e592b945d798ed01c457e3864252302eb33a3..0000000000000000000000000000000000000000
--- a/transformers/models/detr/convert_detr_to_pytorch.py
+++ /dev/null
@@ -1,386 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DETR checkpoints with native (Transformers) backbone."""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor, ResNetConfig
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_detr_config(model_name):
- # initialize config
- if "resnet-50" in model_name:
- backbone_config = ResNetConfig.from_pretrained("microsoft/resnet-50")
- elif "resnet-101" in model_name:
- backbone_config = ResNetConfig.from_pretrained("microsoft/resnet-101")
- else:
- raise ValueError("Model name should include either resnet50 or resnet101")
-
- config = DetrConfig(use_timm_backbone=False, backbone_config=backbone_config)
-
- # set label attributes
- is_panoptic = "panoptic" in model_name
- if is_panoptic:
- config.num_labels = 250
- else:
- config.num_labels = 91
- repo_id = "huggingface/label-files"
- filename = "coco-detection-id2label.json"
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
-
- return config, is_panoptic
-
-
-def create_rename_keys(config):
- # here we list all keys to be renamed (original name on the left, our name on the right)
- rename_keys = []
-
- # stem
- # fmt: off
- rename_keys.append(("backbone.0.body.conv1.weight", "backbone.conv_encoder.model.embedder.embedder.convolution.weight"))
- rename_keys.append(("backbone.0.body.bn1.weight", "backbone.conv_encoder.model.embedder.embedder.normalization.weight"))
- rename_keys.append(("backbone.0.body.bn1.bias", "backbone.conv_encoder.model.embedder.embedder.normalization.bias"))
- rename_keys.append(("backbone.0.body.bn1.running_mean", "backbone.conv_encoder.model.embedder.embedder.normalization.running_mean"))
- rename_keys.append(("backbone.0.body.bn1.running_var", "backbone.conv_encoder.model.embedder.embedder.normalization.running_var"))
- # stages
- for stage_idx in range(len(config.backbone_config.depths)):
- for layer_idx in range(config.backbone_config.depths[stage_idx]):
- # shortcut
- if layer_idx == 0:
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var",
- )
- )
- # 3 convs
- for i in range(3):
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean",
- )
- )
- rename_keys.append(
- (
- f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var",
- f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var",
- )
- )
- # fmt: on
-
- for i in range(config.encoder_layers):
- # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
- rename_keys.append(
- (
- f"transformer.encoder.layers.{i}.self_attn.out_proj.weight",
- f"encoder.layers.{i}.self_attn.out_proj.weight",
- )
- )
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.self_attn.out_proj.bias", f"encoder.layers.{i}.self_attn.out_proj.bias")
- )
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.weight", f"encoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear1.bias", f"encoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.weight", f"encoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.encoder.layers.{i}.linear2.bias", f"encoder.layers.{i}.fc2.bias"))
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.norm1.weight", f"encoder.layers.{i}.self_attn_layer_norm.weight")
- )
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.norm1.bias", f"encoder.layers.{i}.self_attn_layer_norm.bias")
- )
- rename_keys.append(
- (f"transformer.encoder.layers.{i}.norm2.weight", f"encoder.layers.{i}.final_layer_norm.weight")
- )
- rename_keys.append((f"transformer.encoder.layers.{i}.norm2.bias", f"encoder.layers.{i}.final_layer_norm.bias"))
- # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
- rename_keys.append(
- (
- f"transformer.decoder.layers.{i}.self_attn.out_proj.weight",
- f"decoder.layers.{i}.self_attn.out_proj.weight",
- )
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.self_attn.out_proj.bias", f"decoder.layers.{i}.self_attn.out_proj.bias")
- )
- rename_keys.append(
- (
- f"transformer.decoder.layers.{i}.multihead_attn.out_proj.weight",
- f"decoder.layers.{i}.encoder_attn.out_proj.weight",
- )
- )
- rename_keys.append(
- (
- f"transformer.decoder.layers.{i}.multihead_attn.out_proj.bias",
- f"decoder.layers.{i}.encoder_attn.out_proj.bias",
- )
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.weight", f"decoder.layers.{i}.fc1.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear1.bias", f"decoder.layers.{i}.fc1.bias"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.weight", f"decoder.layers.{i}.fc2.weight"))
- rename_keys.append((f"transformer.decoder.layers.{i}.linear2.bias", f"decoder.layers.{i}.fc2.bias"))
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm1.weight", f"decoder.layers.{i}.self_attn_layer_norm.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm1.bias", f"decoder.layers.{i}.self_attn_layer_norm.bias")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm2.weight", f"decoder.layers.{i}.encoder_attn_layer_norm.weight")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm2.bias", f"decoder.layers.{i}.encoder_attn_layer_norm.bias")
- )
- rename_keys.append(
- (f"transformer.decoder.layers.{i}.norm3.weight", f"decoder.layers.{i}.final_layer_norm.weight")
- )
- rename_keys.append((f"transformer.decoder.layers.{i}.norm3.bias", f"decoder.layers.{i}.final_layer_norm.bias"))
-
- # convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
- rename_keys.extend(
- [
- ("input_proj.weight", "input_projection.weight"),
- ("input_proj.bias", "input_projection.bias"),
- ("query_embed.weight", "query_position_embeddings.weight"),
- ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
- ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
- ("class_embed.weight", "class_labels_classifier.weight"),
- ("class_embed.bias", "class_labels_classifier.bias"),
- ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
- ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
- ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
- ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
- ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
- ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
- ]
- )
-
- return rename_keys
-
-
-def rename_key(state_dict, old, new):
- val = state_dict.pop(old)
- state_dict[new] = val
-
-
-def read_in_q_k_v(state_dict, is_panoptic=False):
- prefix = ""
- if is_panoptic:
- prefix = "detr."
-
- # first: transformer encoder
- for i in range(6):
- # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"encoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
- state_dict[f"encoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
- state_dict[f"encoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
- state_dict[f"encoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
- state_dict[f"encoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
- state_dict[f"encoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
- # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
- for i in range(6):
- # read in weights + bias of input projection layer of self-attention
- in_proj_weight = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"decoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
- state_dict[f"decoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
- state_dict[f"decoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
- state_dict[f"decoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
- state_dict[f"decoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
- state_dict[f"decoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
- # read in weights + bias of input projection layer of cross-attention
- in_proj_weight_cross_attn = state_dict.pop(
- f"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"
- )
- in_proj_bias_cross_attn = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias")
- # next, add query, keys and values (in that order) of cross-attention to the state dict
- state_dict[f"decoder.layers.{i}.encoder_attn.q_proj.weight"] = in_proj_weight_cross_attn[:256, :]
- state_dict[f"decoder.layers.{i}.encoder_attn.q_proj.bias"] = in_proj_bias_cross_attn[:256]
- state_dict[f"decoder.layers.{i}.encoder_attn.k_proj.weight"] = in_proj_weight_cross_attn[256:512, :]
- state_dict[f"decoder.layers.{i}.encoder_attn.k_proj.bias"] = in_proj_bias_cross_attn[256:512]
- state_dict[f"decoder.layers.{i}.encoder_attn.v_proj.weight"] = in_proj_weight_cross_attn[-256:, :]
- state_dict[f"decoder.layers.{i}.encoder_attn.v_proj.bias"] = in_proj_bias_cross_attn[-256:]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
-
- return im
-
-
-@torch.no_grad()
-def convert_detr_checkpoint(model_name, pytorch_dump_folder_path=None, push_to_hub=False):
- """
- Copy/paste/tweak model's weights to our DETR structure.
- """
-
- # load default config
- config, is_panoptic = get_detr_config(model_name)
-
- # load original model from torch hub
- model_name_to_original_name = {
- "detr-resnet-50": "detr_resnet50",
- "detr-resnet-101": "detr_resnet101",
- }
- logger.info(f"Converting model {model_name}...")
- detr = torch.hub.load("facebookresearch/detr", model_name_to_original_name[model_name], pretrained=True).eval()
- state_dict = detr.state_dict()
- # rename keys
- for src, dest in create_rename_keys(config):
- if is_panoptic:
- src = "detr." + src
- rename_key(state_dict, src, dest)
- # query, key and value matrices need special treatment
- read_in_q_k_v(state_dict, is_panoptic=is_panoptic)
- # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
- prefix = "detr.model." if is_panoptic else "model."
- for key in state_dict.copy().keys():
- if is_panoptic:
- if (
- key.startswith("detr")
- and not key.startswith("class_labels_classifier")
- and not key.startswith("bbox_predictor")
- ):
- val = state_dict.pop(key)
- state_dict["detr.model" + key[4:]] = val
- elif "class_labels_classifier" in key or "bbox_predictor" in key:
- val = state_dict.pop(key)
- state_dict["detr." + key] = val
- elif key.startswith("bbox_attention") or key.startswith("mask_head"):
- continue
- else:
- val = state_dict.pop(key)
- state_dict[prefix + key] = val
- else:
- if not key.startswith("class_labels_classifier") and not key.startswith("bbox_predictor"):
- val = state_dict.pop(key)
- state_dict[prefix + key] = val
-
- # finally, create HuggingFace model and load state dict
- model = DetrForSegmentation(config) if is_panoptic else DetrForObjectDetection(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- # verify our conversion on an image
- format = "coco_panoptic" if is_panoptic else "coco_detection"
- processor = DetrImageProcessor(format=format)
-
- encoding = processor(images=prepare_img(), return_tensors="pt")
- pixel_values = encoding["pixel_values"]
-
- original_outputs = detr(pixel_values)
- outputs = model(pixel_values)
-
- assert torch.allclose(outputs.logits, original_outputs["pred_logits"], atol=1e-3)
- assert torch.allclose(outputs.pred_boxes, original_outputs["pred_boxes"], atol=1e-3)
- if is_panoptic:
- assert torch.allclose(outputs.pred_masks, original_outputs["pred_masks"], atol=1e-4)
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- # Save model and image processor
- logger.info(f"Saving PyTorch model and image processor to {pytorch_dump_folder_path}...")
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- # Upload model and image processor to the hub
- logger.info("Uploading PyTorch model and image processor to the hub...")
- model.push_to_hub(f"nielsr/{model_name}")
- processor.push_to_hub(f"nielsr/{model_name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- "--model_name",
- default="detr-resnet-50",
- type=str,
- choices=["detr-resnet-50", "detr-resnet-101"],
- help="Name of the DETR model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
- )
- parser.add_argument("--push_to_hub", action="store_true", help="Whether to push the model to the hub or not.")
- args = parser.parse_args()
- convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/detr/feature_extraction_detr.py b/transformers/models/detr/feature_extraction_detr.py
deleted file mode 100644
index 6ea33666466f9a11cc074051510f0c52a2e19278..0000000000000000000000000000000000000000
--- a/transformers/models/detr/feature_extraction_detr.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for DETR."""
-
-import warnings
-
-from ...image_transforms import rgb_to_id as _rgb_to_id
-from ...utils import logging
-from .image_processing_detr import DetrImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-def rgb_to_id(x):
- warnings.warn(
- "rgb_to_id has moved and will not be importable from this module from v5. "
- "Please import from transformers.image_transforms instead.",
- FutureWarning,
- )
- return _rgb_to_id(x)
-
-
-class DetrFeatureExtractor(DetrImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class DetrFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
- " Please use DetrImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/detr/image_processing_detr.py b/transformers/models/detr/image_processing_detr.py
deleted file mode 100644
index e0e59cbc7c40c678021114f4bfe4b2aeaf418131..0000000000000000000000000000000000000000
--- a/transformers/models/detr/image_processing_detr.py
+++ /dev/null
@@ -1,1965 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for DETR."""
-
-import io
-import pathlib
-from collections import defaultdict
-from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
-from ...image_transforms import (
- PaddingMode,
- center_to_corners_format,
- corners_to_center_format,
- id_to_rgb,
- pad,
- rescale,
- resize,
- rgb_to_id,
- to_channel_dimension_format,
-)
-from ...image_utils import (
- IMAGENET_DEFAULT_MEAN,
- IMAGENET_DEFAULT_STD,
- AnnotationFormat,
- AnnotationType,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- get_image_size,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_annotations,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import (
- TensorType,
- is_flax_available,
- is_jax_tensor,
- is_scipy_available,
- is_tf_available,
- is_tf_tensor,
- is_torch_available,
- is_torch_tensor,
- is_vision_available,
- logging,
-)
-
-
-if is_torch_available():
- import torch
- from torch import nn
-
-
-if is_vision_available():
- import PIL
-
-
-if is_scipy_available():
- import scipy.special
- import scipy.stats
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-SUPPORTED_ANNOTATION_FORMATS = (AnnotationFormat.COCO_DETECTION, AnnotationFormat.COCO_PANOPTIC)
-
-
-# From the original repo: https://github.com/facebookresearch/detr/blob/3af9fa878e73b6894ce3596450a8d9b89d918ca9/datasets/transforms.py#L76
-def get_size_with_aspect_ratio(image_size, size, max_size=None) -> Tuple[int, int]:
- """
- Computes the output image size given the input image size and the desired output size.
-
- Args:
- image_size (`Tuple[int, int]`):
- The input image size.
- size (`int`):
- The desired output size.
- max_size (`int`, *optional*):
- The maximum allowed output size.
- """
- height, width = image_size
- if max_size is not None:
- min_original_size = float(min((height, width)))
- max_original_size = float(max((height, width)))
- if max_original_size / min_original_size * size > max_size:
- size = int(round(max_size * min_original_size / max_original_size))
-
- if (height <= width and height == size) or (width <= height and width == size):
- return height, width
-
- if width < height:
- ow = size
- oh = int(size * height / width)
- else:
- oh = size
- ow = int(size * width / height)
- return (oh, ow)
-
-
-def get_resize_output_image_size(
- input_image: np.ndarray,
- size: Union[int, Tuple[int, int], List[int]],
- max_size: Optional[int] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
-) -> Tuple[int, int]:
- """
- Computes the output image size given the input image size and the desired output size. If the desired output size
- is a tuple or list, the output image size is returned as is. If the desired output size is an integer, the output
- image size is computed by keeping the aspect ratio of the input image size.
-
- Args:
- input_image (`np.ndarray`):
- The image to resize.
- size (`int` or `Tuple[int, int]` or `List[int]`):
- The desired output size.
- max_size (`int`, *optional*):
- The maximum allowed output size.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred from the input image.
- """
- image_size = get_image_size(input_image, input_data_format)
- if isinstance(size, (list, tuple)):
- return size
-
- return get_size_with_aspect_ratio(image_size, size, max_size)
-
-
-def get_numpy_to_framework_fn(arr) -> Callable:
- """
- Returns a function that converts a numpy array to the framework of the input array.
-
- Args:
- arr (`np.ndarray`): The array to convert.
- """
- if isinstance(arr, np.ndarray):
- return np.array
- if is_tf_available() and is_tf_tensor(arr):
- import tensorflow as tf
-
- return tf.convert_to_tensor
- if is_torch_available() and is_torch_tensor(arr):
- import torch
-
- return torch.tensor
- if is_flax_available() and is_jax_tensor(arr):
- import jax.numpy as jnp
-
- return jnp.array
- raise ValueError(f"Cannot convert arrays of type {type(arr)}")
-
-
-def safe_squeeze(arr: np.ndarray, axis: Optional[int] = None) -> np.ndarray:
- """
- Squeezes an array, but only if the axis specified has dim 1.
- """
- if axis is None:
- return arr.squeeze()
-
- try:
- return arr.squeeze(axis=axis)
- except ValueError:
- return arr
-
-
-def normalize_annotation(annotation: Dict, image_size: Tuple[int, int]) -> Dict:
- image_height, image_width = image_size
- norm_annotation = {}
- for key, value in annotation.items():
- if key == "boxes":
- boxes = value
- boxes = corners_to_center_format(boxes)
- boxes /= np.asarray([image_width, image_height, image_width, image_height], dtype=np.float32)
- norm_annotation[key] = boxes
- else:
- norm_annotation[key] = value
- return norm_annotation
-
-
-# Copied from transformers.models.vilt.image_processing_vilt.max_across_indices
-def max_across_indices(values: Iterable[Any]) -> List[Any]:
- """
- Return the maximum value across all indices of an iterable of values.
- """
- return [max(values_i) for values_i in zip(*values)]
-
-
-# Copied from transformers.models.vilt.image_processing_vilt.get_max_height_width
-def get_max_height_width(
- images: List[np.ndarray], input_data_format: Optional[Union[str, ChannelDimension]] = None
-) -> List[int]:
- """
- Get the maximum height and width across all images in a batch.
- """
- if input_data_format is None:
- input_data_format = infer_channel_dimension_format(images[0])
-
- if input_data_format == ChannelDimension.FIRST:
- _, max_height, max_width = max_across_indices([img.shape for img in images])
- elif input_data_format == ChannelDimension.LAST:
- max_height, max_width, _ = max_across_indices([img.shape for img in images])
- else:
- raise ValueError(f"Invalid channel dimension format: {input_data_format}")
- return (max_height, max_width)
-
-
-# Copied from transformers.models.vilt.image_processing_vilt.make_pixel_mask
-def make_pixel_mask(
- image: np.ndarray, output_size: Tuple[int, int], input_data_format: Optional[Union[str, ChannelDimension]] = None
-) -> np.ndarray:
- """
- Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.
-
- Args:
- image (`np.ndarray`):
- Image to make the pixel mask for.
- output_size (`Tuple[int, int]`):
- Output size of the mask.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- mask = np.zeros(output_size, dtype=np.int64)
- mask[:input_height, :input_width] = 1
- return mask
-
-
-# inspired by https://github.com/facebookresearch/detr/blob/master/datasets/coco.py#L33
-def convert_coco_poly_to_mask(segmentations, height: int, width: int) -> np.ndarray:
- """
- Convert a COCO polygon annotation to a mask.
-
- Args:
- segmentations (`List[List[float]]`):
- List of polygons, each polygon represented by a list of x-y coordinates.
- height (`int`):
- Height of the mask.
- width (`int`):
- Width of the mask.
- """
- try:
- from pycocotools import mask as coco_mask
- except ImportError:
- raise ImportError("Pycocotools is not installed in your environment.")
-
- masks = []
- for polygons in segmentations:
- rles = coco_mask.frPyObjects(polygons, height, width)
- mask = coco_mask.decode(rles)
- if len(mask.shape) < 3:
- mask = mask[..., None]
- mask = np.asarray(mask, dtype=np.uint8)
- mask = np.any(mask, axis=2)
- masks.append(mask)
- if masks:
- masks = np.stack(masks, axis=0)
- else:
- masks = np.zeros((0, height, width), dtype=np.uint8)
-
- return masks
-
-
-# inspired by https://github.com/facebookresearch/detr/blob/master/datasets/coco.py#L50
-def prepare_coco_detection_annotation(
- image,
- target,
- return_segmentation_masks: bool = False,
- input_data_format: Optional[Union[ChannelDimension, str]] = None,
-):
- """
- Convert the target in COCO format into the format expected by DETR.
- """
- image_height, image_width = get_image_size(image, channel_dim=input_data_format)
-
- image_id = target["image_id"]
- image_id = np.asarray([image_id], dtype=np.int64)
-
- # Get all COCO annotations for the given image.
- annotations = target["annotations"]
- annotations = [obj for obj in annotations if "iscrowd" not in obj or obj["iscrowd"] == 0]
-
- classes = [obj["category_id"] for obj in annotations]
- classes = np.asarray(classes, dtype=np.int64)
-
- # for conversion to coco api
- area = np.asarray([obj["area"] for obj in annotations], dtype=np.float32)
- iscrowd = np.asarray([obj["iscrowd"] if "iscrowd" in obj else 0 for obj in annotations], dtype=np.int64)
-
- boxes = [obj["bbox"] for obj in annotations]
- # guard against no boxes via resizing
- boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4)
- boxes[:, 2:] += boxes[:, :2]
- boxes[:, 0::2] = boxes[:, 0::2].clip(min=0, max=image_width)
- boxes[:, 1::2] = boxes[:, 1::2].clip(min=0, max=image_height)
-
- keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
-
- new_target = {}
- new_target["image_id"] = image_id
- new_target["class_labels"] = classes[keep]
- new_target["boxes"] = boxes[keep]
- new_target["area"] = area[keep]
- new_target["iscrowd"] = iscrowd[keep]
- new_target["orig_size"] = np.asarray([int(image_height), int(image_width)], dtype=np.int64)
-
- if annotations and "keypoints" in annotations[0]:
- keypoints = [obj["keypoints"] for obj in annotations]
- # Converting the filtered keypoints list to a numpy array
- keypoints = np.asarray(keypoints, dtype=np.float32)
- # Apply the keep mask here to filter the relevant annotations
- keypoints = keypoints[keep]
- num_keypoints = keypoints.shape[0]
- keypoints = keypoints.reshape((-1, 3)) if num_keypoints else keypoints
- new_target["keypoints"] = keypoints
-
- if return_segmentation_masks:
- segmentation_masks = [obj["segmentation"] for obj in annotations]
- masks = convert_coco_poly_to_mask(segmentation_masks, image_height, image_width)
- new_target["masks"] = masks[keep]
-
- return new_target
-
-
-def masks_to_boxes(masks: np.ndarray) -> np.ndarray:
- """
- Compute the bounding boxes around the provided panoptic segmentation masks.
-
- Args:
- masks: masks in format `[number_masks, height, width]` where N is the number of masks
-
- Returns:
- boxes: bounding boxes in format `[number_masks, 4]` in xyxy format
- """
- if masks.size == 0:
- return np.zeros((0, 4))
-
- h, w = masks.shape[-2:]
- y = np.arange(0, h, dtype=np.float32)
- x = np.arange(0, w, dtype=np.float32)
- # see https://github.com/pytorch/pytorch/issues/50276
- y, x = np.meshgrid(y, x, indexing="ij")
-
- x_mask = masks * np.expand_dims(x, axis=0)
- x_max = x_mask.reshape(x_mask.shape[0], -1).max(-1)
- x = np.ma.array(x_mask, mask=~(np.array(masks, dtype=bool)))
- x_min = x.filled(fill_value=1e8)
- x_min = x_min.reshape(x_min.shape[0], -1).min(-1)
-
- y_mask = masks * np.expand_dims(y, axis=0)
- y_max = y_mask.reshape(x_mask.shape[0], -1).max(-1)
- y = np.ma.array(y_mask, mask=~(np.array(masks, dtype=bool)))
- y_min = y.filled(fill_value=1e8)
- y_min = y_min.reshape(y_min.shape[0], -1).min(-1)
-
- return np.stack([x_min, y_min, x_max, y_max], 1)
-
-
-def prepare_coco_panoptic_annotation(
- image: np.ndarray,
- target: Dict,
- masks_path: Union[str, pathlib.Path],
- return_masks: bool = True,
- input_data_format: Union[ChannelDimension, str] = None,
-) -> Dict:
- """
- Prepare a coco panoptic annotation for DETR.
- """
- image_height, image_width = get_image_size(image, channel_dim=input_data_format)
- annotation_path = pathlib.Path(masks_path) / target["file_name"]
-
- new_target = {}
- new_target["image_id"] = np.asarray([target["image_id"] if "image_id" in target else target["id"]], dtype=np.int64)
- new_target["size"] = np.asarray([image_height, image_width], dtype=np.int64)
- new_target["orig_size"] = np.asarray([image_height, image_width], dtype=np.int64)
-
- if "segments_info" in target:
- masks = np.asarray(PIL.Image.open(annotation_path), dtype=np.uint32)
- masks = rgb_to_id(masks)
-
- ids = np.array([segment_info["id"] for segment_info in target["segments_info"]])
- masks = masks == ids[:, None, None]
- masks = masks.astype(np.uint8)
- if return_masks:
- new_target["masks"] = masks
- new_target["boxes"] = masks_to_boxes(masks)
- new_target["class_labels"] = np.array(
- [segment_info["category_id"] for segment_info in target["segments_info"]], dtype=np.int64
- )
- new_target["iscrowd"] = np.asarray(
- [segment_info["iscrowd"] for segment_info in target["segments_info"]], dtype=np.int64
- )
- new_target["area"] = np.asarray(
- [segment_info["area"] for segment_info in target["segments_info"]], dtype=np.float32
- )
-
- return new_target
-
-
-def get_segmentation_image(
- masks: np.ndarray, input_size: Tuple, target_size: Tuple, stuff_equiv_classes, deduplicate=False
-):
- h, w = input_size
- final_h, final_w = target_size
-
- m_id = scipy.special.softmax(masks.transpose(0, 1), -1)
-
- if m_id.shape[-1] == 0:
- # We didn't detect any mask :(
- m_id = np.zeros((h, w), dtype=np.int64)
- else:
- m_id = m_id.argmax(-1).reshape(h, w)
-
- if deduplicate:
- # Merge the masks corresponding to the same stuff class
- for equiv in stuff_equiv_classes.values():
- for eq_id in equiv:
- m_id[m_id == eq_id] = equiv[0]
-
- seg_img = id_to_rgb(m_id)
- seg_img = resize(seg_img, (final_w, final_h), resample=PILImageResampling.NEAREST)
- return seg_img
-
-
-def get_mask_area(seg_img: np.ndarray, target_size: Tuple[int, int], n_classes: int) -> np.ndarray:
- final_h, final_w = target_size
- np_seg_img = seg_img.astype(np.uint8)
- np_seg_img = np_seg_img.reshape(final_h, final_w, 3)
- m_id = rgb_to_id(np_seg_img)
- area = [(m_id == i).sum() for i in range(n_classes)]
- return area
-
-
-def score_labels_from_class_probabilities(logits: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
- probs = scipy.special.softmax(logits, axis=-1)
- labels = probs.argmax(-1, keepdims=True)
- scores = np.take_along_axis(probs, labels, axis=-1)
- scores, labels = scores.squeeze(-1), labels.squeeze(-1)
- return scores, labels
-
-
-def post_process_panoptic_sample(
- out_logits: np.ndarray,
- masks: np.ndarray,
- boxes: np.ndarray,
- processed_size: Tuple[int, int],
- target_size: Tuple[int, int],
- is_thing_map: Dict,
- threshold=0.85,
-) -> Dict:
- """
- Converts the output of [`DetrForSegmentation`] into panoptic segmentation predictions for a single sample.
-
- Args:
- out_logits (`torch.Tensor`):
- The logits for this sample.
- masks (`torch.Tensor`):
- The predicted segmentation masks for this sample.
- boxes (`torch.Tensor`):
- The prediced bounding boxes for this sample. The boxes are in the normalized format `(center_x, center_y,
- width, height)` and values between `[0, 1]`, relative to the size the image (disregarding padding).
- processed_size (`Tuple[int, int]`):
- The processed size of the image `(height, width)`, as returned by the preprocessing step i.e. the size
- after data augmentation but before batching.
- target_size (`Tuple[int, int]`):
- The target size of the image, `(height, width)` corresponding to the requested final size of the
- prediction.
- is_thing_map (`Dict`):
- A dictionary mapping class indices to a boolean value indicating whether the class is a thing or not.
- threshold (`float`, *optional*, defaults to 0.85):
- The threshold used to binarize the segmentation masks.
- """
- # we filter empty queries and detection below threshold
- scores, labels = score_labels_from_class_probabilities(out_logits)
- keep = (labels != out_logits.shape[-1] - 1) & (scores > threshold)
-
- cur_scores = scores[keep]
- cur_classes = labels[keep]
- cur_boxes = center_to_corners_format(boxes[keep])
-
- if len(cur_boxes) != len(cur_classes):
- raise ValueError("Not as many boxes as there are classes")
-
- cur_masks = masks[keep]
- cur_masks = resize(cur_masks[:, None], processed_size, resample=PILImageResampling.BILINEAR)
- cur_masks = safe_squeeze(cur_masks, 1)
- b, h, w = cur_masks.shape
-
- # It may be that we have several predicted masks for the same stuff class.
- # In the following, we track the list of masks ids for each stuff class (they are merged later on)
- cur_masks = cur_masks.reshape(b, -1)
- stuff_equiv_classes = defaultdict(list)
- for k, label in enumerate(cur_classes):
- if not is_thing_map[label]:
- stuff_equiv_classes[label].append(k)
-
- seg_img = get_segmentation_image(cur_masks, processed_size, target_size, stuff_equiv_classes, deduplicate=True)
- area = get_mask_area(cur_masks, processed_size, n_classes=len(cur_scores))
-
- # We filter out any mask that is too small
- if cur_classes.size() > 0:
- # We know filter empty masks as long as we find some
- filtered_small = np.array([a <= 4 for a in area], dtype=bool)
- while filtered_small.any():
- cur_masks = cur_masks[~filtered_small]
- cur_scores = cur_scores[~filtered_small]
- cur_classes = cur_classes[~filtered_small]
- seg_img = get_segmentation_image(cur_masks, (h, w), target_size, stuff_equiv_classes, deduplicate=True)
- area = get_mask_area(seg_img, target_size, n_classes=len(cur_scores))
- filtered_small = np.array([a <= 4 for a in area], dtype=bool)
- else:
- cur_classes = np.ones((1, 1), dtype=np.int64)
-
- segments_info = [
- {"id": i, "isthing": is_thing_map[cat], "category_id": int(cat), "area": a}
- for i, (cat, a) in enumerate(zip(cur_classes, area))
- ]
- del cur_classes
-
- with io.BytesIO() as out:
- PIL.Image.fromarray(seg_img).save(out, format="PNG")
- predictions = {"png_string": out.getvalue(), "segments_info": segments_info}
-
- return predictions
-
-
-def resize_annotation(
- annotation: Dict[str, Any],
- orig_size: Tuple[int, int],
- target_size: Tuple[int, int],
- threshold: float = 0.5,
- resample: PILImageResampling = PILImageResampling.NEAREST,
-):
- """
- Resizes an annotation to a target size.
-
- Args:
- annotation (`Dict[str, Any]`):
- The annotation dictionary.
- orig_size (`Tuple[int, int]`):
- The original size of the input image.
- target_size (`Tuple[int, int]`):
- The target size of the image, as returned by the preprocessing `resize` step.
- threshold (`float`, *optional*, defaults to 0.5):
- The threshold used to binarize the segmentation masks.
- resample (`PILImageResampling`, defaults to `PILImageResampling.NEAREST`):
- The resampling filter to use when resizing the masks.
- """
- ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(target_size, orig_size))
- ratio_height, ratio_width = ratios
-
- new_annotation = {}
- new_annotation["size"] = target_size
-
- for key, value in annotation.items():
- if key == "boxes":
- boxes = value
- scaled_boxes = boxes * np.asarray([ratio_width, ratio_height, ratio_width, ratio_height], dtype=np.float32)
- new_annotation["boxes"] = scaled_boxes
- elif key == "area":
- area = value
- scaled_area = area * (ratio_width * ratio_height)
- new_annotation["area"] = scaled_area
- elif key == "masks":
- masks = value[:, None]
- masks = np.array([resize(mask, target_size, resample=resample) for mask in masks])
- masks = masks.astype(np.float32)
- masks = masks[:, 0] > threshold
- new_annotation["masks"] = masks
- elif key == "size":
- new_annotation["size"] = target_size
- else:
- new_annotation[key] = value
-
- return new_annotation
-
-
-# TODO - (Amy) make compatible with other frameworks
-def binary_mask_to_rle(mask):
- """
- Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format.
-
- Args:
- mask (`torch.Tensor` or `numpy.array`):
- A binary mask tensor of shape `(height, width)` where 0 denotes background and 1 denotes the target
- segment_id or class_id.
- Returns:
- `List`: Run-length encoded list of the binary mask. Refer to COCO API for more information about the RLE
- format.
- """
- if is_torch_tensor(mask):
- mask = mask.numpy()
-
- pixels = mask.flatten()
- pixels = np.concatenate([[0], pixels, [0]])
- runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
- runs[1::2] -= runs[::2]
- return list(runs)
-
-
-# TODO - (Amy) make compatible with other frameworks
-def convert_segmentation_to_rle(segmentation):
- """
- Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format.
-
- Args:
- segmentation (`torch.Tensor` or `numpy.array`):
- A segmentation map of shape `(height, width)` where each value denotes a segment or class id.
- Returns:
- `List[List]`: A list of lists, where each list is the run-length encoding of a segment / class id.
- """
- segment_ids = torch.unique(segmentation)
-
- run_length_encodings = []
- for idx in segment_ids:
- mask = torch.where(segmentation == idx, 1, 0)
- rle = binary_mask_to_rle(mask)
- run_length_encodings.append(rle)
-
- return run_length_encodings
-
-
-def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels):
- """
- Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and
- `labels`.
-
- Args:
- masks (`torch.Tensor`):
- A tensor of shape `(num_queries, height, width)`.
- scores (`torch.Tensor`):
- A tensor of shape `(num_queries)`.
- labels (`torch.Tensor`):
- A tensor of shape `(num_queries)`.
- object_mask_threshold (`float`):
- A number between 0 and 1 used to binarize the masks.
- Raises:
- `ValueError`: Raised when the first dimension doesn't match in all input tensors.
- Returns:
- `Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region
- < `object_mask_threshold`.
- """
- if not (masks.shape[0] == scores.shape[0] == labels.shape[0]):
- raise ValueError("mask, scores and labels must have the same shape!")
-
- to_keep = labels.ne(num_labels) & (scores > object_mask_threshold)
-
- return masks[to_keep], scores[to_keep], labels[to_keep]
-
-
-def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8):
- # Get the mask associated with the k class
- mask_k = mask_labels == k
- mask_k_area = mask_k.sum()
-
- # Compute the area of all the stuff in query k
- original_area = (mask_probs[k] >= mask_threshold).sum()
- mask_exists = mask_k_area > 0 and original_area > 0
-
- # Eliminate disconnected tiny segments
- if mask_exists:
- area_ratio = mask_k_area / original_area
- if not area_ratio.item() > overlap_mask_area_threshold:
- mask_exists = False
-
- return mask_exists, mask_k
-
-
-def compute_segments(
- mask_probs,
- pred_scores,
- pred_labels,
- mask_threshold: float = 0.5,
- overlap_mask_area_threshold: float = 0.8,
- label_ids_to_fuse: Optional[Set[int]] = None,
- target_size: Tuple[int, int] = None,
-):
- height = mask_probs.shape[1] if target_size is None else target_size[0]
- width = mask_probs.shape[2] if target_size is None else target_size[1]
-
- segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device)
- segments: List[Dict] = []
-
- if target_size is not None:
- mask_probs = nn.functional.interpolate(
- mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False
- )[0]
-
- current_segment_id = 0
-
- # Weigh each mask by its prediction score
- mask_probs *= pred_scores.view(-1, 1, 1)
- mask_labels = mask_probs.argmax(0) # [height, width]
-
- # Keep track of instances of each class
- stuff_memory_list: Dict[str, int] = {}
- for k in range(pred_labels.shape[0]):
- pred_class = pred_labels[k].item()
- should_fuse = pred_class in label_ids_to_fuse
-
- # Check if mask exists and large enough to be a segment
- mask_exists, mask_k = check_segment_validity(
- mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold
- )
-
- if mask_exists:
- if pred_class in stuff_memory_list:
- current_segment_id = stuff_memory_list[pred_class]
- else:
- current_segment_id += 1
-
- # Add current object segment to final segmentation map
- segmentation[mask_k] = current_segment_id
- segment_score = round(pred_scores[k].item(), 6)
- segments.append(
- {
- "id": current_segment_id,
- "label_id": pred_class,
- "was_fused": should_fuse,
- "score": segment_score,
- }
- )
- if should_fuse:
- stuff_memory_list[pred_class] = current_segment_id
-
- return segmentation, segments
-
-
-class DetrImageProcessor(BaseImageProcessor):
- r"""
- Constructs a Detr image processor.
-
- Args:
- format (`str`, *optional*, defaults to `"coco_detection"`):
- Data format of the annotations. One of "coco_detection" or "coco_panoptic".
- do_resize (`bool`, *optional*, defaults to `True`):
- Controls whether to resize the image's `(height, width)` dimensions to the specified `size`. Can be
- overridden by the `do_resize` parameter in the `preprocess` method.
- size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 800, "longest_edge": 1333}`):
- Size of the image's `(height, width)` dimensions after resizing. Can be overridden by the `size` parameter
- in the `preprocess` method.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Controls whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the
- `do_rescale` parameter in the `preprocess` method.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the
- `preprocess` method.
- do_normalize (`bool`, *optional*, defaults to True):
- Controls whether to normalize the image. Can be overridden by the `do_normalize` parameter in the
- `preprocess` method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_DEFAULT_MEAN`):
- Mean values to use when normalizing the image. Can be a single value or a list of values, one for each
- channel. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_DEFAULT_STD`):
- Standard deviation values to use when normalizing the image. Can be a single value or a list of values, one
- for each channel. Can be overridden by the `image_std` parameter in the `preprocess` method.
- do_convert_annotations (`bool`, *optional*, defaults to `True`):
- Controls whether to convert the annotations to the format expected by the DETR model. Converts the
- bounding boxes to the format `(center_x, center_y, width, height)` and in the range `[0, 1]`.
- Can be overridden by the `do_convert_annotations` parameter in the `preprocess` method.
- do_pad (`bool`, *optional*, defaults to `True`):
- Controls whether to pad the image. Can be overridden by the `do_pad` parameter in the `preprocess`
- method. If `True` will pad the images in the batch to the largest height and width in the batch.
- Padding will be applied to the bottom and right of the image with zeros.
- """
-
- model_input_names = ["pixel_values", "pixel_mask"]
-
- def __init__(
- self,
- format: Union[str, AnnotationFormat] = AnnotationFormat.COCO_DETECTION,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- do_normalize: bool = True,
- image_mean: Union[float, List[float]] = None,
- image_std: Union[float, List[float]] = None,
- do_convert_annotations: Optional[bool] = None,
- do_pad: bool = True,
- **kwargs,
- ) -> None:
- if "pad_and_return_pixel_mask" in kwargs:
- do_pad = kwargs.pop("pad_and_return_pixel_mask")
-
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` parameter is deprecated and will be removed in v4.26. "
- "Please specify in `size['longest_edge'] instead`.",
- )
- max_size = kwargs.pop("max_size")
- else:
- max_size = None if size is None else 1333
-
- size = size if size is not None else {"shortest_edge": 800, "longest_edge": 1333}
- size = get_size_dict(size, max_size=max_size, default_to_square=False)
-
- # Backwards compatibility
- if do_convert_annotations is None:
- do_convert_annotations = do_normalize
-
- super().__init__(**kwargs)
- self.format = format
- self.do_resize = do_resize
- self.size = size
- self.resample = resample
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_normalize = do_normalize
- self.do_convert_annotations = do_convert_annotations
- self.image_mean = image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
- self.do_pad = do_pad
- self._valid_processor_keys = [
- "images",
- "annotations",
- "return_segmentation_masks",
- "masks_path",
- "do_resize",
- "size",
- "resample",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "do_convert_annotations",
- "image_mean",
- "image_std",
- "do_pad",
- "format",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- @classmethod
- def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
- """
- Overrides the `from_dict` method from the base class to make sure parameters are updated if image processor is
- created using from_dict and kwargs e.g. `DetrImageProcessor.from_pretrained(checkpoint, size=600,
- max_size=800)`
- """
- image_processor_dict = image_processor_dict.copy()
- if "max_size" in kwargs:
- image_processor_dict["max_size"] = kwargs.pop("max_size")
- if "pad_and_return_pixel_mask" in kwargs:
- image_processor_dict["pad_and_return_pixel_mask"] = kwargs.pop("pad_and_return_pixel_mask")
- return super().from_dict(image_processor_dict, **kwargs)
-
- def prepare_annotation(
- self,
- image: np.ndarray,
- target: Dict,
- format: Optional[AnnotationFormat] = None,
- return_segmentation_masks: bool = None,
- masks_path: Optional[Union[str, pathlib.Path]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> Dict:
- """
- Prepare an annotation for feeding into DETR model.
- """
- format = format if format is not None else self.format
-
- if format == AnnotationFormat.COCO_DETECTION:
- return_segmentation_masks = False if return_segmentation_masks is None else return_segmentation_masks
- target = prepare_coco_detection_annotation(
- image, target, return_segmentation_masks, input_data_format=input_data_format
- )
- elif format == AnnotationFormat.COCO_PANOPTIC:
- return_segmentation_masks = True if return_segmentation_masks is None else return_segmentation_masks
- target = prepare_coco_panoptic_annotation(
- image,
- target,
- masks_path=masks_path,
- return_masks=return_segmentation_masks,
- input_data_format=input_data_format,
- )
- else:
- raise ValueError(f"Format {format} is not supported.")
- return target
-
- def prepare(self, image, target, return_segmentation_masks=None, masks_path=None):
- logger.warning_once(
- "The `prepare` method is deprecated and will be removed in a v4.33. "
- "Please use `prepare_annotation` instead. Note: the `prepare_annotation` method "
- "does not return the image anymore.",
- )
- target = self.prepare_annotation(image, target, return_segmentation_masks, masks_path, self.format)
- return image, target
-
- def convert_coco_poly_to_mask(self, *args, **kwargs):
- logger.warning_once("The `convert_coco_poly_to_mask` method is deprecated and will be removed in v4.33. ")
- return convert_coco_poly_to_mask(*args, **kwargs)
-
- def prepare_coco_detection(self, *args, **kwargs):
- logger.warning_once("The `prepare_coco_detection` method is deprecated and will be removed in v4.33. ")
- return prepare_coco_detection_annotation(*args, **kwargs)
-
- def prepare_coco_panoptic(self, *args, **kwargs):
- logger.warning_once("The `prepare_coco_panoptic` method is deprecated and will be removed in v4.33. ")
- return prepare_coco_panoptic_annotation(*args, **kwargs)
-
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize the image to the given size. Size can be `min_size` (scalar) or `(height, width)` tuple. If size is an
- int, smaller edge of the image will be matched to this number.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary containing the size to resize to. Can contain the keys `shortest_edge` and `longest_edge` or
- `height` and `width`.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` parameter is deprecated and will be removed in v4.26. "
- "Please specify in `size['longest_edge'] instead`.",
- )
- max_size = kwargs.pop("max_size")
- else:
- max_size = None
- size = get_size_dict(size, max_size=max_size, default_to_square=False)
- if "shortest_edge" in size and "longest_edge" in size:
- size = get_resize_output_image_size(
- image, size["shortest_edge"], size["longest_edge"], input_data_format=input_data_format
- )
- elif "height" in size and "width" in size:
- size = (size["height"], size["width"])
- else:
- raise ValueError(
- "Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got"
- f" {size.keys()}."
- )
- image = resize(
- image, size=size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs
- )
- return image
-
- def resize_annotation(
- self,
- annotation,
- orig_size,
- size,
- resample: PILImageResampling = PILImageResampling.NEAREST,
- ) -> Dict:
- """
- Resize the annotation to match the resized image. If size is an int, smaller edge of the mask will be matched
- to this number.
- """
- return resize_annotation(annotation, orig_size=orig_size, target_size=size, resample=resample)
-
- # TODO (Amy) - update to use `rescale_factor` instead of `scale`
- def rescale(
- self,
- image: np.ndarray,
- rescale_factor: float,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.ndarray:
- """
- Rescale the image by the given factor. image = image * rescale_factor.
-
- Args:
- image (`np.ndarray`):
- Image to rescale.
- rescale_factor (`float`):
- The value to use for rescaling.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format for the input image. If unset, is inferred from the input image. Can be
- one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- """
- return rescale(image, rescale_factor, data_format=data_format, input_data_format=input_data_format)
-
- def normalize_annotation(self, annotation: Dict, image_size: Tuple[int, int]) -> Dict:
- """
- Normalize the boxes in the annotation from `[top_left_x, top_left_y, bottom_right_x, bottom_right_y]` to
- `[center_x, center_y, width, height]` format and from absolute to relative pixel values.
- """
- return normalize_annotation(annotation, image_size=image_size)
-
- def _update_annotation_for_padded_image(
- self,
- annotation: Dict,
- input_image_size: Tuple[int, int],
- output_image_size: Tuple[int, int],
- padding,
- update_bboxes,
- ) -> Dict:
- """
- Update the annotation for a padded image.
- """
- new_annotation = {}
- new_annotation["size"] = output_image_size
-
- for key, value in annotation.items():
- if key == "masks":
- masks = value
- masks = pad(
- masks,
- padding,
- mode=PaddingMode.CONSTANT,
- constant_values=0,
- input_data_format=ChannelDimension.FIRST,
- )
- masks = safe_squeeze(masks, 1)
- new_annotation["masks"] = masks
- elif key == "boxes" and update_bboxes:
- boxes = value
- boxes *= np.asarray(
- [
- input_image_size[1] / output_image_size[1],
- input_image_size[0] / output_image_size[0],
- input_image_size[1] / output_image_size[1],
- input_image_size[0] / output_image_size[0],
- ]
- )
- new_annotation["boxes"] = boxes
- elif key == "size":
- new_annotation["size"] = output_image_size
- else:
- new_annotation[key] = value
- return new_annotation
-
- def _pad_image(
- self,
- image: np.ndarray,
- output_size: Tuple[int, int],
- annotation: Optional[Dict[str, Any]] = None,
- constant_values: Union[float, Iterable[float]] = 0,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- update_bboxes: bool = True,
- ) -> np.ndarray:
- """
- Pad an image with zeros to the given size.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- output_height, output_width = output_size
-
- pad_bottom = output_height - input_height
- pad_right = output_width - input_width
- padding = ((0, pad_bottom), (0, pad_right))
- padded_image = pad(
- image,
- padding,
- mode=PaddingMode.CONSTANT,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- )
- if annotation is not None:
- annotation = self._update_annotation_for_padded_image(
- annotation, (input_height, input_width), (output_height, output_width), padding, update_bboxes
- )
- return padded_image, annotation
-
- def pad(
- self,
- images: List[np.ndarray],
- annotations: Optional[Union[AnnotationType, List[AnnotationType]]] = None,
- constant_values: Union[float, Iterable[float]] = 0,
- return_pixel_mask: bool = True,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- update_bboxes: bool = True,
- ) -> BatchFeature:
- """
- Pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width
- in the batch and optionally returns their corresponding pixel mask.
-
- Args:
- images (List[`np.ndarray`]):
- Images to pad.
- annotations (`AnnotationType` or `List[AnnotationType]`, *optional*):
- Annotations to transform according to the padding that is applied to the images.
- constant_values (`float` or `Iterable[float]`, *optional*):
- The value to use for the padding if `mode` is `"constant"`.
- return_pixel_mask (`bool`, *optional*, defaults to `True`):
- Whether to return a pixel mask.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the image. If not provided, it will be the same as the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- update_bboxes (`bool`, *optional*, defaults to `True`):
- Whether to update the bounding boxes in the annotations to match the padded images. If the
- bounding boxes have not been converted to relative coordinates and `(centre_x, centre_y, width, height)`
- format, the bounding boxes will not be updated.
- """
- pad_size = get_max_height_width(images, input_data_format=input_data_format)
-
- annotation_list = annotations if annotations is not None else [None] * len(images)
- padded_images = []
- padded_annotations = []
- for image, annotation in zip(images, annotation_list):
- padded_image, padded_annotation = self._pad_image(
- image,
- pad_size,
- annotation,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- update_bboxes=update_bboxes,
- )
- padded_images.append(padded_image)
- padded_annotations.append(padded_annotation)
-
- data = {"pixel_values": padded_images}
-
- if return_pixel_mask:
- masks = [
- make_pixel_mask(image=image, output_size=pad_size, input_data_format=input_data_format)
- for image in images
- ]
- data["pixel_mask"] = masks
-
- encoded_inputs = BatchFeature(data=data, tensor_type=return_tensors)
-
- if annotations is not None:
- encoded_inputs["labels"] = [
- BatchFeature(annotation, tensor_type=return_tensors) for annotation in padded_annotations
- ]
-
- return encoded_inputs
-
- def preprocess(
- self,
- images: ImageInput,
- annotations: Optional[Union[AnnotationType, List[AnnotationType]]] = None,
- return_segmentation_masks: bool = None,
- masks_path: Optional[Union[str, pathlib.Path]] = None,
- do_resize: Optional[bool] = None,
- size: Optional[Dict[str, int]] = None,
- resample=None, # PILImageResampling
- do_rescale: Optional[bool] = None,
- rescale_factor: Optional[Union[int, float]] = None,
- do_normalize: Optional[bool] = None,
- do_convert_annotations: Optional[bool] = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- do_pad: Optional[bool] = None,
- format: Optional[Union[str, AnnotationFormat]] = None,
- return_tensors: Optional[Union[TensorType, str]] = None,
- data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> BatchFeature:
- """
- Preprocess an image or a batch of images so that it can be used by the model.
-
- Args:
- images (`ImageInput`):
- Image or batch of images to preprocess. Expects a single or batch of images with pixel values ranging
- from 0 to 255. If passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- annotations (`AnnotationType` or `List[AnnotationType]`, *optional*):
- List of annotations associated with the image or batch of images. If annotation is for object
- detection, the annotations should be a dictionary with the following keys:
- - "image_id" (`int`): The image id.
- - "annotations" (`List[Dict]`): List of annotations for an image. Each annotation should be a
- dictionary. An image can have no annotations, in which case the list should be empty.
- If annotation is for segmentation, the annotations should be a dictionary with the following keys:
- - "image_id" (`int`): The image id.
- - "segments_info" (`List[Dict]`): List of segments for an image. Each segment should be a dictionary.
- An image can have no segments, in which case the list should be empty.
- - "file_name" (`str`): The file name of the image.
- return_segmentation_masks (`bool`, *optional*, defaults to self.return_segmentation_masks):
- Whether to return segmentation masks.
- masks_path (`str` or `pathlib.Path`, *optional*):
- Path to the directory containing the segmentation masks.
- do_resize (`bool`, *optional*, defaults to self.do_resize):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to self.size):
- Size of the image after resizing.
- resample (`PILImageResampling`, *optional*, defaults to self.resample):
- Resampling filter to use when resizing the image.
- do_rescale (`bool`, *optional*, defaults to self.do_rescale):
- Whether to rescale the image.
- rescale_factor (`float`, *optional*, defaults to self.rescale_factor):
- Rescale factor to use when rescaling the image.
- do_normalize (`bool`, *optional*, defaults to self.do_normalize):
- Whether to normalize the image.
- do_convert_annotations (`bool`, *optional*, defaults to self.do_convert_annotations):
- Whether to convert the annotations to the format expected by the model. Converts the bounding
- boxes from the format `(top_left_x, top_left_y, width, height)` to `(center_x, center_y, width, height)`
- and in relative coordinates.
- image_mean (`float` or `List[float]`, *optional*, defaults to self.image_mean):
- Mean to use when normalizing the image.
- image_std (`float` or `List[float]`, *optional*, defaults to self.image_std):
- Standard deviation to use when normalizing the image.
- do_pad (`bool`, *optional*, defaults to self.do_pad):
- Whether to pad the image. If `True` will pad the images in the batch to the largest image in the batch
- and create a pixel mask. Padding will be applied to the bottom and right of the image with zeros.
- format (`str` or `AnnotationFormat`, *optional*, defaults to self.format):
- Format of the annotations.
- return_tensors (`str` or `TensorType`, *optional*, defaults to self.return_tensors):
- Type of tensors to return. If `None`, will return the list of images.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - Unset: Use the channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- if "pad_and_return_pixel_mask" in kwargs:
- logger.warning_once(
- "The `pad_and_return_pixel_mask` argument is deprecated and will be removed in a future version, "
- "use `do_pad` instead."
- )
- do_pad = kwargs.pop("pad_and_return_pixel_mask")
-
- max_size = None
- if "max_size" in kwargs:
- logger.warning_once(
- "The `max_size` argument is deprecated and will be removed in a future version, use"
- " `size['longest_edge']` instead."
- )
- size = kwargs.pop("max_size")
-
- do_resize = self.do_resize if do_resize is None else do_resize
- size = self.size if size is None else size
- size = get_size_dict(size=size, max_size=max_size, default_to_square=False)
- resample = self.resample if resample is None else resample
- do_rescale = self.do_rescale if do_rescale is None else do_rescale
- rescale_factor = self.rescale_factor if rescale_factor is None else rescale_factor
- do_normalize = self.do_normalize if do_normalize is None else do_normalize
- image_mean = self.image_mean if image_mean is None else image_mean
- image_std = self.image_std if image_std is None else image_std
- do_convert_annotations = (
- self.do_convert_annotations if do_convert_annotations is None else do_convert_annotations
- )
- do_pad = self.do_pad if do_pad is None else do_pad
- format = self.format if format is None else format
-
- images = make_list_of_images(images)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- # Here, the pad() method pads to the maximum of (width, height). It does not need to be validated.
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
-
- if annotations is not None and isinstance(annotations, dict):
- annotations = [annotations]
-
- if annotations is not None and len(images) != len(annotations):
- raise ValueError(
- f"The number of images ({len(images)}) and annotations ({len(annotations)}) do not match."
- )
-
- format = AnnotationFormat(format)
- if annotations is not None:
- validate_annotations(format, SUPPORTED_ANNOTATION_FORMATS, annotations)
-
- if (
- masks_path is not None
- and format == AnnotationFormat.COCO_PANOPTIC
- and not isinstance(masks_path, (pathlib.Path, str))
- ):
- raise ValueError(
- "The path to the directory containing the mask PNG files should be provided as a"
- f" `pathlib.Path` or string object, but is {type(masks_path)} instead."
- )
-
- # All transformations expect numpy arrays
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- # prepare (COCO annotations as a list of Dict -> DETR target as a single Dict per image)
- if annotations is not None:
- prepared_images = []
- prepared_annotations = []
- for image, target in zip(images, annotations):
- target = self.prepare_annotation(
- image,
- target,
- format,
- return_segmentation_masks=return_segmentation_masks,
- masks_path=masks_path,
- input_data_format=input_data_format,
- )
- prepared_images.append(image)
- prepared_annotations.append(target)
- images = prepared_images
- annotations = prepared_annotations
- del prepared_images, prepared_annotations
-
- # transformations
- if do_resize:
- if annotations is not None:
- resized_images, resized_annotations = [], []
- for image, target in zip(images, annotations):
- orig_size = get_image_size(image, input_data_format)
- resized_image = self.resize(
- image, size=size, max_size=max_size, resample=resample, input_data_format=input_data_format
- )
- resized_annotation = self.resize_annotation(
- target, orig_size, get_image_size(resized_image, input_data_format)
- )
- resized_images.append(resized_image)
- resized_annotations.append(resized_annotation)
- images = resized_images
- annotations = resized_annotations
- del resized_images, resized_annotations
- else:
- images = [
- self.resize(image, size=size, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_rescale:
- images = [self.rescale(image, rescale_factor, input_data_format=input_data_format) for image in images]
-
- if do_normalize:
- images = [
- self.normalize(image, image_mean, image_std, input_data_format=input_data_format) for image in images
- ]
-
- if do_convert_annotations and annotations is not None:
- annotations = [
- self.normalize_annotation(annotation, get_image_size(image, input_data_format))
- for annotation, image in zip(annotations, images)
- ]
-
- if do_pad:
- # Pads images and returns their mask: {'pixel_values': ..., 'pixel_mask': ...}
- encoded_inputs = self.pad(
- images,
- annotations=annotations,
- return_pixel_mask=True,
- data_format=data_format,
- input_data_format=input_data_format,
- update_bboxes=do_convert_annotations,
- return_tensors=return_tensors,
- )
- else:
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
- for image in images
- ]
- encoded_inputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)
- if annotations is not None:
- encoded_inputs["labels"] = [
- BatchFeature(annotation, tensor_type=return_tensors) for annotation in annotations
- ]
-
- return encoded_inputs
-
- # POSTPROCESSING METHODS - TODO: add support for other frameworks
- # inspired by https://github.com/facebookresearch/detr/blob/master/models/detr.py#L258
- def post_process(self, outputs, target_sizes):
- """
- Converts the raw output of [`DetrForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y,
- bottom_right_x, bottom_right_y) format. Only supports PyTorch.
-
- Args:
- outputs ([`DetrObjectDetectionOutput`]):
- Raw outputs of the model.
- target_sizes (`torch.Tensor` of shape `(batch_size, 2)`):
- Tensor containing the size (height, width) of each image of the batch. For evaluation, this must be the
- original image size (before any data augmentation). For visualization, this should be the image size
- after data augment, but before padding.
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
- in the batch as predicted by the model.
- """
- logger.warning_once(
- "`post_process` is deprecated and will be removed in v5 of Transformers, please use"
- " `post_process_object_detection` instead, with `threshold=0.` for equivalent results.",
- )
-
- out_logits, out_bbox = outputs.logits, outputs.pred_boxes
-
- if len(out_logits) != len(target_sizes):
- raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the logits")
- if target_sizes.shape[1] != 2:
- raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch")
-
- prob = nn.functional.softmax(out_logits, -1)
- scores, labels = prob[..., :-1].max(-1)
-
- # convert to [x0, y0, x1, y1] format
- boxes = center_to_corners_format(out_bbox)
- # and from relative [0, 1] to absolute [0, height] coordinates
- img_h, img_w = target_sizes.unbind(1)
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device)
- boxes = boxes * scale_fct[:, None, :]
-
- results = [{"scores": s, "labels": l, "boxes": b} for s, l, b in zip(scores, labels, boxes)]
- return results
-
- def post_process_segmentation(self, outputs, target_sizes, threshold=0.9, mask_threshold=0.5):
- """
- Converts the output of [`DetrForSegmentation`] into image segmentation predictions. Only supports PyTorch.
-
- Args:
- outputs ([`DetrSegmentationOutput`]):
- Raw outputs of the model.
- target_sizes (`torch.Tensor` of shape `(batch_size, 2)` or `List[Tuple]` of length `batch_size`):
- Torch Tensor (or list) corresponding to the requested final size (h, w) of each prediction.
- threshold (`float`, *optional*, defaults to 0.9):
- Threshold to use to filter out queries.
- mask_threshold (`float`, *optional*, defaults to 0.5):
- Threshold to use when turning the predicted masks into binary values.
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels, and masks for an image
- in the batch as predicted by the model.
- """
- logger.warning_once(
- "`post_process_segmentation` is deprecated and will be removed in v5 of Transformers, please use"
- " `post_process_semantic_segmentation`.",
- )
- out_logits, raw_masks = outputs.logits, outputs.pred_masks
- empty_label = out_logits.shape[-1] - 1
- preds = []
-
- def to_tuple(tup):
- if isinstance(tup, tuple):
- return tup
- return tuple(tup.cpu().tolist())
-
- for cur_logits, cur_masks, size in zip(out_logits, raw_masks, target_sizes):
- # we filter empty queries and detection below threshold
- cur_scores, cur_labels = cur_logits.softmax(-1).max(-1)
- keep = cur_labels.ne(empty_label) & (cur_scores > threshold)
- cur_scores = cur_scores[keep]
- cur_labels = cur_labels[keep]
- cur_masks = cur_masks[keep]
- cur_masks = nn.functional.interpolate(cur_masks[:, None], to_tuple(size), mode="bilinear").squeeze(1)
- cur_masks = (cur_masks.sigmoid() > mask_threshold) * 1
-
- predictions = {"scores": cur_scores, "labels": cur_labels, "masks": cur_masks}
- preds.append(predictions)
- return preds
-
- # inspired by https://github.com/facebookresearch/detr/blob/master/models/segmentation.py#L218
- def post_process_instance(self, results, outputs, orig_target_sizes, max_target_sizes, threshold=0.5):
- """
- Converts the output of [`DetrForSegmentation`] into actual instance segmentation predictions. Only supports
- PyTorch.
-
- Args:
- results (`List[Dict]`):
- Results list obtained by [`~DetrImageProcessor.post_process`], to which "masks" results will be added.
- outputs ([`DetrSegmentationOutput`]):
- Raw outputs of the model.
- orig_target_sizes (`torch.Tensor` of shape `(batch_size, 2)`):
- Tensor containing the size (h, w) of each image of the batch. For evaluation, this must be the original
- image size (before any data augmentation).
- max_target_sizes (`torch.Tensor` of shape `(batch_size, 2)`):
- Tensor containing the maximum size (h, w) of each image of the batch. For evaluation, this must be the
- original image size (before any data augmentation).
- threshold (`float`, *optional*, defaults to 0.5):
- Threshold to use when turning the predicted masks into binary values.
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels, boxes and masks for an
- image in the batch as predicted by the model.
- """
- logger.warning_once(
- "`post_process_instance` is deprecated and will be removed in v5 of Transformers, please use"
- " `post_process_instance_segmentation`.",
- )
-
- if len(orig_target_sizes) != len(max_target_sizes):
- raise ValueError("Make sure to pass in as many orig_target_sizes as max_target_sizes")
- max_h, max_w = max_target_sizes.max(0)[0].tolist()
- outputs_masks = outputs.pred_masks.squeeze(2)
- outputs_masks = nn.functional.interpolate(
- outputs_masks, size=(max_h, max_w), mode="bilinear", align_corners=False
- )
- outputs_masks = (outputs_masks.sigmoid() > threshold).cpu()
-
- for i, (cur_mask, t, tt) in enumerate(zip(outputs_masks, max_target_sizes, orig_target_sizes)):
- img_h, img_w = t[0], t[1]
- results[i]["masks"] = cur_mask[:, :img_h, :img_w].unsqueeze(1)
- results[i]["masks"] = nn.functional.interpolate(
- results[i]["masks"].float(), size=tuple(tt.tolist()), mode="nearest"
- ).byte()
-
- return results
-
- # inspired by https://github.com/facebookresearch/detr/blob/master/models/segmentation.py#L241
- def post_process_panoptic(self, outputs, processed_sizes, target_sizes=None, is_thing_map=None, threshold=0.85):
- """
- Converts the output of [`DetrForSegmentation`] into actual panoptic predictions. Only supports PyTorch.
-
- Args:
- outputs ([`DetrSegmentationOutput`]):
- Raw outputs of the model.
- processed_sizes (`torch.Tensor` of shape `(batch_size, 2)` or `List[Tuple]` of length `batch_size`):
- Torch Tensor (or list) containing the size (h, w) of each image of the batch, i.e. the size after data
- augmentation but before batching.
- target_sizes (`torch.Tensor` of shape `(batch_size, 2)` or `List[Tuple]` of length `batch_size`, *optional*):
- Torch Tensor (or list) corresponding to the requested final size `(height, width)` of each prediction.
- If left to None, it will default to the `processed_sizes`.
- is_thing_map (`torch.Tensor` of shape `(batch_size, 2)`, *optional*):
- Dictionary mapping class indices to either True or False, depending on whether or not they are a thing.
- If not set, defaults to the `is_thing_map` of COCO panoptic.
- threshold (`float`, *optional*, defaults to 0.85):
- Threshold to use to filter out queries.
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing a PNG string and segments_info values for
- an image in the batch as predicted by the model.
- """
- logger.warning_once(
- "`post_process_panoptic is deprecated and will be removed in v5 of Transformers, please use"
- " `post_process_panoptic_segmentation`.",
- )
- if target_sizes is None:
- target_sizes = processed_sizes
- if len(processed_sizes) != len(target_sizes):
- raise ValueError("Make sure to pass in as many processed_sizes as target_sizes")
-
- if is_thing_map is None:
- # default to is_thing_map of COCO panoptic
- is_thing_map = {i: i <= 90 for i in range(201)}
-
- out_logits, raw_masks, raw_boxes = outputs.logits, outputs.pred_masks, outputs.pred_boxes
- if not len(out_logits) == len(raw_masks) == len(target_sizes):
- raise ValueError(
- "Make sure that you pass in as many target sizes as the batch dimension of the logits and masks"
- )
- empty_label = out_logits.shape[-1] - 1
- preds = []
-
- def to_tuple(tup):
- if isinstance(tup, tuple):
- return tup
- return tuple(tup.cpu().tolist())
-
- for cur_logits, cur_masks, cur_boxes, size, target_size in zip(
- out_logits, raw_masks, raw_boxes, processed_sizes, target_sizes
- ):
- # we filter empty queries and detection below threshold
- cur_scores, cur_labels = cur_logits.softmax(-1).max(-1)
- keep = cur_labels.ne(empty_label) & (cur_scores > threshold)
- cur_scores = cur_scores[keep]
- cur_labels = cur_labels[keep]
- cur_masks = cur_masks[keep]
- cur_masks = nn.functional.interpolate(cur_masks[:, None], to_tuple(size), mode="bilinear").squeeze(1)
- cur_boxes = center_to_corners_format(cur_boxes[keep])
-
- h, w = cur_masks.shape[-2:]
- if len(cur_boxes) != len(cur_labels):
- raise ValueError("Not as many boxes as there are classes")
-
- # It may be that we have several predicted masks for the same stuff class.
- # In the following, we track the list of masks ids for each stuff class (they are merged later on)
- cur_masks = cur_masks.flatten(1)
- stuff_equiv_classes = defaultdict(lambda: [])
- for k, label in enumerate(cur_labels):
- if not is_thing_map[label.item()]:
- stuff_equiv_classes[label.item()].append(k)
-
- def get_ids_area(masks, scores, dedup=False):
- # This helper function creates the final panoptic segmentation image
- # It also returns the area of the masks that appears on the image
-
- m_id = masks.transpose(0, 1).softmax(-1)
-
- if m_id.shape[-1] == 0:
- # We didn't detect any mask :(
- m_id = torch.zeros((h, w), dtype=torch.long, device=m_id.device)
- else:
- m_id = m_id.argmax(-1).view(h, w)
-
- if dedup:
- # Merge the masks corresponding to the same stuff class
- for equiv in stuff_equiv_classes.values():
- if len(equiv) > 1:
- for eq_id in equiv:
- m_id.masked_fill_(m_id.eq(eq_id), equiv[0])
-
- final_h, final_w = to_tuple(target_size)
-
- seg_img = PIL.Image.fromarray(id_to_rgb(m_id.view(h, w).cpu().numpy()))
- seg_img = seg_img.resize(size=(final_w, final_h), resample=PILImageResampling.NEAREST)
-
- np_seg_img = torch.ByteTensor(torch.ByteStorage.from_buffer(seg_img.tobytes()))
- np_seg_img = np_seg_img.view(final_h, final_w, 3)
- np_seg_img = np_seg_img.numpy()
-
- m_id = torch.from_numpy(rgb_to_id(np_seg_img))
-
- area = []
- for i in range(len(scores)):
- area.append(m_id.eq(i).sum().item())
- return area, seg_img
-
- area, seg_img = get_ids_area(cur_masks, cur_scores, dedup=True)
- if cur_labels.numel() > 0:
- # We know filter empty masks as long as we find some
- while True:
- filtered_small = torch.as_tensor(
- [area[i] <= 4 for i, c in enumerate(cur_labels)], dtype=torch.bool, device=keep.device
- )
- if filtered_small.any().item():
- cur_scores = cur_scores[~filtered_small]
- cur_labels = cur_labels[~filtered_small]
- cur_masks = cur_masks[~filtered_small]
- area, seg_img = get_ids_area(cur_masks, cur_scores)
- else:
- break
-
- else:
- cur_labels = torch.ones(1, dtype=torch.long, device=cur_labels.device)
-
- segments_info = []
- for i, a in enumerate(area):
- cat = cur_labels[i].item()
- segments_info.append({"id": i, "isthing": is_thing_map[cat], "category_id": cat, "area": a})
- del cur_labels
-
- with io.BytesIO() as out:
- seg_img.save(out, format="PNG")
- predictions = {"png_string": out.getvalue(), "segments_info": segments_info}
- preds.append(predictions)
- return preds
-
- # inspired by https://github.com/facebookresearch/detr/blob/master/models/detr.py#L258
- def post_process_object_detection(
- self, outputs, threshold: float = 0.5, target_sizes: Union[TensorType, List[Tuple]] = None
- ):
- """
- Converts the raw output of [`DetrForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y,
- bottom_right_x, bottom_right_y) format. Only supports PyTorch.
-
- Args:
- outputs ([`DetrObjectDetectionOutput`]):
- Raw outputs of the model.
- threshold (`float`, *optional*):
- Score threshold to keep object detection predictions.
- target_sizes (`torch.Tensor` or `List[Tuple[int, int]]`, *optional*):
- Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
- `(height, width)` of each image in the batch. If unset, predictions will not be resized.
- Returns:
- `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
- in the batch as predicted by the model.
- """
- out_logits, out_bbox = outputs.logits, outputs.pred_boxes
-
- if target_sizes is not None:
- if len(out_logits) != len(target_sizes):
- raise ValueError(
- "Make sure that you pass in as many target sizes as the batch dimension of the logits"
- )
-
- prob = nn.functional.softmax(out_logits, -1)
- scores, labels = prob[..., :-1].max(-1)
-
- # Convert to [x0, y0, x1, y1] format
- boxes = center_to_corners_format(out_bbox)
-
- # Convert from relative [0, 1] to absolute [0, height] coordinates
- if target_sizes is not None:
- if isinstance(target_sizes, List):
- img_h = torch.Tensor([i[0] for i in target_sizes])
- img_w = torch.Tensor([i[1] for i in target_sizes])
- else:
- img_h, img_w = target_sizes.unbind(1)
-
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device)
- boxes = boxes * scale_fct[:, None, :]
-
- results = []
- for s, l, b in zip(scores, labels, boxes):
- score = s[s > threshold]
- label = l[s > threshold]
- box = b[s > threshold]
- results.append({"scores": score, "labels": label, "boxes": box})
-
- return results
-
- def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple[int, int]] = None):
- """
- Converts the output of [`DetrForSegmentation`] into semantic segmentation maps. Only supports PyTorch.
-
- Args:
- outputs ([`DetrForSegmentation`]):
- Raw outputs of the model.
- target_sizes (`List[Tuple[int, int]]`, *optional*):
- A list of tuples (`Tuple[int, int]`) containing the target size (height, width) of each image in the
- batch. If unset, predictions will not be resized.
- Returns:
- `List[torch.Tensor]`:
- A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width)
- corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each
- `torch.Tensor` correspond to a semantic class id.
- """
- class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1]
- masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width]
-
- # Remove the null class `[..., :-1]`
- masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
- masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width]
-
- # Semantic segmentation logits of shape (batch_size, num_classes, height, width)
- segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
- batch_size = class_queries_logits.shape[0]
-
- # Resize logits and compute semantic segmentation maps
- if target_sizes is not None:
- if batch_size != len(target_sizes):
- raise ValueError(
- "Make sure that you pass in as many target sizes as the batch dimension of the logits"
- )
-
- semantic_segmentation = []
- for idx in range(batch_size):
- resized_logits = nn.functional.interpolate(
- segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
- )
- semantic_map = resized_logits[0].argmax(dim=0)
- semantic_segmentation.append(semantic_map)
- else:
- semantic_segmentation = segmentation.argmax(dim=1)
- semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]
-
- return semantic_segmentation
-
- # inspired by https://github.com/facebookresearch/detr/blob/master/models/segmentation.py#L218
- def post_process_instance_segmentation(
- self,
- outputs,
- threshold: float = 0.5,
- mask_threshold: float = 0.5,
- overlap_mask_area_threshold: float = 0.8,
- target_sizes: Optional[List[Tuple[int, int]]] = None,
- return_coco_annotation: Optional[bool] = False,
- ) -> List[Dict]:
- """
- Converts the output of [`DetrForSegmentation`] into instance segmentation predictions. Only supports PyTorch.
-
- Args:
- outputs ([`DetrForSegmentation`]):
- Raw outputs of the model.
- threshold (`float`, *optional*, defaults to 0.5):
- The probability score threshold to keep predicted instance masks.
- mask_threshold (`float`, *optional*, defaults to 0.5):
- Threshold to use when turning the predicted masks into binary values.
- overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8):
- The overlap mask area threshold to merge or discard small disconnected parts within each binary
- instance mask.
- target_sizes (`List[Tuple]`, *optional*):
- List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested
- final size (height, width) of each prediction. If unset, predictions will not be resized.
- return_coco_annotation (`bool`, *optional*):
- Defaults to `False`. If set to `True`, segmentation maps are returned in COCO run-length encoding (RLE)
- format.
- Returns:
- `List[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys:
- - **segmentation** -- A tensor of shape `(height, width)` where each pixel represents a `segment_id` or
- `List[List]` run-length encoding (RLE) of the segmentation map if return_coco_annotation is set to
- `True`. Set to `None` if no mask if found above `threshold`.
- - **segments_info** -- A dictionary that contains additional information on each segment.
- - **id** -- An integer representing the `segment_id`.
- - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`.
- - **score** -- Prediction score of segment with `segment_id`.
- """
- class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1]
- masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width]
-
- batch_size = class_queries_logits.shape[0]
- num_labels = class_queries_logits.shape[-1] - 1
-
- mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width]
-
- # Predicted label and score of each query (batch_size, num_queries)
- pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1)
-
- # Loop over items in batch size
- results: List[Dict[str, TensorType]] = []
-
- for i in range(batch_size):
- mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects(
- mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels
- )
-
- # No mask found
- if mask_probs_item.shape[0] <= 0:
- height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:]
- segmentation = torch.zeros((height, width)) - 1
- results.append({"segmentation": segmentation, "segments_info": []})
- continue
-
- # Get segmentation map and segment information of batch item
- target_size = target_sizes[i] if target_sizes is not None else None
- segmentation, segments = compute_segments(
- mask_probs=mask_probs_item,
- pred_scores=pred_scores_item,
- pred_labels=pred_labels_item,
- mask_threshold=mask_threshold,
- overlap_mask_area_threshold=overlap_mask_area_threshold,
- label_ids_to_fuse=[],
- target_size=target_size,
- )
-
- # Return segmentation map in run-length encoding (RLE) format
- if return_coco_annotation:
- segmentation = convert_segmentation_to_rle(segmentation)
-
- results.append({"segmentation": segmentation, "segments_info": segments})
- return results
-
- # inspired by https://github.com/facebookresearch/detr/blob/master/models/segmentation.py#L241
- def post_process_panoptic_segmentation(
- self,
- outputs,
- threshold: float = 0.5,
- mask_threshold: float = 0.5,
- overlap_mask_area_threshold: float = 0.8,
- label_ids_to_fuse: Optional[Set[int]] = None,
- target_sizes: Optional[List[Tuple[int, int]]] = None,
- ) -> List[Dict]:
- """
- Converts the output of [`DetrForSegmentation`] into image panoptic segmentation predictions. Only supports
- PyTorch.
-
- Args:
- outputs ([`DetrForSegmentation`]):
- The outputs from [`DetrForSegmentation`].
- threshold (`float`, *optional*, defaults to 0.5):
- The probability score threshold to keep predicted instance masks.
- mask_threshold (`float`, *optional*, defaults to 0.5):
- Threshold to use when turning the predicted masks into binary values.
- overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8):
- The overlap mask area threshold to merge or discard small disconnected parts within each binary
- instance mask.
- label_ids_to_fuse (`Set[int]`, *optional*):
- The labels in this state will have all their instances be fused together. For instance we could say
- there can only be one sky in an image, but several persons, so the label ID for sky would be in that
- set, but not the one for person.
- target_sizes (`List[Tuple]`, *optional*):
- List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested
- final size (height, width) of each prediction in batch. If unset, predictions will not be resized.
- Returns:
- `List[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys:
- - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id` or
- `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized to
- the corresponding `target_sizes` entry.
- - **segments_info** -- A dictionary that contains additional information on each segment.
- - **id** -- an integer representing the `segment_id`.
- - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`.
- - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise.
- Multiple instances of the same class / label were fused and assigned a single `segment_id`.
- - **score** -- Prediction score of segment with `segment_id`.
- """
-
- if label_ids_to_fuse is None:
- logger.warning_once("`label_ids_to_fuse` unset. No instance will be fused.")
- label_ids_to_fuse = set()
-
- class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1]
- masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width]
-
- batch_size = class_queries_logits.shape[0]
- num_labels = class_queries_logits.shape[-1] - 1
-
- mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width]
-
- # Predicted label and score of each query (batch_size, num_queries)
- pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1)
-
- # Loop over items in batch size
- results: List[Dict[str, TensorType]] = []
-
- for i in range(batch_size):
- mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects(
- mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels
- )
-
- # No mask found
- if mask_probs_item.shape[0] <= 0:
- height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:]
- segmentation = torch.zeros((height, width)) - 1
- results.append({"segmentation": segmentation, "segments_info": []})
- continue
-
- # Get segmentation map and segment information of batch item
- target_size = target_sizes[i] if target_sizes is not None else None
- segmentation, segments = compute_segments(
- mask_probs=mask_probs_item,
- pred_scores=pred_scores_item,
- pred_labels=pred_labels_item,
- mask_threshold=mask_threshold,
- overlap_mask_area_threshold=overlap_mask_area_threshold,
- label_ids_to_fuse=label_ids_to_fuse,
- target_size=target_size,
- )
-
- results.append({"segmentation": segmentation, "segments_info": segments})
- return results
diff --git a/transformers/models/detr/modeling_detr.py b/transformers/models/detr/modeling_detr.py
deleted file mode 100644
index d7fcdfc5bc7e8393048621bcec39198c34b447c6..0000000000000000000000000000000000000000
--- a/transformers/models/detr/modeling_detr.py
+++ /dev/null
@@ -1,2451 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Facebook AI Research The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DETR model."""
-
-
-import math
-from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple, Union
-
-import torch
-from torch import Tensor, nn
-
-from ...activations import ACT2FN
-from ...modeling_attn_mask_utils import _prepare_4d_attention_mask
-from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithCrossAttentions, Seq2SeqModelOutput
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_accelerate_available,
- is_scipy_available,
- is_timm_available,
- is_vision_available,
- logging,
- replace_return_docstrings,
- requires_backends,
-)
-from ...utils.backbone_utils import load_backbone
-from .configuration_detr import DetrConfig
-
-
-if is_accelerate_available():
- from accelerate import PartialState
- from accelerate.utils import reduce
-
-if is_scipy_available():
- from scipy.optimize import linear_sum_assignment
-
-if is_timm_available():
- from timm import create_model
-
-if is_vision_available():
- from transformers.image_transforms import center_to_corners_format
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "DetrConfig"
-_CHECKPOINT_FOR_DOC = "facebook/detr-resnet-50"
-
-
-from ..deprecated._archive_maps import DETR_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class DetrDecoderOutput(BaseModelOutputWithCrossAttentions):
- """
- Base class for outputs of the DETR decoder. This class adds one attribute to BaseModelOutputWithCrossAttentions,
- namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them
- gone through a layernorm. This is useful when training the model with auxiliary decoding losses.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
- the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`):
- Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a
- layernorm.
- """
-
- intermediate_hidden_states: Optional[torch.FloatTensor] = None
-
-
-@dataclass
-class DetrModelOutput(Seq2SeqModelOutput):
- """
- Base class for outputs of the DETR encoder-decoder model. This class adds one attribute to Seq2SeqModelOutput,
- namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them
- gone through a layernorm. This is useful when training the model with auxiliary decoding losses.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the decoder at the output of each
- layer plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the encoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, sequence_length, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`):
- Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a
- layernorm.
- """
-
- intermediate_hidden_states: Optional[torch.FloatTensor] = None
-
-
-@dataclass
-class DetrObjectDetectionOutput(ModelOutput):
- """
- Output type of [`DetrForObjectDetection`].
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
- Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
- bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
- scale-invariant IoU loss.
- loss_dict (`Dict`, *optional*):
- A dictionary containing the individual losses. Useful for logging.
- logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
- Classification logits (including no-object) for all queries.
- pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
- values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
- possible padding). You can use [`~DetrImageProcessor.post_process_object_detection`] to retrieve the
- unnormalized bounding boxes.
- auxiliary_outputs (`list[Dict]`, *optional*):
- Optional, only returned when auxilary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
- and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
- `pred_boxes`) for each decoder layer.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the decoder at the output of each
- layer plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the encoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- loss_dict: Optional[Dict] = None
- logits: torch.FloatTensor = None
- pred_boxes: torch.FloatTensor = None
- auxiliary_outputs: Optional[List[Dict]] = None
- last_hidden_state: Optional[torch.FloatTensor] = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class DetrSegmentationOutput(ModelOutput):
- """
- Output type of [`DetrForSegmentation`].
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
- Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
- bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
- scale-invariant IoU loss.
- loss_dict (`Dict`, *optional*):
- A dictionary containing the individual losses. Useful for logging.
- logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
- Classification logits (including no-object) for all queries.
- pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
- Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
- values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
- possible padding). You can use [`~DetrImageProcessor.post_process_object_detection`] to retrieve the
- unnormalized bounding boxes.
- pred_masks (`torch.FloatTensor` of shape `(batch_size, num_queries, height/4, width/4)`):
- Segmentation masks logits for all queries. See also
- [`~DetrImageProcessor.post_process_semantic_segmentation`] or
- [`~DetrImageProcessor.post_process_instance_segmentation`]
- [`~DetrImageProcessor.post_process_panoptic_segmentation`] to evaluate semantic, instance and panoptic
- segmentation masks respectively.
- auxiliary_outputs (`list[Dict]`, *optional*):
- Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
- and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
- `pred_boxes`) for each decoder layer.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the decoder of the model.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the decoder at the output of each
- layer plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
- used to compute the weighted average in the cross-attention heads.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
- layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights of the encoder, after the attention softmax, used to compute the
- weighted average in the self-attention heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- loss_dict: Optional[Dict] = None
- logits: torch.FloatTensor = None
- pred_boxes: torch.FloatTensor = None
- pred_masks: torch.FloatTensor = None
- auxiliary_outputs: Optional[List[Dict]] = None
- last_hidden_state: Optional[torch.FloatTensor] = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-# BELOW: utilities copied from
-# https://github.com/facebookresearch/detr/blob/master/backbone.py
-class DetrFrozenBatchNorm2d(nn.Module):
- """
- BatchNorm2d where the batch statistics and the affine parameters are fixed.
-
- Copy-paste from torchvision.misc.ops with added eps before rqsrt, without which any other models than
- torchvision.models.resnet[18,34,50,101] produce nans.
- """
-
- def __init__(self, n):
- super().__init__()
- self.register_buffer("weight", torch.ones(n))
- self.register_buffer("bias", torch.zeros(n))
- self.register_buffer("running_mean", torch.zeros(n))
- self.register_buffer("running_var", torch.ones(n))
-
- def _load_from_state_dict(
- self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
- ):
- num_batches_tracked_key = prefix + "num_batches_tracked"
- if num_batches_tracked_key in state_dict:
- del state_dict[num_batches_tracked_key]
-
- super()._load_from_state_dict(
- state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
- )
-
- def forward(self, x):
- # move reshapes to the beginning
- # to make it user-friendly
- weight = self.weight.reshape(1, -1, 1, 1)
- bias = self.bias.reshape(1, -1, 1, 1)
- running_var = self.running_var.reshape(1, -1, 1, 1)
- running_mean = self.running_mean.reshape(1, -1, 1, 1)
- epsilon = 1e-5
- scale = weight * (running_var + epsilon).rsqrt()
- bias = bias - running_mean * scale
- return x * scale + bias
-
-
-def replace_batch_norm(model):
- r"""
- Recursively replace all `torch.nn.BatchNorm2d` with `DetrFrozenBatchNorm2d`.
-
- Args:
- model (torch.nn.Module):
- input model
- """
- for name, module in model.named_children():
- if isinstance(module, nn.BatchNorm2d):
- new_module = DetrFrozenBatchNorm2d(module.num_features)
-
- if not module.weight.device == torch.device("meta"):
- new_module.weight.data.copy_(module.weight)
- new_module.bias.data.copy_(module.bias)
- new_module.running_mean.data.copy_(module.running_mean)
- new_module.running_var.data.copy_(module.running_var)
-
- model._modules[name] = new_module
-
- if len(list(module.children())) > 0:
- replace_batch_norm(module)
-
-
-class DetrConvEncoder(nn.Module):
- """
- Convolutional backbone, using either the AutoBackbone API or one from the timm library.
-
- nn.BatchNorm2d layers are replaced by DetrFrozenBatchNorm2d as defined above.
-
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.config = config
-
- if config.use_timm_backbone:
- requires_backends(self, ["timm"])
- kwargs = {}
- if config.dilation:
- kwargs["output_stride"] = 16
- backbone = create_model(
- config.backbone,
- pretrained=config.use_pretrained_backbone,
- features_only=True,
- out_indices=(1, 2, 3, 4),
- in_chans=config.num_channels,
- **kwargs,
- )
- else:
- backbone = load_backbone(config)
-
- # replace batch norm by frozen batch norm
- with torch.no_grad():
- replace_batch_norm(backbone)
- self.model = backbone
- self.intermediate_channel_sizes = (
- self.model.feature_info.channels() if config.use_timm_backbone else self.model.channels
- )
-
- backbone_model_type = config.backbone if config.use_timm_backbone else config.backbone_config.model_type
- if "resnet" in backbone_model_type:
- for name, parameter in self.model.named_parameters():
- if config.use_timm_backbone:
- if "layer2" not in name and "layer3" not in name and "layer4" not in name:
- parameter.requires_grad_(False)
- else:
- if "stage.1" not in name and "stage.2" not in name and "stage.3" not in name:
- parameter.requires_grad_(False)
-
- def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
- # send pixel_values through the model to get list of feature maps
- features = self.model(pixel_values) if self.config.use_timm_backbone else self.model(pixel_values).feature_maps
-
- out = []
- for feature_map in features:
- # downsample pixel_mask to match shape of corresponding feature_map
- mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
- out.append((feature_map, mask))
- return out
-
-
-class DetrConvModel(nn.Module):
- """
- This module adds 2D position embeddings to all intermediate feature maps of the convolutional encoder.
- """
-
- def __init__(self, conv_encoder, position_embedding):
- super().__init__()
- self.conv_encoder = conv_encoder
- self.position_embedding = position_embedding
-
- def forward(self, pixel_values, pixel_mask):
- # send pixel_values and pixel_mask through backbone to get list of (feature_map, pixel_mask) tuples
- out = self.conv_encoder(pixel_values, pixel_mask)
- pos = []
- for feature_map, mask in out:
- # position encoding
- pos.append(self.position_embedding(feature_map, mask).to(feature_map.dtype))
-
- return out, pos
-
-
-class DetrSinePositionEmbedding(nn.Module):
- """
- This is a more standard version of the position embedding, very similar to the one used by the Attention is all you
- need paper, generalized to work on images.
- """
-
- def __init__(self, embedding_dim=64, temperature=10000, normalize=False, scale=None):
- super().__init__()
- self.embedding_dim = embedding_dim
- self.temperature = temperature
- self.normalize = normalize
- if scale is not None and normalize is False:
- raise ValueError("normalize should be True if scale is passed")
- if scale is None:
- scale = 2 * math.pi
- self.scale = scale
-
- def forward(self, pixel_values, pixel_mask):
- if pixel_mask is None:
- raise ValueError("No pixel mask provided")
- y_embed = pixel_mask.cumsum(1, dtype=torch.float32)
- x_embed = pixel_mask.cumsum(2, dtype=torch.float32)
- if self.normalize:
- y_embed = y_embed / (y_embed[:, -1:, :] + 1e-6) * self.scale
- x_embed = x_embed / (x_embed[:, :, -1:] + 1e-6) * self.scale
-
- dim_t = torch.arange(self.embedding_dim, dtype=torch.int64, device=pixel_values.device).float()
- dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.embedding_dim)
-
- pos_x = x_embed[:, :, :, None] / dim_t
- pos_y = y_embed[:, :, :, None] / dim_t
- pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
- return pos
-
-
-class DetrLearnedPositionEmbedding(nn.Module):
- """
- This module learns positional embeddings up to a fixed maximum size.
- """
-
- def __init__(self, embedding_dim=256):
- super().__init__()
- self.row_embeddings = nn.Embedding(50, embedding_dim)
- self.column_embeddings = nn.Embedding(50, embedding_dim)
-
- def forward(self, pixel_values, pixel_mask=None):
- height, width = pixel_values.shape[-2:]
- width_values = torch.arange(width, device=pixel_values.device)
- height_values = torch.arange(height, device=pixel_values.device)
- x_emb = self.column_embeddings(width_values)
- y_emb = self.row_embeddings(height_values)
- pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1)
- pos = pos.permute(2, 0, 1)
- pos = pos.unsqueeze(0)
- pos = pos.repeat(pixel_values.shape[0], 1, 1, 1)
- return pos
-
-
-def build_position_encoding(config):
- n_steps = config.d_model // 2
- if config.position_embedding_type == "sine":
- # TODO find a better way of exposing other arguments
- position_embedding = DetrSinePositionEmbedding(n_steps, normalize=True)
- elif config.position_embedding_type == "learned":
- position_embedding = DetrLearnedPositionEmbedding(n_steps)
- else:
- raise ValueError(f"Not supported {config.position_embedding_type}")
-
- return position_embedding
-
-
-class DetrAttention(nn.Module):
- """
- Multi-headed attention from 'Attention Is All You Need' paper.
-
- Here, we add position embeddings to the queries and keys (as explained in the DETR paper).
- """
-
- def __init__(
- self,
- embed_dim: int,
- num_heads: int,
- dropout: float = 0.0,
- bias: bool = True,
- ):
- super().__init__()
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.dropout = dropout
- self.head_dim = embed_dim // num_heads
- if self.head_dim * num_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {num_heads})."
- )
- self.scaling = self.head_dim**-0.5
-
- self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
-
- def _shape(self, tensor: torch.Tensor, seq_len: int, batch_size: int):
- return tensor.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def with_pos_embed(self, tensor: torch.Tensor, object_queries: Optional[Tensor], **kwargs):
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- return tensor if object_queries is None else tensor + object_queries
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- object_queries: Optional[torch.Tensor] = None,
- key_value_states: Optional[torch.Tensor] = None,
- spatial_position_embeddings: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- **kwargs,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel"""
-
- position_embeddings = kwargs.pop("position_ebmeddings", None)
- key_value_position_embeddings = kwargs.pop("key_value_position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if key_value_position_embeddings is not None and spatial_position_embeddings is not None:
- raise ValueError(
- "Cannot specify both key_value_position_embeddings and spatial_position_embeddings. Please use just spatial_position_embeddings"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- if key_value_position_embeddings is not None:
- logger.warning_once(
- "key_value_position_embeddings has been deprecated and will be removed in v4.34. Please use spatial_position_embeddings instead"
- )
- spatial_position_embeddings = key_value_position_embeddings
-
- # if key_value_states are provided this layer is used as a cross-attention layer
- # for the decoder
- is_cross_attention = key_value_states is not None
- batch_size, target_len, embed_dim = hidden_states.size()
-
- # add position embeddings to the hidden states before projecting to queries and keys
- if object_queries is not None:
- hidden_states_original = hidden_states
- hidden_states = self.with_pos_embed(hidden_states, object_queries)
-
- # add key-value position embeddings to the key value states
- if spatial_position_embeddings is not None:
- key_value_states_original = key_value_states
- key_value_states = self.with_pos_embed(key_value_states, spatial_position_embeddings)
-
- # get query proj
- query_states = self.q_proj(hidden_states) * self.scaling
- # get key, value proj
- if is_cross_attention:
- # cross_attentions
- key_states = self._shape(self.k_proj(key_value_states), -1, batch_size)
- value_states = self._shape(self.v_proj(key_value_states_original), -1, batch_size)
- else:
- # self_attention
- key_states = self._shape(self.k_proj(hidden_states), -1, batch_size)
- value_states = self._shape(self.v_proj(hidden_states_original), -1, batch_size)
-
- proj_shape = (batch_size * self.num_heads, -1, self.head_dim)
- query_states = self._shape(query_states, target_len, batch_size).view(*proj_shape)
- key_states = key_states.view(*proj_shape)
- value_states = value_states.view(*proj_shape)
-
- source_len = key_states.size(1)
-
- attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-
- if attn_weights.size() != (batch_size * self.num_heads, target_len, source_len):
- raise ValueError(
- f"Attention weights should be of size {(batch_size * self.num_heads, target_len, source_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- if attention_mask is not None:
- if attention_mask.size() != (batch_size, 1, target_len, source_len):
- raise ValueError(
- f"Attention mask should be of size {(batch_size, 1, target_len, source_len)}, but is"
- f" {attention_mask.size()}"
- )
- attn_weights = attn_weights.view(batch_size, self.num_heads, target_len, source_len) + attention_mask
- attn_weights = attn_weights.view(batch_size * self.num_heads, target_len, source_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if output_attentions:
- # this operation is a bit awkward, but it's required to
- # make sure that attn_weights keeps its gradient.
- # In order to do so, attn_weights have to reshaped
- # twice and have to be reused in the following
- attn_weights_reshaped = attn_weights.view(batch_size, self.num_heads, target_len, source_len)
- attn_weights = attn_weights_reshaped.view(batch_size * self.num_heads, target_len, source_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-
- attn_output = torch.bmm(attn_probs, value_states)
-
- if attn_output.size() != (batch_size * self.num_heads, target_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(batch_size, self.num_heads, target_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.view(batch_size, self.num_heads, target_len, self.head_dim)
- attn_output = attn_output.transpose(1, 2)
- attn_output = attn_output.reshape(batch_size, target_len, embed_dim)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped
-
-
-class DetrEncoderLayer(nn.Module):
- def __init__(self, config: DetrConfig):
- super().__init__()
- self.embed_dim = config.d_model
- self.self_attn = DetrAttention(
- embed_dim=self.embed_dim,
- num_heads=config.encoder_attention_heads,
- dropout=config.attention_dropout,
- )
- self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
- self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
- self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: torch.Tensor,
- object_queries: torch.Tensor = None,
- output_attentions: bool = False,
- **kwargs,
- ):
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`): attention mask of size
- `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
- values.
- object_queries (`torch.FloatTensor`, *optional*):
- Object queries (also called content embeddings), to be added to the hidden states.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- residual = hidden_states
- hidden_states, attn_weights = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- object_queries=object_queries,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- residual = hidden_states
- hidden_states = self.activation_fn(self.fc1(hidden_states))
- hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
-
- hidden_states = self.fc2(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
- hidden_states = residual + hidden_states
- hidden_states = self.final_layer_norm(hidden_states)
-
- if self.training:
- if torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any():
- clamp_value = torch.finfo(hidden_states.dtype).max - 1000
- hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
-
-class DetrDecoderLayer(nn.Module):
- def __init__(self, config: DetrConfig):
- super().__init__()
- self.embed_dim = config.d_model
-
- self.self_attn = DetrAttention(
- embed_dim=self.embed_dim,
- num_heads=config.decoder_attention_heads,
- dropout=config.attention_dropout,
- )
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
-
- self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- self.encoder_attn = DetrAttention(
- self.embed_dim,
- config.decoder_attention_heads,
- dropout=config.attention_dropout,
- )
- self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
- self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
- self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- object_queries: Optional[torch.Tensor] = None,
- query_position_embeddings: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- **kwargs,
- ):
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`): attention mask of size
- `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
- values.
- object_queries (`torch.FloatTensor`, *optional*):
- object_queries that are added to the hidden states
- in the cross-attention layer.
- query_position_embeddings (`torch.FloatTensor`, *optional*):
- position embeddings that are added to the queries and keys
- in the self-attention layer.
- encoder_hidden_states (`torch.FloatTensor`):
- cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
- encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
- `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
- values.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- residual = hidden_states
-
- # Self Attention
- hidden_states, self_attn_weights = self.self_attn(
- hidden_states=hidden_states,
- object_queries=query_position_embeddings,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- # Cross-Attention Block
- cross_attn_weights = None
- if encoder_hidden_states is not None:
- residual = hidden_states
-
- hidden_states, cross_attn_weights = self.encoder_attn(
- hidden_states=hidden_states,
- object_queries=query_position_embeddings,
- key_value_states=encoder_hidden_states,
- attention_mask=encoder_attention_mask,
- spatial_position_embeddings=object_queries,
- output_attentions=output_attentions,
- )
-
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.encoder_attn_layer_norm(hidden_states)
-
- # Fully Connected
- residual = hidden_states
- hidden_states = self.activation_fn(self.fc1(hidden_states))
- hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
- hidden_states = self.fc2(hidden_states)
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
- hidden_states = residual + hidden_states
- hidden_states = self.final_layer_norm(hidden_states)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights, cross_attn_weights)
-
- return outputs
-
-
-class DetrClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, input_dim: int, inner_dim: int, num_classes: int, pooler_dropout: float):
- super().__init__()
- self.dense = nn.Linear(input_dim, inner_dim)
- self.dropout = nn.Dropout(p=pooler_dropout)
- self.out_proj = nn.Linear(inner_dim, num_classes)
-
- def forward(self, hidden_states: torch.Tensor):
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.dense(hidden_states)
- hidden_states = torch.tanh(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.out_proj(hidden_states)
- return hidden_states
-
-
-class DetrPreTrainedModel(PreTrainedModel):
- config_class = DetrConfig
- base_model_prefix = "model"
- main_input_name = "pixel_values"
- _no_split_modules = [r"DetrConvEncoder", r"DetrEncoderLayer", r"DetrDecoderLayer"]
-
- def _init_weights(self, module):
- std = self.config.init_std
- xavier_std = self.config.init_xavier_std
-
- if isinstance(module, DetrMHAttentionMap):
- nn.init.zeros_(module.k_linear.bias)
- nn.init.zeros_(module.q_linear.bias)
- nn.init.xavier_uniform_(module.k_linear.weight, gain=xavier_std)
- nn.init.xavier_uniform_(module.q_linear.weight, gain=xavier_std)
- elif isinstance(module, DetrLearnedPositionEmbedding):
- nn.init.uniform_(module.row_embeddings.weight)
- nn.init.uniform_(module.column_embeddings.weight)
- if isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
-
-DETR_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`DetrConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DETR_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Padding will be ignored by default should you provide it.
-
- Pixel values can be obtained using [`AutoImageProcessor`]. See [`DetrImageProcessor.__call__`] for details.
-
- pixel_mask (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
- Mask to avoid performing attention on padding pixel values. Mask values selected in `[0, 1]`:
-
- - 1 for pixels that are real (i.e. **not masked**),
- - 0 for pixels that are padding (i.e. **masked**).
-
- [What are attention masks?](../glossary#attention-mask)
-
- decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*):
- Not used by default. Can be used to mask object queries.
- encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*):
- Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
- `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of
- hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you
- can choose to directly pass a flattened representation of an image.
- decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
- Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an
- embedded representation.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class DetrEncoder(DetrPreTrainedModel):
- """
- Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
- [`DetrEncoderLayer`].
-
- The encoder updates the flattened feature map through multiple self-attention layers.
-
- Small tweak for DETR:
-
- - object_queries are added to the forward pass.
-
- Args:
- config: DetrConfig
- """
-
- def __init__(self, config: DetrConfig):
- super().__init__(config)
-
- self.dropout = config.dropout
- self.layerdrop = config.encoder_layerdrop
-
- self.layers = nn.ModuleList([DetrEncoderLayer(config) for _ in range(config.encoder_layers)])
-
- # in the original DETR, no layernorm is used at the end of the encoder, as "normalize_before" is set to False by default
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def forward(
- self,
- inputs_embeds=None,
- attention_mask=None,
- object_queries=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- **kwargs,
- ):
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Flattened feature map (output of the backbone + projection layer) that is passed to the encoder.
-
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`:
-
- - 1 for pixel features that are real (i.e. **not masked**),
- - 0 for pixel features that are padding (i.e. **masked**).
-
- [What are attention masks?](../glossary#attention-mask)
-
- object_queries (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Object queries that are added to the queries in each self-attention layer.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- """
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- hidden_states = inputs_embeds
- hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
- # expand attention_mask
- if attention_mask is not None:
- # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len]
- attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
-
- encoder_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
- for i, encoder_layer in enumerate(self.layers):
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
- # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
- to_drop = False
- if self.training:
- dropout_probability = torch.rand([])
- if dropout_probability < self.layerdrop: # skip the layer
- to_drop = True
-
- if to_drop:
- layer_outputs = (None, None)
- else:
- # we add object_queries as extra input to the encoder_layer
- layer_outputs = encoder_layer(
- hidden_states,
- attention_mask,
- object_queries=object_queries,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
- )
-
-
-class DetrDecoder(DetrPreTrainedModel):
- """
- Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DetrDecoderLayer`].
-
- The decoder updates the query embeddings through multiple self-attention and cross-attention layers.
-
- Some small tweaks for DETR:
-
- - object_queries and query_position_embeddings are added to the forward pass.
- - if self.config.auxiliary_loss is set to True, also returns a stack of activations from all decoding layers.
-
- Args:
- config: DetrConfig
- """
-
- def __init__(self, config: DetrConfig):
- super().__init__(config)
- self.dropout = config.dropout
- self.layerdrop = config.decoder_layerdrop
-
- self.layers = nn.ModuleList([DetrDecoderLayer(config) for _ in range(config.decoder_layers)])
- # in DETR, the decoder uses layernorm after the last decoder layer output
- self.layernorm = nn.LayerNorm(config.d_model)
-
- self.gradient_checkpointing = False
- # Initialize weights and apply final processing
- self.post_init()
-
- def forward(
- self,
- inputs_embeds=None,
- attention_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- object_queries=None,
- query_position_embeddings=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- **kwargs,
- ):
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- The query embeddings that are passed into the decoder.
-
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on certain queries. Mask values selected in `[0, 1]`:
-
- - 1 for queries that are **not masked**,
- - 0 for queries that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
- of the decoder.
- encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
- Mask to avoid performing cross-attention on padding pixel_values of the encoder. Mask values selected
- in `[0, 1]`:
-
- - 1 for pixels that are real (i.e. **not masked**),
- - 0 for pixels that are padding (i.e. **masked**).
-
- object_queries (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Object queries that are added to the queries and keys in each cross-attention layer.
- query_position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
- , *optional*): Position embeddings that are added to the values and keys in each self-attention layer.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- """
- position_embeddings = kwargs.pop("position_embeddings", None)
-
- if kwargs:
- raise ValueError(f"Unexpected arguments {kwargs.keys()}")
-
- if position_embeddings is not None and object_queries is not None:
- raise ValueError(
- "Cannot specify both position_embeddings and object_queries. Please use just object_queries"
- )
-
- if position_embeddings is not None:
- logger.warning_once(
- "position_embeddings has been deprecated and will be removed in v4.34. Please use object_queries instead"
- )
- object_queries = position_embeddings
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if inputs_embeds is not None:
- hidden_states = inputs_embeds
- input_shape = inputs_embeds.size()[:-1]
-
- combined_attention_mask = None
-
- if attention_mask is not None and combined_attention_mask is not None:
- # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len]
- combined_attention_mask = combined_attention_mask + _prepare_4d_attention_mask(
- attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
- )
-
- # expand encoder attention mask
- if encoder_hidden_states is not None and encoder_attention_mask is not None:
- # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len]
- encoder_attention_mask = _prepare_4d_attention_mask(
- encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
- )
-
- # optional intermediate hidden states
- intermediate = () if self.config.auxiliary_loss else None
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
-
- for idx, decoder_layer in enumerate(self.layers):
- # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
- if self.training:
- dropout_probability = torch.rand([])
- if dropout_probability < self.layerdrop:
- continue
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- decoder_layer.__call__,
- hidden_states,
- combined_attention_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- None,
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- attention_mask=combined_attention_mask,
- object_queries=object_queries,
- query_position_embeddings=query_position_embeddings,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- if self.config.auxiliary_loss:
- hidden_states = self.layernorm(hidden_states)
- intermediate += (hidden_states,)
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- if encoder_hidden_states is not None:
- all_cross_attentions += (layer_outputs[2],)
-
- # finally, apply layernorm
- hidden_states = self.layernorm(hidden_states)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- # stack intermediate decoder activations
- if self.config.auxiliary_loss:
- intermediate = torch.stack(intermediate)
-
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, all_hidden_states, all_self_attns, all_cross_attentions, intermediate]
- if v is not None
- )
- return DetrDecoderOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- cross_attentions=all_cross_attentions,
- intermediate_hidden_states=intermediate,
- )
-
-
-@add_start_docstrings(
- """
- The bare DETR Model (consisting of a backbone and encoder-decoder Transformer) outputting raw hidden-states without
- any specific head on top.
- """,
- DETR_START_DOCSTRING,
-)
-class DetrModel(DetrPreTrainedModel):
- def __init__(self, config: DetrConfig):
- super().__init__(config)
-
- # Create backbone + positional encoding
- backbone = DetrConvEncoder(config)
- object_queries = build_position_encoding(config)
- self.backbone = DetrConvModel(backbone, object_queries)
-
- # Create projection layer
- self.input_projection = nn.Conv2d(backbone.intermediate_channel_sizes[-1], config.d_model, kernel_size=1)
-
- self.query_position_embeddings = nn.Embedding(config.num_queries, config.d_model)
-
- self.encoder = DetrEncoder(config)
- self.decoder = DetrDecoder(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_encoder(self):
- return self.encoder
-
- def get_decoder(self):
- return self.decoder
-
- def freeze_backbone(self):
- for name, param in self.backbone.conv_encoder.model.named_parameters():
- param.requires_grad_(False)
-
- def unfreeze_backbone(self):
- for name, param in self.backbone.conv_encoder.model.named_parameters():
- param.requires_grad_(True)
-
- @add_start_docstrings_to_model_forward(DETR_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DetrModelOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.FloatTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], DetrModelOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, DetrModel
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")
- >>> model = DetrModel.from_pretrained("facebook/detr-resnet-50")
-
- >>> # prepare image for the model
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> # forward pass
- >>> outputs = model(**inputs)
-
- >>> # the last hidden states are the final query embeddings of the Transformer decoder
- >>> # these are of shape (batch_size, num_queries, hidden_size)
- >>> last_hidden_states = outputs.last_hidden_state
- >>> list(last_hidden_states.shape)
- [1, 100, 256]
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- batch_size, num_channels, height, width = pixel_values.shape
- device = pixel_values.device
-
- if pixel_mask is None:
- pixel_mask = torch.ones(((batch_size, height, width)), device=device)
-
- # First, sent pixel_values + pixel_mask through Backbone to obtain the features
- # pixel_values should be of shape (batch_size, num_channels, height, width)
- # pixel_mask should be of shape (batch_size, height, width)
- features, object_queries_list = self.backbone(pixel_values, pixel_mask)
-
- # get final feature map and downsampled mask
- feature_map, mask = features[-1]
-
- if mask is None:
- raise ValueError("Backbone does not return downsampled pixel mask")
-
- # Second, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
- projected_feature_map = self.input_projection(feature_map)
-
- # Third, flatten the feature map + position embeddings of shape NxCxHxW to NxCxHW, and permute it to NxHWxC
- # In other words, turn their shape into (batch_size, sequence_length, hidden_size)
- flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1)
- object_queries = object_queries_list[-1].flatten(2).permute(0, 2, 1)
-
- flattened_mask = mask.flatten(1)
-
- # Fourth, sent flattened_features + flattened_mask + position embeddings through encoder
- # flattened_features is a Tensor of shape (batch_size, heigth*width, hidden_size)
- # flattened_mask is a Tensor of shape (batch_size, heigth*width)
- if encoder_outputs is None:
- encoder_outputs = self.encoder(
- inputs_embeds=flattened_features,
- attention_mask=flattened_mask,
- object_queries=object_queries,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
- elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
- encoder_outputs = BaseModelOutput(
- last_hidden_state=encoder_outputs[0],
- hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
- attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
- )
-
- # Fifth, sent query embeddings + object_queries through the decoder (which is conditioned on the encoder output)
- query_position_embeddings = self.query_position_embeddings.weight.unsqueeze(0).repeat(batch_size, 1, 1)
- queries = torch.zeros_like(query_position_embeddings)
-
- # decoder outputs consists of (dec_features, dec_hidden, dec_attn)
- decoder_outputs = self.decoder(
- inputs_embeds=queries,
- attention_mask=None,
- object_queries=object_queries,
- query_position_embeddings=query_position_embeddings,
- encoder_hidden_states=encoder_outputs[0],
- encoder_attention_mask=flattened_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- return decoder_outputs + encoder_outputs
-
- return DetrModelOutput(
- last_hidden_state=decoder_outputs.last_hidden_state,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- DETR Model (consisting of a backbone and encoder-decoder Transformer) with object detection heads on top, for tasks
- such as COCO detection.
- """,
- DETR_START_DOCSTRING,
-)
-class DetrForObjectDetection(DetrPreTrainedModel):
- def __init__(self, config: DetrConfig):
- super().__init__(config)
-
- # DETR encoder-decoder model
- self.model = DetrModel(config)
-
- # Object detection heads
- self.class_labels_classifier = nn.Linear(
- config.d_model, config.num_labels + 1
- ) # We add one for the "no object" class
- self.bbox_predictor = DetrMLPPredictionHead(
- input_dim=config.d_model, hidden_dim=config.d_model, output_dim=4, num_layers=3
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- # taken from https://github.com/facebookresearch/detr/blob/master/models/detr.py
- @torch.jit.unused
- def _set_aux_loss(self, outputs_class, outputs_coord):
- # this is a workaround to make torchscript happy, as torchscript
- # doesn't support dictionary with non-homogeneous values, such
- # as a dict having both a Tensor and a list.
- return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
-
- @add_start_docstrings_to_model_forward(DETR_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DetrObjectDetectionOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.FloatTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[List[dict]] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], DetrObjectDetectionOutput]:
- r"""
- labels (`List[Dict]` of len `(batch_size,)`, *optional*):
- Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
- following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
- respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
- in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, DetrForObjectDetection
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")
- >>> model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
- >>> target_sizes = torch.tensor([image.size[::-1]])
- >>> results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[
- ... 0
- ... ]
-
- >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
- ... box = [round(i, 2) for i in box.tolist()]
- ... print(
- ... f"Detected {model.config.id2label[label.item()]} with confidence "
- ... f"{round(score.item(), 3)} at location {box}"
- ... )
- Detected remote with confidence 0.998 at location [40.16, 70.81, 175.55, 117.98]
- Detected remote with confidence 0.996 at location [333.24, 72.55, 368.33, 187.66]
- Detected couch with confidence 0.995 at location [-0.02, 1.15, 639.73, 473.76]
- Detected cat with confidence 0.999 at location [13.24, 52.05, 314.02, 470.93]
- Detected cat with confidence 0.999 at location [345.4, 23.85, 640.37, 368.72]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # First, sent images through DETR base model to obtain encoder + decoder outputs
- outputs = self.model(
- pixel_values,
- pixel_mask=pixel_mask,
- decoder_attention_mask=decoder_attention_mask,
- encoder_outputs=encoder_outputs,
- inputs_embeds=inputs_embeds,
- decoder_inputs_embeds=decoder_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- # class logits + predicted bounding boxes
- logits = self.class_labels_classifier(sequence_output)
- pred_boxes = self.bbox_predictor(sequence_output).sigmoid()
-
- loss, loss_dict, auxiliary_outputs = None, None, None
- if labels is not None:
- # First: create the matcher
- matcher = DetrHungarianMatcher(
- class_cost=self.config.class_cost, bbox_cost=self.config.bbox_cost, giou_cost=self.config.giou_cost
- )
- # Second: create the criterion
- losses = ["labels", "boxes", "cardinality"]
- criterion = DetrLoss(
- matcher=matcher,
- num_classes=self.config.num_labels,
- eos_coef=self.config.eos_coefficient,
- losses=losses,
- )
- criterion.to(self.device)
- # Third: compute the losses, based on outputs and labels
- outputs_loss = {}
- outputs_loss["logits"] = logits
- outputs_loss["pred_boxes"] = pred_boxes
- if self.config.auxiliary_loss:
- intermediate = outputs.intermediate_hidden_states if return_dict else outputs[4]
- outputs_class = self.class_labels_classifier(intermediate)
- outputs_coord = self.bbox_predictor(intermediate).sigmoid()
- auxiliary_outputs = self._set_aux_loss(outputs_class, outputs_coord)
- outputs_loss["auxiliary_outputs"] = auxiliary_outputs
-
- loss_dict = criterion(outputs_loss, labels)
- # Fourth: compute total loss, as a weighted sum of the various losses
- weight_dict = {"loss_ce": 1, "loss_bbox": self.config.bbox_loss_coefficient}
- weight_dict["loss_giou"] = self.config.giou_loss_coefficient
- if self.config.auxiliary_loss:
- aux_weight_dict = {}
- for i in range(self.config.decoder_layers - 1):
- aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
- weight_dict.update(aux_weight_dict)
- loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
-
- if not return_dict:
- if auxiliary_outputs is not None:
- output = (logits, pred_boxes) + auxiliary_outputs + outputs
- else:
- output = (logits, pred_boxes) + outputs
- return ((loss, loss_dict) + output) if loss is not None else output
-
- return DetrObjectDetectionOutput(
- loss=loss,
- loss_dict=loss_dict,
- logits=logits,
- pred_boxes=pred_boxes,
- auxiliary_outputs=auxiliary_outputs,
- last_hidden_state=outputs.last_hidden_state,
- decoder_hidden_states=outputs.decoder_hidden_states,
- decoder_attentions=outputs.decoder_attentions,
- cross_attentions=outputs.cross_attentions,
- encoder_last_hidden_state=outputs.encoder_last_hidden_state,
- encoder_hidden_states=outputs.encoder_hidden_states,
- encoder_attentions=outputs.encoder_attentions,
- )
-
-
-@add_start_docstrings(
- """
- DETR Model (consisting of a backbone and encoder-decoder Transformer) with a segmentation head on top, for tasks
- such as COCO panoptic.
-
- """,
- DETR_START_DOCSTRING,
-)
-class DetrForSegmentation(DetrPreTrainedModel):
- def __init__(self, config: DetrConfig):
- super().__init__(config)
-
- # object detection model
- self.detr = DetrForObjectDetection(config)
-
- # segmentation head
- hidden_size, number_of_heads = config.d_model, config.encoder_attention_heads
- intermediate_channel_sizes = self.detr.model.backbone.conv_encoder.intermediate_channel_sizes
-
- self.mask_head = DetrMaskHeadSmallConv(
- hidden_size + number_of_heads, intermediate_channel_sizes[::-1][-3:], hidden_size
- )
-
- self.bbox_attention = DetrMHAttentionMap(
- hidden_size, hidden_size, number_of_heads, dropout=0.0, std=config.init_xavier_std
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DETR_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DetrSegmentationOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- pixel_mask: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.FloatTensor] = None,
- encoder_outputs: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[List[dict]] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], DetrSegmentationOutput]:
- r"""
- labels (`List[Dict]` of len `(batch_size,)`, *optional*):
- Labels for computing the bipartite matching loss, DICE/F-1 loss and Focal loss. List of dicts, each
- dictionary containing at least the following 3 keys: 'class_labels', 'boxes' and 'masks' (the class labels,
- bounding boxes and segmentation masks of an image in the batch respectively). The class labels themselves
- should be a `torch.LongTensor` of len `(number of bounding boxes in the image,)`, the boxes a
- `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)` and the masks a
- `torch.FloatTensor` of shape `(number of bounding boxes in the image, height, width)`.
-
- Returns:
-
- Examples:
-
- ```python
- >>> import io
- >>> import requests
- >>> from PIL import Image
- >>> import torch
- >>> import numpy
-
- >>> from transformers import AutoImageProcessor, DetrForSegmentation
- >>> from transformers.image_transforms import rgb_to_id
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50-panoptic")
- >>> model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
-
- >>> # prepare image for the model
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> # forward pass
- >>> outputs = model(**inputs)
-
- >>> # Use the `post_process_panoptic_segmentation` method of the `image_processor` to retrieve post-processed panoptic segmentation maps
- >>> # Segmentation results are returned as a list of dictionaries
- >>> result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[(300, 500)])
-
- >>> # A tensor of shape (height, width) where each value denotes a segment id, filled with -1 if no segment is found
- >>> panoptic_seg = result[0]["segmentation"]
- >>> # Get prediction score and segment_id to class_id mapping of each segment
- >>> panoptic_segments_info = result[0]["segments_info"]
- ```"""
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- batch_size, num_channels, height, width = pixel_values.shape
- device = pixel_values.device
-
- if pixel_mask is None:
- pixel_mask = torch.ones((batch_size, height, width), device=device)
-
- # First, get list of feature maps and position embeddings
- features, object_queries_list = self.detr.model.backbone(pixel_values, pixel_mask=pixel_mask)
-
- # Second, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
- feature_map, mask = features[-1]
- batch_size, num_channels, height, width = feature_map.shape
- projected_feature_map = self.detr.model.input_projection(feature_map)
-
- # Third, flatten the feature map + position embeddings of shape NxCxHxW to NxCxHW, and permute it to NxHWxC
- # In other words, turn their shape into (batch_size, sequence_length, hidden_size)
- flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1)
- object_queries = object_queries_list[-1].flatten(2).permute(0, 2, 1)
-
- flattened_mask = mask.flatten(1)
-
- # Fourth, sent flattened_features + flattened_mask + position embeddings through encoder
- # flattened_features is a Tensor of shape (batch_size, heigth*width, hidden_size)
- # flattened_mask is a Tensor of shape (batch_size, heigth*width)
- if encoder_outputs is None:
- encoder_outputs = self.detr.model.encoder(
- inputs_embeds=flattened_features,
- attention_mask=flattened_mask,
- object_queries=object_queries,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
- elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
- encoder_outputs = BaseModelOutput(
- last_hidden_state=encoder_outputs[0],
- hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
- attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
- )
-
- # Fifth, sent query embeddings + position embeddings through the decoder (which is conditioned on the encoder output)
- query_position_embeddings = self.detr.model.query_position_embeddings.weight.unsqueeze(0).repeat(
- batch_size, 1, 1
- )
- queries = torch.zeros_like(query_position_embeddings)
-
- # decoder outputs consists of (dec_features, dec_hidden, dec_attn)
- decoder_outputs = self.detr.model.decoder(
- inputs_embeds=queries,
- attention_mask=None,
- object_queries=object_queries,
- query_position_embeddings=query_position_embeddings,
- encoder_hidden_states=encoder_outputs[0],
- encoder_attention_mask=flattened_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = decoder_outputs[0]
-
- # Sixth, compute logits, pred_boxes and pred_masks
- logits = self.detr.class_labels_classifier(sequence_output)
- pred_boxes = self.detr.bbox_predictor(sequence_output).sigmoid()
-
- memory = encoder_outputs[0].permute(0, 2, 1).view(batch_size, self.config.d_model, height, width)
- mask = flattened_mask.view(batch_size, height, width)
-
- # FIXME h_boxes takes the last one computed, keep this in mind
- # important: we need to reverse the mask, since in the original implementation the mask works reversed
- # bbox_mask is of shape (batch_size, num_queries, number_of_attention_heads in bbox_attention, height/32, width/32)
- bbox_mask = self.bbox_attention(sequence_output, memory, mask=~mask)
-
- seg_masks = self.mask_head(projected_feature_map, bbox_mask, [features[2][0], features[1][0], features[0][0]])
-
- pred_masks = seg_masks.view(batch_size, self.detr.config.num_queries, seg_masks.shape[-2], seg_masks.shape[-1])
-
- loss, loss_dict, auxiliary_outputs = None, None, None
- if labels is not None:
- # First: create the matcher
- matcher = DetrHungarianMatcher(
- class_cost=self.config.class_cost, bbox_cost=self.config.bbox_cost, giou_cost=self.config.giou_cost
- )
- # Second: create the criterion
- losses = ["labels", "boxes", "cardinality", "masks"]
- criterion = DetrLoss(
- matcher=matcher,
- num_classes=self.config.num_labels,
- eos_coef=self.config.eos_coefficient,
- losses=losses,
- )
- criterion.to(self.device)
- # Third: compute the losses, based on outputs and labels
- outputs_loss = {}
- outputs_loss["logits"] = logits
- outputs_loss["pred_boxes"] = pred_boxes
- outputs_loss["pred_masks"] = pred_masks
- if self.config.auxiliary_loss:
- intermediate = decoder_outputs.intermediate_hidden_states if return_dict else decoder_outputs[-1]
- outputs_class = self.detr.class_labels_classifier(intermediate)
- outputs_coord = self.detr.bbox_predictor(intermediate).sigmoid()
- auxiliary_outputs = self.detr._set_aux_loss(outputs_class, outputs_coord)
- outputs_loss["auxiliary_outputs"] = auxiliary_outputs
-
- loss_dict = criterion(outputs_loss, labels)
- # Fourth: compute total loss, as a weighted sum of the various losses
- weight_dict = {"loss_ce": 1, "loss_bbox": self.config.bbox_loss_coefficient}
- weight_dict["loss_giou"] = self.config.giou_loss_coefficient
- weight_dict["loss_mask"] = self.config.mask_loss_coefficient
- weight_dict["loss_dice"] = self.config.dice_loss_coefficient
- if self.config.auxiliary_loss:
- aux_weight_dict = {}
- for i in range(self.config.decoder_layers - 1):
- aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
- weight_dict.update(aux_weight_dict)
- loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
-
- if not return_dict:
- if auxiliary_outputs is not None:
- output = (logits, pred_boxes, pred_masks) + auxiliary_outputs + decoder_outputs + encoder_outputs
- else:
- output = (logits, pred_boxes, pred_masks) + decoder_outputs + encoder_outputs
- return ((loss, loss_dict) + output) if loss is not None else output
-
- return DetrSegmentationOutput(
- loss=loss,
- loss_dict=loss_dict,
- logits=logits,
- pred_boxes=pred_boxes,
- pred_masks=pred_masks,
- auxiliary_outputs=auxiliary_outputs,
- last_hidden_state=decoder_outputs.last_hidden_state,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- )
-
-
-def _expand(tensor, length: int):
- return tensor.unsqueeze(1).repeat(1, int(length), 1, 1, 1).flatten(0, 1)
-
-
-# taken from https://github.com/facebookresearch/detr/blob/master/models/segmentation.py
-class DetrMaskHeadSmallConv(nn.Module):
- """
- Simple convolutional head, using group norm. Upsampling is done using a FPN approach
- """
-
- def __init__(self, dim, fpn_dims, context_dim):
- super().__init__()
-
- if dim % 8 != 0:
- raise ValueError(
- "The hidden_size + number of attention heads must be divisible by 8 as the number of groups in"
- " GroupNorm is set to 8"
- )
-
- inter_dims = [dim, context_dim // 2, context_dim // 4, context_dim // 8, context_dim // 16, context_dim // 64]
-
- self.lay1 = nn.Conv2d(dim, dim, 3, padding=1)
- self.gn1 = nn.GroupNorm(8, dim)
- self.lay2 = nn.Conv2d(dim, inter_dims[1], 3, padding=1)
- self.gn2 = nn.GroupNorm(min(8, inter_dims[1]), inter_dims[1])
- self.lay3 = nn.Conv2d(inter_dims[1], inter_dims[2], 3, padding=1)
- self.gn3 = nn.GroupNorm(min(8, inter_dims[2]), inter_dims[2])
- self.lay4 = nn.Conv2d(inter_dims[2], inter_dims[3], 3, padding=1)
- self.gn4 = nn.GroupNorm(min(8, inter_dims[3]), inter_dims[3])
- self.lay5 = nn.Conv2d(inter_dims[3], inter_dims[4], 3, padding=1)
- self.gn5 = nn.GroupNorm(min(8, inter_dims[4]), inter_dims[4])
- self.out_lay = nn.Conv2d(inter_dims[4], 1, 3, padding=1)
-
- self.dim = dim
-
- self.adapter1 = nn.Conv2d(fpn_dims[0], inter_dims[1], 1)
- self.adapter2 = nn.Conv2d(fpn_dims[1], inter_dims[2], 1)
- self.adapter3 = nn.Conv2d(fpn_dims[2], inter_dims[3], 1)
-
- for m in self.modules():
- if isinstance(m, nn.Conv2d):
- nn.init.kaiming_uniform_(m.weight, a=1)
- nn.init.constant_(m.bias, 0)
-
- def forward(self, x: Tensor, bbox_mask: Tensor, fpns: List[Tensor]):
- # here we concatenate x, the projected feature map, of shape (batch_size, d_model, heigth/32, width/32) with
- # the bbox_mask = the attention maps of shape (batch_size, n_queries, n_heads, height/32, width/32).
- # We expand the projected feature map to match the number of heads.
- x = torch.cat([_expand(x, bbox_mask.shape[1]), bbox_mask.flatten(0, 1)], 1)
-
- x = self.lay1(x)
- x = self.gn1(x)
- x = nn.functional.relu(x)
- x = self.lay2(x)
- x = self.gn2(x)
- x = nn.functional.relu(x)
-
- cur_fpn = self.adapter1(fpns[0])
- if cur_fpn.size(0) != x.size(0):
- cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
- x = cur_fpn + nn.functional.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
- x = self.lay3(x)
- x = self.gn3(x)
- x = nn.functional.relu(x)
-
- cur_fpn = self.adapter2(fpns[1])
- if cur_fpn.size(0) != x.size(0):
- cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
- x = cur_fpn + nn.functional.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
- x = self.lay4(x)
- x = self.gn4(x)
- x = nn.functional.relu(x)
-
- cur_fpn = self.adapter3(fpns[2])
- if cur_fpn.size(0) != x.size(0):
- cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
- x = cur_fpn + nn.functional.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
- x = self.lay5(x)
- x = self.gn5(x)
- x = nn.functional.relu(x)
-
- x = self.out_lay(x)
- return x
-
-
-class DetrMHAttentionMap(nn.Module):
- """This is a 2D attention module, which only returns the attention softmax (no multiplication by value)"""
-
- def __init__(self, query_dim, hidden_dim, num_heads, dropout=0.0, bias=True, std=None):
- super().__init__()
- self.num_heads = num_heads
- self.hidden_dim = hidden_dim
- self.dropout = nn.Dropout(dropout)
-
- self.q_linear = nn.Linear(query_dim, hidden_dim, bias=bias)
- self.k_linear = nn.Linear(query_dim, hidden_dim, bias=bias)
-
- self.normalize_fact = float(hidden_dim / self.num_heads) ** -0.5
-
- def forward(self, q, k, mask: Optional[Tensor] = None):
- q = self.q_linear(q)
- k = nn.functional.conv2d(k, self.k_linear.weight.unsqueeze(-1).unsqueeze(-1), self.k_linear.bias)
- queries_per_head = q.view(q.shape[0], q.shape[1], self.num_heads, self.hidden_dim // self.num_heads)
- keys_per_head = k.view(k.shape[0], self.num_heads, self.hidden_dim // self.num_heads, k.shape[-2], k.shape[-1])
- weights = torch.einsum("bqnc,bnchw->bqnhw", queries_per_head * self.normalize_fact, keys_per_head)
-
- if mask is not None:
- weights.masked_fill_(mask.unsqueeze(1).unsqueeze(1), torch.finfo(weights.dtype).min)
- weights = nn.functional.softmax(weights.flatten(2), dim=-1).view(weights.size())
- weights = self.dropout(weights)
- return weights
-
-
-def dice_loss(inputs, targets, num_boxes):
- """
- Compute the DICE loss, similar to generalized IOU for masks
-
- Args:
- inputs: A float tensor of arbitrary shape.
- The predictions for each example.
- targets: A float tensor with the same shape as inputs. Stores the binary
- classification label for each element in inputs (0 for the negative class and 1 for the positive
- class).
- """
- inputs = inputs.sigmoid()
- inputs = inputs.flatten(1)
- numerator = 2 * (inputs * targets).sum(1)
- denominator = inputs.sum(-1) + targets.sum(-1)
- loss = 1 - (numerator + 1) / (denominator + 1)
- return loss.sum() / num_boxes
-
-
-def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2):
- """
- Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
-
- Args:
- inputs (`torch.FloatTensor` of arbitrary shape):
- The predictions for each example.
- targets (`torch.FloatTensor` with the same shape as `inputs`)
- A tensor storing the binary classification label for each element in the `inputs` (0 for the negative class
- and 1 for the positive class).
- alpha (`float`, *optional*, defaults to `0.25`):
- Optional weighting factor in the range (0,1) to balance positive vs. negative examples.
- gamma (`int`, *optional*, defaults to `2`):
- Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples.
-
- Returns:
- Loss tensor
- """
- prob = inputs.sigmoid()
- ce_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
- # add modulating factor
- p_t = prob * targets + (1 - prob) * (1 - targets)
- loss = ce_loss * ((1 - p_t) ** gamma)
-
- if alpha >= 0:
- alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
- loss = alpha_t * loss
-
- return loss.mean(1).sum() / num_boxes
-
-
-# taken from https://github.com/facebookresearch/detr/blob/master/models/detr.py
-class DetrLoss(nn.Module):
- """
- This class computes the losses for DetrForObjectDetection/DetrForSegmentation. The process happens in two steps: 1)
- we compute hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair
- of matched ground-truth / prediction (supervise class and box).
-
- A note on the `num_classes` argument (copied from original repo in detr.py): "the naming of the `num_classes`
- parameter of the criterion is somewhat misleading. It indeed corresponds to `max_obj_id` + 1, where `max_obj_id` is
- the maximum id for a class in your dataset. For example, COCO has a `max_obj_id` of 90, so we pass `num_classes` to
- be 91. As another example, for a dataset that has a single class with `id` 1, you should pass `num_classes` to be 2
- (`max_obj_id` + 1). For more details on this, check the following discussion
- https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223"
-
-
- Args:
- matcher (`DetrHungarianMatcher`):
- Module able to compute a matching between targets and proposals.
- num_classes (`int`):
- Number of object categories, omitting the special no-object category.
- eos_coef (`float`):
- Relative classification weight applied to the no-object category.
- losses (`List[str]`):
- List of all the losses to be applied. See `get_loss` for a list of all available losses.
- """
-
- def __init__(self, matcher, num_classes, eos_coef, losses):
- super().__init__()
- self.matcher = matcher
- self.num_classes = num_classes
- self.eos_coef = eos_coef
- self.losses = losses
- empty_weight = torch.ones(self.num_classes + 1)
- empty_weight[-1] = self.eos_coef
- self.register_buffer("empty_weight", empty_weight)
-
- # removed logging parameter, which was part of the original implementation
- def loss_labels(self, outputs, targets, indices, num_boxes):
- """
- Classification loss (NLL) targets dicts must contain the key "class_labels" containing a tensor of dim
- [nb_target_boxes]
- """
- if "logits" not in outputs:
- raise KeyError("No logits were found in the outputs")
- source_logits = outputs["logits"]
-
- idx = self._get_source_permutation_idx(indices)
- target_classes_o = torch.cat([t["class_labels"][J] for t, (_, J) in zip(targets, indices)])
- target_classes = torch.full(
- source_logits.shape[:2], self.num_classes, dtype=torch.int64, device=source_logits.device
- )
- target_classes[idx] = target_classes_o
-
- loss_ce = nn.functional.cross_entropy(source_logits.transpose(1, 2), target_classes, self.empty_weight)
- losses = {"loss_ce": loss_ce}
-
- return losses
-
- @torch.no_grad()
- def loss_cardinality(self, outputs, targets, indices, num_boxes):
- """
- Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes.
-
- This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients.
- """
- logits = outputs["logits"]
- device = logits.device
- target_lengths = torch.as_tensor([len(v["class_labels"]) for v in targets], device=device)
- # Count the number of predictions that are NOT "no-object" (which is the last class)
- card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1)
- card_err = nn.functional.l1_loss(card_pred.float(), target_lengths.float())
- losses = {"cardinality_error": card_err}
- return losses
-
- def loss_boxes(self, outputs, targets, indices, num_boxes):
- """
- Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss.
-
- Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes
- are expected in format (center_x, center_y, w, h), normalized by the image size.
- """
- if "pred_boxes" not in outputs:
- raise KeyError("No predicted boxes found in outputs")
- idx = self._get_source_permutation_idx(indices)
- source_boxes = outputs["pred_boxes"][idx]
- target_boxes = torch.cat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0)
-
- loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none")
-
- losses = {}
- losses["loss_bbox"] = loss_bbox.sum() / num_boxes
-
- loss_giou = 1 - torch.diag(
- generalized_box_iou(center_to_corners_format(source_boxes), center_to_corners_format(target_boxes))
- )
- losses["loss_giou"] = loss_giou.sum() / num_boxes
- return losses
-
- def loss_masks(self, outputs, targets, indices, num_boxes):
- """
- Compute the losses related to the masks: the focal loss and the dice loss.
-
- Targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w].
- """
- if "pred_masks" not in outputs:
- raise KeyError("No predicted masks found in outputs")
-
- source_idx = self._get_source_permutation_idx(indices)
- target_idx = self._get_target_permutation_idx(indices)
- source_masks = outputs["pred_masks"]
- source_masks = source_masks[source_idx]
- masks = [t["masks"] for t in targets]
- # TODO use valid to mask invalid areas due to padding in loss
- target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
- target_masks = target_masks.to(source_masks)
- target_masks = target_masks[target_idx]
-
- # upsample predictions to the target size
- source_masks = nn.functional.interpolate(
- source_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False
- )
- source_masks = source_masks[:, 0].flatten(1)
-
- target_masks = target_masks.flatten(1)
- target_masks = target_masks.view(source_masks.shape)
- losses = {
- "loss_mask": sigmoid_focal_loss(source_masks, target_masks, num_boxes),
- "loss_dice": dice_loss(source_masks, target_masks, num_boxes),
- }
- return losses
-
- def _get_source_permutation_idx(self, indices):
- # permute predictions following indices
- batch_idx = torch.cat([torch.full_like(source, i) for i, (source, _) in enumerate(indices)])
- source_idx = torch.cat([source for (source, _) in indices])
- return batch_idx, source_idx
-
- def _get_target_permutation_idx(self, indices):
- # permute targets following indices
- batch_idx = torch.cat([torch.full_like(target, i) for i, (_, target) in enumerate(indices)])
- target_idx = torch.cat([target for (_, target) in indices])
- return batch_idx, target_idx
-
- def get_loss(self, loss, outputs, targets, indices, num_boxes):
- loss_map = {
- "labels": self.loss_labels,
- "cardinality": self.loss_cardinality,
- "boxes": self.loss_boxes,
- "masks": self.loss_masks,
- }
- if loss not in loss_map:
- raise ValueError(f"Loss {loss} not supported")
- return loss_map[loss](outputs, targets, indices, num_boxes)
-
- def forward(self, outputs, targets):
- """
- This performs the loss computation.
-
- Args:
- outputs (`dict`, *optional*):
- Dictionary of tensors, see the output specification of the model for the format.
- targets (`List[dict]`, *optional*):
- List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the
- losses applied, see each loss' doc.
- """
- outputs_without_aux = {k: v for k, v in outputs.items() if k != "auxiliary_outputs"}
-
- # Retrieve the matching between the outputs of the last layer and the targets
- indices = self.matcher(outputs_without_aux, targets)
-
- # Compute the average number of target boxes across all nodes, for normalization purposes
- num_boxes = sum(len(t["class_labels"]) for t in targets)
- num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
- world_size = 1
- if is_accelerate_available():
- if PartialState._shared_state != {}:
- num_boxes = reduce(num_boxes)
- world_size = PartialState().num_processes
- num_boxes = torch.clamp(num_boxes / world_size, min=1).item()
-
- # Compute all the requested losses
- losses = {}
- for loss in self.losses:
- losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))
-
- # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
- if "auxiliary_outputs" in outputs:
- for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]):
- indices = self.matcher(auxiliary_outputs, targets)
- for loss in self.losses:
- if loss == "masks":
- # Intermediate masks losses are too costly to compute, we ignore them.
- continue
- l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes)
- l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
- losses.update(l_dict)
-
- return losses
-
-
-# taken from https://github.com/facebookresearch/detr/blob/master/models/detr.py
-class DetrMLPPredictionHead(nn.Module):
- """
- Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates,
- height and width of a bounding box w.r.t. an image.
-
- Copied from https://github.com/facebookresearch/detr/blob/master/models/detr.py
-
- """
-
- def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
- super().__init__()
- self.num_layers = num_layers
- h = [hidden_dim] * (num_layers - 1)
- self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
-
- def forward(self, x):
- for i, layer in enumerate(self.layers):
- x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
- return x
-
-
-# taken from https://github.com/facebookresearch/detr/blob/master/models/matcher.py
-class DetrHungarianMatcher(nn.Module):
- """
- This class computes an assignment between the targets and the predictions of the network.
-
- For efficiency reasons, the targets don't include the no_object. Because of this, in general, there are more
- predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, while the others are
- un-matched (and thus treated as non-objects).
-
- Args:
- class_cost:
- The relative weight of the classification error in the matching cost.
- bbox_cost:
- The relative weight of the L1 error of the bounding box coordinates in the matching cost.
- giou_cost:
- The relative weight of the giou loss of the bounding box in the matching cost.
- """
-
- def __init__(self, class_cost: float = 1, bbox_cost: float = 1, giou_cost: float = 1):
- super().__init__()
- requires_backends(self, ["scipy"])
-
- self.class_cost = class_cost
- self.bbox_cost = bbox_cost
- self.giou_cost = giou_cost
- if class_cost == 0 and bbox_cost == 0 and giou_cost == 0:
- raise ValueError("All costs of the Matcher can't be 0")
-
- @torch.no_grad()
- def forward(self, outputs, targets):
- """
- Args:
- outputs (`dict`):
- A dictionary that contains at least these entries:
- * "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
- * "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates.
- targets (`List[dict]`):
- A list of targets (len(targets) = batch_size), where each target is a dict containing:
- * "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of
- ground-truth
- objects in the target) containing the class labels
- * "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates.
-
- Returns:
- `List[Tuple]`: A list of size `batch_size`, containing tuples of (index_i, index_j) where:
- - index_i is the indices of the selected predictions (in order)
- - index_j is the indices of the corresponding selected targets (in order)
- For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
- """
- batch_size, num_queries = outputs["logits"].shape[:2]
-
- # We flatten to compute the cost matrices in a batch
- out_prob = outputs["logits"].flatten(0, 1).softmax(-1) # [batch_size * num_queries, num_classes]
- out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4]
-
- # Also concat the target labels and boxes
- target_ids = torch.cat([v["class_labels"] for v in targets])
- target_bbox = torch.cat([v["boxes"] for v in targets])
-
- # Compute the classification cost. Contrary to the loss, we don't use the NLL,
- # but approximate it in 1 - proba[target class].
- # The 1 is a constant that doesn't change the matching, it can be ommitted.
- class_cost = -out_prob[:, target_ids]
-
- # Compute the L1 cost between boxes
- bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)
-
- # Compute the giou cost between boxes
- giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))
-
- # Final cost matrix
- cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
- cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()
-
- sizes = [len(v["boxes"]) for v in targets]
- indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]
- return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
-
-
-# below: bounding box utilities taken from https://github.com/facebookresearch/detr/blob/master/util/box_ops.py
-
-
-def _upcast(t: Tensor) -> Tensor:
- # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
- if t.is_floating_point():
- return t if t.dtype in (torch.float32, torch.float64) else t.float()
- else:
- return t if t.dtype in (torch.int32, torch.int64) else t.int()
-
-
-def box_area(boxes: Tensor) -> Tensor:
- """
- Computes the area of a set of bounding boxes, which are specified by its (x1, y1, x2, y2) coordinates.
-
- Args:
- boxes (`torch.FloatTensor` of shape `(number_of_boxes, 4)`):
- Boxes for which the area will be computed. They are expected to be in (x1, y1, x2, y2) format with `0 <= x1
- < x2` and `0 <= y1 < y2`.
-
- Returns:
- `torch.FloatTensor`: a tensor containing the area for each box.
- """
- boxes = _upcast(boxes)
- return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
-
-
-# modified from torchvision to also return the union
-def box_iou(boxes1, boxes2):
- area1 = box_area(boxes1)
- area2 = box_area(boxes2)
-
- left_top = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
- right_bottom = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
-
- width_height = (right_bottom - left_top).clamp(min=0) # [N,M,2]
- inter = width_height[:, :, 0] * width_height[:, :, 1] # [N,M]
-
- union = area1[:, None] + area2 - inter
-
- iou = inter / union
- return iou, union
-
-
-def generalized_box_iou(boxes1, boxes2):
- """
- Generalized IoU from https://giou.stanford.edu/. The boxes should be in [x0, y0, x1, y1] (corner) format.
-
- Returns:
- `torch.FloatTensor`: a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2)
- """
- # degenerate boxes gives inf / nan results
- # so do an early check
- if not (boxes1[:, 2:] >= boxes1[:, :2]).all():
- raise ValueError(f"boxes1 must be in [x0, y0, x1, y1] (corner) format, but got {boxes1}")
- if not (boxes2[:, 2:] >= boxes2[:, :2]).all():
- raise ValueError(f"boxes2 must be in [x0, y0, x1, y1] (corner) format, but got {boxes2}")
- iou, union = box_iou(boxes1, boxes2)
-
- top_left = torch.min(boxes1[:, None, :2], boxes2[:, :2])
- bottom_right = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
-
- width_height = (bottom_right - top_left).clamp(min=0) # [N,M,2]
- area = width_height[:, :, 0] * width_height[:, :, 1]
-
- return iou - (area - union) / area
-
-
-# below: taken from https://github.com/facebookresearch/detr/blob/master/util/misc.py#L306
-def _max_by_axis(the_list):
- # type: (List[List[int]]) -> List[int]
- maxes = the_list[0]
- for sublist in the_list[1:]:
- for index, item in enumerate(sublist):
- maxes[index] = max(maxes[index], item)
- return maxes
-
-
-class NestedTensor(object):
- def __init__(self, tensors, mask: Optional[Tensor]):
- self.tensors = tensors
- self.mask = mask
-
- def to(self, device):
- cast_tensor = self.tensors.to(device)
- mask = self.mask
- if mask is not None:
- cast_mask = mask.to(device)
- else:
- cast_mask = None
- return NestedTensor(cast_tensor, cast_mask)
-
- def decompose(self):
- return self.tensors, self.mask
-
- def __repr__(self):
- return str(self.tensors)
-
-
-def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
- if tensor_list[0].ndim == 3:
- max_size = _max_by_axis([list(img.shape) for img in tensor_list])
- batch_shape = [len(tensor_list)] + max_size
- batch_size, num_channels, height, width = batch_shape
- dtype = tensor_list[0].dtype
- device = tensor_list[0].device
- tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
- mask = torch.ones((batch_size, height, width), dtype=torch.bool, device=device)
- for img, pad_img, m in zip(tensor_list, tensor, mask):
- pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
- m[: img.shape[1], : img.shape[2]] = False
- else:
- raise ValueError("Only 3-dimensional tensors are supported")
- return NestedTensor(tensor, mask)
diff --git a/transformers/models/dialogpt/__init__.py b/transformers/models/dialogpt/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/transformers/models/dialogpt/__pycache__/__init__.cpython-310.pyc b/transformers/models/dialogpt/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index c13082c95edf2da198682e28c693b2e23319601c..0000000000000000000000000000000000000000
Binary files a/transformers/models/dialogpt/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dialogpt/__pycache__/convert_dialogpt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/dialogpt/__pycache__/convert_dialogpt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index c1fa9049850f87cd01689c667f2c4c2fbf7ce8c6..0000000000000000000000000000000000000000
Binary files a/transformers/models/dialogpt/__pycache__/convert_dialogpt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dialogpt/convert_dialogpt_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/dialogpt/convert_dialogpt_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index fbf34012924b901f3a074d36ed9be7b1fc32913b..0000000000000000000000000000000000000000
--- a/transformers/models/dialogpt/convert_dialogpt_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import os
-
-import torch
-
-from transformers.utils import WEIGHTS_NAME
-
-
-DIALOGPT_MODELS = ["small", "medium", "large"]
-
-OLD_KEY = "lm_head.decoder.weight"
-NEW_KEY = "lm_head.weight"
-
-
-def convert_dialogpt_checkpoint(checkpoint_path: str, pytorch_dump_folder_path: str):
- d = torch.load(checkpoint_path)
- d[NEW_KEY] = d.pop(OLD_KEY)
- os.makedirs(pytorch_dump_folder_path, exist_ok=True)
- torch.save(d, os.path.join(pytorch_dump_folder_path, WEIGHTS_NAME))
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument("--dialogpt_path", default=".", type=str)
- args = parser.parse_args()
- for MODEL in DIALOGPT_MODELS:
- checkpoint_path = os.path.join(args.dialogpt_path, f"{MODEL}_ft.pkl")
- pytorch_dump_folder_path = f"./DialoGPT-{MODEL}"
- convert_dialogpt_checkpoint(
- checkpoint_path,
- pytorch_dump_folder_path,
- )
diff --git a/transformers/models/dinat/__init__.py b/transformers/models/dinat/__init__.py
deleted file mode 100644
index 88470f1ca9f9bd68a2f89691cfe5b9031e3cae66..0000000000000000000000000000000000000000
--- a/transformers/models/dinat/__init__.py
+++ /dev/null
@@ -1,56 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {"configuration_dinat": ["DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP", "DinatConfig"]}
-
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_dinat"] = [
- "DINAT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DinatForImageClassification",
- "DinatModel",
- "DinatPreTrainedModel",
- "DinatBackbone",
- ]
-
-if TYPE_CHECKING:
- from .configuration_dinat import DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP, DinatConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_dinat import (
- DINAT_PRETRAINED_MODEL_ARCHIVE_LIST,
- DinatBackbone,
- DinatForImageClassification,
- DinatModel,
- DinatPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/dinat/__pycache__/__init__.cpython-310.pyc b/transformers/models/dinat/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 7b87d3c8cc51ded080662015157a5b1761c6ea09..0000000000000000000000000000000000000000
Binary files a/transformers/models/dinat/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dinat/__pycache__/configuration_dinat.cpython-310.pyc b/transformers/models/dinat/__pycache__/configuration_dinat.cpython-310.pyc
deleted file mode 100644
index de188ac0a4259c955581f8c4b15263af85335dd1..0000000000000000000000000000000000000000
Binary files a/transformers/models/dinat/__pycache__/configuration_dinat.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dinat/__pycache__/modeling_dinat.cpython-310.pyc b/transformers/models/dinat/__pycache__/modeling_dinat.cpython-310.pyc
deleted file mode 100644
index d838c5af4b3def3d106b9cbf7eb89b501cc5b648..0000000000000000000000000000000000000000
Binary files a/transformers/models/dinat/__pycache__/modeling_dinat.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dinat/configuration_dinat.py b/transformers/models/dinat/configuration_dinat.py
deleted file mode 100644
index 4bd38c73857a972db6d66a05b48ca9049bbe1a12..0000000000000000000000000000000000000000
--- a/transformers/models/dinat/configuration_dinat.py
+++ /dev/null
@@ -1,152 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Dilated Neighborhood Attention Transformer model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DinatConfig(BackboneConfigMixin, PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DinatModel`]. It is used to instantiate a Dinat
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the Dinat
- [shi-labs/dinat-mini-in1k-224](https://huggingface.co/shi-labs/dinat-mini-in1k-224) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- patch_size (`int`, *optional*, defaults to 4):
- The size (resolution) of each patch. NOTE: Only patch size of 4 is supported at the moment.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- embed_dim (`int`, *optional*, defaults to 64):
- Dimensionality of patch embedding.
- depths (`List[int]`, *optional*, defaults to `[3, 4, 6, 5]`):
- Number of layers in each level of the encoder.
- num_heads (`List[int]`, *optional*, defaults to `[2, 4, 8, 16]`):
- Number of attention heads in each layer of the Transformer encoder.
- kernel_size (`int`, *optional*, defaults to 7):
- Neighborhood Attention kernel size.
- dilations (`List[List[int]]`, *optional*, defaults to `[[1, 8, 1], [1, 4, 1, 4], [1, 2, 1, 2, 1, 2], [1, 1, 1, 1, 1]]`):
- Dilation value of each NA layer in the Transformer encoder.
- mlp_ratio (`float`, *optional*, defaults to 3.0):
- Ratio of MLP hidden dimensionality to embedding dimensionality.
- qkv_bias (`bool`, *optional*, defaults to `True`):
- Whether or not a learnable bias should be added to the queries, keys and values.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings and encoder.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- drop_path_rate (`float`, *optional*, defaults to 0.1):
- Stochastic depth rate.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder. If string, `"gelu"`, `"relu"`,
- `"selu"` and `"gelu_new"` are supported.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the layer normalization layers.
- layer_scale_init_value (`float`, *optional*, defaults to 0.0):
- The initial value for the layer scale. Disabled if <=0.
- out_features (`List[str]`, *optional*):
- If used as backbone, list of features to output. Can be any of `"stem"`, `"stage1"`, `"stage2"`, etc.
- (depending on how many stages the model has). If unset and `out_indices` is set, will default to the
- corresponding stages. If unset and `out_indices` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
- out_indices (`List[int]`, *optional*):
- If used as backbone, list of indices of features to output. Can be any of 0, 1, 2, etc. (depending on how
- many stages the model has). If unset and `out_features` is set, will default to the corresponding stages.
- If unset and `out_features` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
-
- Example:
-
- ```python
- >>> from transformers import DinatConfig, DinatModel
-
- >>> # Initializing a Dinat shi-labs/dinat-mini-in1k-224 style configuration
- >>> configuration = DinatConfig()
-
- >>> # Initializing a model (with random weights) from the shi-labs/dinat-mini-in1k-224 style configuration
- >>> model = DinatModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "dinat"
-
- attribute_map = {
- "num_attention_heads": "num_heads",
- "num_hidden_layers": "num_layers",
- }
-
- def __init__(
- self,
- patch_size=4,
- num_channels=3,
- embed_dim=64,
- depths=[3, 4, 6, 5],
- num_heads=[2, 4, 8, 16],
- kernel_size=7,
- dilations=[[1, 8, 1], [1, 4, 1, 4], [1, 2, 1, 2, 1, 2], [1, 1, 1, 1, 1]],
- mlp_ratio=3.0,
- qkv_bias=True,
- hidden_dropout_prob=0.0,
- attention_probs_dropout_prob=0.0,
- drop_path_rate=0.1,
- hidden_act="gelu",
- initializer_range=0.02,
- layer_norm_eps=1e-5,
- layer_scale_init_value=0.0,
- out_features=None,
- out_indices=None,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.embed_dim = embed_dim
- self.depths = depths
- self.num_layers = len(depths)
- self.num_heads = num_heads
- self.kernel_size = kernel_size
- self.dilations = dilations
- self.mlp_ratio = mlp_ratio
- self.qkv_bias = qkv_bias
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.drop_path_rate = drop_path_rate
- self.hidden_act = hidden_act
- self.layer_norm_eps = layer_norm_eps
- self.initializer_range = initializer_range
- # we set the hidden_size attribute in order to make Dinat work with VisionEncoderDecoderModel
- # this indicates the channel dimension after the last stage of the model
- self.hidden_size = int(embed_dim * 2 ** (len(depths) - 1))
- self.layer_scale_init_value = layer_scale_init_value
- self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(depths) + 1)]
- self._out_features, self._out_indices = get_aligned_output_features_output_indices(
- out_features=out_features, out_indices=out_indices, stage_names=self.stage_names
- )
diff --git a/transformers/models/dinat/modeling_dinat.py b/transformers/models/dinat/modeling_dinat.py
deleted file mode 100644
index 72bf6d1170094c8bfbc9b92ca3a102fbdec331ff..0000000000000000000000000000000000000000
--- a/transformers/models/dinat/modeling_dinat.py
+++ /dev/null
@@ -1,976 +0,0 @@
-# coding=utf-8
-# Copyright 2022 SHI Labs and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Dilated Neighborhood Attention Transformer model."""
-
-
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import BackboneOutput
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- ModelOutput,
- OptionalDependencyNotAvailable,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_natten_available,
- logging,
- replace_return_docstrings,
- requires_backends,
-)
-from ...utils.backbone_utils import BackboneMixin
-from .configuration_dinat import DinatConfig
-
-
-if is_natten_available():
- from natten.functional import natten2dav, natten2dqkrpb
-else:
-
- def natten2dqkrpb(*args, **kwargs):
- raise OptionalDependencyNotAvailable()
-
- def natten2dav(*args, **kwargs):
- raise OptionalDependencyNotAvailable()
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "DinatConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "shi-labs/dinat-mini-in1k-224"
-_EXPECTED_OUTPUT_SHAPE = [1, 7, 7, 512]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "shi-labs/dinat-mini-in1k-224"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# drop_path and DinatDropPath are from the timm library.
-
-
-@dataclass
-# Copied from transformers.models.nat.modeling_nat.NatEncoderOutput with Nat->Dinat
-class DinatEncoderOutput(ModelOutput):
- """
- Dinat encoder's outputs, with potential hidden states and attentions.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each stage) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- last_hidden_state: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-@dataclass
-# Copied from transformers.models.nat.modeling_nat.NatModelOutput with Nat->Dinat
-class DinatModelOutput(ModelOutput):
- """
- Dinat model's outputs that also contains a pooling of the last hidden states.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*, returned when `add_pooling_layer=True` is passed):
- Average pooling of the last layer hidden-state.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each stage) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- last_hidden_state: torch.FloatTensor = None
- pooler_output: Optional[torch.FloatTensor] = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-@dataclass
-# Copied from transformers.models.nat.modeling_nat.NatImageClassifierOutput with Nat->Dinat
-class DinatImageClassifierOutput(ModelOutput):
- """
- Dinat outputs for image classification.
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
- Classification (or regression if config.num_labels==1) loss.
- logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Classification (or regression if config.num_labels==1) scores (before SoftMax).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each stage) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- loss: Optional[torch.FloatTensor] = None
- logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-# Copied from transformers.models.nat.modeling_nat.NatEmbeddings with Nat->Dinat
-class DinatEmbeddings(nn.Module):
- """
- Construct the patch and position embeddings.
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.patch_embeddings = DinatPatchEmbeddings(config)
-
- self.norm = nn.LayerNorm(config.embed_dim)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, pixel_values: Optional[torch.FloatTensor]) -> Tuple[torch.Tensor]:
- embeddings = self.patch_embeddings(pixel_values)
- embeddings = self.norm(embeddings)
-
- embeddings = self.dropout(embeddings)
-
- return embeddings
-
-
-# Copied from transformers.models.nat.modeling_nat.NatPatchEmbeddings with Nat->Dinat
-class DinatPatchEmbeddings(nn.Module):
- """
- This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
- `hidden_states` (patch embeddings) of shape `(batch_size, height, width, hidden_size)` to be consumed by a
- Transformer.
- """
-
- def __init__(self, config):
- super().__init__()
- patch_size = config.patch_size
- num_channels, hidden_size = config.num_channels, config.embed_dim
- self.num_channels = num_channels
-
- if patch_size == 4:
- pass
- else:
- # TODO: Support arbitrary patch sizes.
- raise ValueError("Dinat only supports patch size of 4 at the moment.")
-
- self.projection = nn.Sequential(
- nn.Conv2d(self.num_channels, hidden_size // 2, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)),
- nn.Conv2d(hidden_size // 2, hidden_size, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)),
- )
-
- def forward(self, pixel_values: Optional[torch.FloatTensor]) -> torch.Tensor:
- _, num_channels, height, width = pixel_values.shape
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- embeddings = self.projection(pixel_values)
- embeddings = embeddings.permute(0, 2, 3, 1)
-
- return embeddings
-
-
-# Copied from transformers.models.nat.modeling_nat.NatDownsampler with Nat->Dinat
-class DinatDownsampler(nn.Module):
- """
- Convolutional Downsampling Layer.
-
- Args:
- dim (`int`):
- Number of input channels.
- norm_layer (`nn.Module`, *optional*, defaults to `nn.LayerNorm`):
- Normalization layer class.
- """
-
- def __init__(self, dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None:
- super().__init__()
- self.dim = dim
- self.reduction = nn.Conv2d(dim, 2 * dim, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
- self.norm = norm_layer(2 * dim)
-
- def forward(self, input_feature: torch.Tensor) -> torch.Tensor:
- input_feature = self.reduction(input_feature.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
- input_feature = self.norm(input_feature)
- return input_feature
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->Dinat
-class DinatDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-class NeighborhoodAttention(nn.Module):
- def __init__(self, config, dim, num_heads, kernel_size, dilation):
- super().__init__()
- if dim % num_heads != 0:
- raise ValueError(
- f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})"
- )
-
- self.num_attention_heads = num_heads
- self.attention_head_size = int(dim / num_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.kernel_size = kernel_size
- self.dilation = dilation
-
- # rpb is learnable relative positional biases; same concept is used Swin.
- self.rpb = nn.Parameter(torch.zeros(num_heads, (2 * self.kernel_size - 1), (2 * self.kernel_size - 1)))
-
- self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
- self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
- self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- # Copied from transformers.models.nat.modeling_nat.NeighborhoodAttention.transpose_for_scores with Nat->Dinat
- def transpose_for_scores(self, x):
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 3, 1, 2, 4)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- query_layer = self.transpose_for_scores(self.query(hidden_states))
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- # Apply the scale factor before computing attention weights. It's usually more efficient because
- # attention weights are typically a bigger tensor compared to query.
- # It gives identical results because scalars are commutable in matrix multiplication.
- query_layer = query_layer / math.sqrt(self.attention_head_size)
-
- # Compute NA between "query" and "key" to get the raw attention scores, and add relative positional biases.
- attention_scores = natten2dqkrpb(query_layer, key_layer, self.rpb, self.kernel_size, self.dilation)
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- context_layer = natten2dav(attention_probs, value_layer, self.kernel_size, self.dilation)
- context_layer = context_layer.permute(0, 2, 3, 1, 4).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-# Copied from transformers.models.nat.modeling_nat.NeighborhoodAttentionOutput
-class NeighborhoodAttentionOutput(nn.Module):
- def __init__(self, config, dim):
- super().__init__()
- self.dense = nn.Linear(dim, dim)
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-class NeighborhoodAttentionModule(nn.Module):
- def __init__(self, config, dim, num_heads, kernel_size, dilation):
- super().__init__()
- self.self = NeighborhoodAttention(config, dim, num_heads, kernel_size, dilation)
- self.output = NeighborhoodAttentionOutput(config, dim)
- self.pruned_heads = set()
-
- # Copied from transformers.models.nat.modeling_nat.NeighborhoodAttentionModule.prune_heads
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- # Copied from transformers.models.nat.modeling_nat.NeighborhoodAttentionModule.forward
- def forward(
- self,
- hidden_states: torch.Tensor,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- self_outputs = self.self(hidden_states, output_attentions)
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.nat.modeling_nat.NatIntermediate with Nat->Dinat
-class DinatIntermediate(nn.Module):
- def __init__(self, config, dim):
- super().__init__()
- self.dense = nn.Linear(dim, int(config.mlp_ratio * dim))
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.nat.modeling_nat.NatOutput with Nat->Dinat
-class DinatOutput(nn.Module):
- def __init__(self, config, dim):
- super().__init__()
- self.dense = nn.Linear(int(config.mlp_ratio * dim), dim)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-class DinatLayer(nn.Module):
- def __init__(self, config, dim, num_heads, dilation, drop_path_rate=0.0):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.kernel_size = config.kernel_size
- self.dilation = dilation
- self.window_size = self.kernel_size * self.dilation
- self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
- self.attention = NeighborhoodAttentionModule(
- config, dim, num_heads, kernel_size=self.kernel_size, dilation=self.dilation
- )
- self.drop_path = DinatDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
- self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
- self.intermediate = DinatIntermediate(config, dim)
- self.output = DinatOutput(config, dim)
- self.layer_scale_parameters = (
- nn.Parameter(config.layer_scale_init_value * torch.ones((2, dim)), requires_grad=True)
- if config.layer_scale_init_value > 0
- else None
- )
-
- def maybe_pad(self, hidden_states, height, width):
- window_size = self.window_size
- pad_values = (0, 0, 0, 0, 0, 0)
- if height < window_size or width < window_size:
- pad_l = pad_t = 0
- pad_r = max(0, window_size - width)
- pad_b = max(0, window_size - height)
- pad_values = (0, 0, pad_l, pad_r, pad_t, pad_b)
- hidden_states = nn.functional.pad(hidden_states, pad_values)
- return hidden_states, pad_values
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- batch_size, height, width, channels = hidden_states.size()
- shortcut = hidden_states
-
- hidden_states = self.layernorm_before(hidden_states)
- # pad hidden_states if they are smaller than kernel size x dilation
- hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)
-
- _, height_pad, width_pad, _ = hidden_states.shape
-
- attention_outputs = self.attention(hidden_states, output_attentions=output_attentions)
-
- attention_output = attention_outputs[0]
-
- was_padded = pad_values[3] > 0 or pad_values[5] > 0
- if was_padded:
- attention_output = attention_output[:, :height, :width, :].contiguous()
-
- if self.layer_scale_parameters is not None:
- attention_output = self.layer_scale_parameters[0] * attention_output
-
- hidden_states = shortcut + self.drop_path(attention_output)
-
- layer_output = self.layernorm_after(hidden_states)
- layer_output = self.output(self.intermediate(layer_output))
-
- if self.layer_scale_parameters is not None:
- layer_output = self.layer_scale_parameters[1] * layer_output
-
- layer_output = hidden_states + self.drop_path(layer_output)
-
- layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)
- return layer_outputs
-
-
-class DinatStage(nn.Module):
- def __init__(self, config, dim, depth, num_heads, dilations, drop_path_rate, downsample):
- super().__init__()
- self.config = config
- self.dim = dim
- self.layers = nn.ModuleList(
- [
- DinatLayer(
- config=config,
- dim=dim,
- num_heads=num_heads,
- dilation=dilations[i],
- drop_path_rate=drop_path_rate[i],
- )
- for i in range(depth)
- ]
- )
-
- # patch merging layer
- if downsample is not None:
- self.downsample = downsample(dim=dim, norm_layer=nn.LayerNorm)
- else:
- self.downsample = None
-
- self.pointing = False
-
- # Copied from transformers.models.nat.modeling_nat.NatStage.forward
- def forward(
- self,
- hidden_states: torch.Tensor,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- _, height, width, _ = hidden_states.size()
- for i, layer_module in enumerate(self.layers):
- layer_outputs = layer_module(hidden_states, output_attentions)
- hidden_states = layer_outputs[0]
-
- hidden_states_before_downsampling = hidden_states
- if self.downsample is not None:
- hidden_states = self.downsample(hidden_states_before_downsampling)
-
- stage_outputs = (hidden_states, hidden_states_before_downsampling)
-
- if output_attentions:
- stage_outputs += layer_outputs[1:]
- return stage_outputs
-
-
-class DinatEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.num_levels = len(config.depths)
- self.config = config
- dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
- self.levels = nn.ModuleList(
- [
- DinatStage(
- config=config,
- dim=int(config.embed_dim * 2**i_layer),
- depth=config.depths[i_layer],
- num_heads=config.num_heads[i_layer],
- dilations=config.dilations[i_layer],
- drop_path_rate=dpr[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])],
- downsample=DinatDownsampler if (i_layer < self.num_levels - 1) else None,
- )
- for i_layer in range(self.num_levels)
- ]
- )
-
- # Copied from transformers.models.nat.modeling_nat.NatEncoder.forward with Nat->Dinat
- def forward(
- self,
- hidden_states: torch.Tensor,
- output_attentions: Optional[bool] = False,
- output_hidden_states: Optional[bool] = False,
- output_hidden_states_before_downsampling: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple, DinatEncoderOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_reshaped_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- if output_hidden_states:
- # rearrange b h w c -> b c h w
- reshaped_hidden_state = hidden_states.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
-
- for i, layer_module in enumerate(self.levels):
- layer_outputs = layer_module(hidden_states, output_attentions)
-
- hidden_states = layer_outputs[0]
- hidden_states_before_downsampling = layer_outputs[1]
-
- if output_hidden_states and output_hidden_states_before_downsampling:
- # rearrange b h w c -> b c h w
- reshaped_hidden_state = hidden_states_before_downsampling.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states_before_downsampling,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
- elif output_hidden_states and not output_hidden_states_before_downsampling:
- # rearrange b h w c -> b c h w
- reshaped_hidden_state = hidden_states.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
-
- if output_attentions:
- all_self_attentions += layer_outputs[2:]
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
-
- return DinatEncoderOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- reshaped_hidden_states=all_reshaped_hidden_states,
- )
-
-
-class DinatPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DinatConfig
- base_model_prefix = "dinat"
- main_input_name = "pixel_values"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-DINAT_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`DinatConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DINAT_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`ViTImageProcessor.__call__`]
- for details.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Dinat Model transformer outputting raw hidden-states without any specific head on top.",
- DINAT_START_DOCSTRING,
-)
-# Copied from transformers.models.nat.modeling_nat.NatModel with Nat->Dinat, NAT->DINAT
-class DinatModel(DinatPreTrainedModel):
- def __init__(self, config, add_pooling_layer=True):
- super().__init__(config)
-
- requires_backends(self, ["natten"])
-
- self.config = config
- self.num_levels = len(config.depths)
- self.num_features = int(config.embed_dim * 2 ** (self.num_levels - 1))
-
- self.embeddings = DinatEmbeddings(config)
- self.encoder = DinatEncoder(config)
-
- self.layernorm = nn.LayerNorm(self.num_features, eps=config.layer_norm_eps)
- self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.patch_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(DINAT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=DinatModelOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, DinatModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- embedding_output = self.embeddings(pixel_values)
-
- encoder_outputs = self.encoder(
- embedding_output,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
-
- pooled_output = None
- if self.pooler is not None:
- pooled_output = self.pooler(sequence_output.flatten(1, 2).transpose(1, 2))
- pooled_output = torch.flatten(pooled_output, 1)
-
- if not return_dict:
- output = (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return output
-
- return DinatModelOutput(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- reshaped_hidden_states=encoder_outputs.reshaped_hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- Dinat Model transformer with an image classification head on top (a linear layer on top of the final hidden state
- of the [CLS] token) e.g. for ImageNet.
- """,
- DINAT_START_DOCSTRING,
-)
-class DinatForImageClassification(DinatPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- requires_backends(self, ["natten"])
-
- self.num_labels = config.num_labels
- self.dinat = DinatModel(config)
-
- # Classifier head
- self.classifier = (
- nn.Linear(self.dinat.num_features, config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DINAT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=DinatImageClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, DinatImageClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.dinat(
- pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return DinatImageClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- reshaped_hidden_states=outputs.reshaped_hidden_states,
- )
-
-
-@add_start_docstrings(
- "NAT backbone, to be used with frameworks like DETR and MaskFormer.",
- DINAT_START_DOCSTRING,
-)
-class DinatBackbone(DinatPreTrainedModel, BackboneMixin):
- def __init__(self, config):
- super().__init__(config)
- super()._init_backbone(config)
-
- requires_backends(self, ["natten"])
-
- self.embeddings = DinatEmbeddings(config)
- self.encoder = DinatEncoder(config)
- self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))]
-
- # Add layer norms to hidden states of out_features
- hidden_states_norms = {}
- for stage, num_channels in zip(self._out_features, self.channels):
- hidden_states_norms[stage] = nn.LayerNorm(num_channels)
- self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.patch_embeddings
-
- @add_start_docstrings_to_model_forward(DINAT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.Tensor,
- output_hidden_states: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> BackboneOutput:
- """
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, AutoBackbone
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> processor = AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224")
- >>> model = AutoBackbone.from_pretrained(
- ... "shi-labs/nat-mini-in1k-224", out_features=["stage1", "stage2", "stage3", "stage4"]
- ... )
-
- >>> inputs = processor(image, return_tensors="pt")
-
- >>> outputs = model(**inputs)
-
- >>> feature_maps = outputs.feature_maps
- >>> list(feature_maps[-1].shape)
- [1, 512, 7, 7]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-
- embedding_output = self.embeddings(pixel_values)
-
- outputs = self.encoder(
- embedding_output,
- output_attentions=output_attentions,
- output_hidden_states=True,
- output_hidden_states_before_downsampling=True,
- return_dict=True,
- )
-
- hidden_states = outputs.reshaped_hidden_states
-
- feature_maps = ()
- for stage, hidden_state in zip(self.stage_names, hidden_states):
- if stage in self.out_features:
- batch_size, num_channels, height, width = hidden_state.shape
- hidden_state = hidden_state.permute(0, 2, 3, 1).contiguous()
- hidden_state = hidden_state.view(batch_size, height * width, num_channels)
- hidden_state = self.hidden_states_norms[stage](hidden_state)
- hidden_state = hidden_state.view(batch_size, height, width, num_channels)
- hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
- feature_maps += (hidden_state,)
-
- if not return_dict:
- output = (feature_maps,)
- if output_hidden_states:
- output += (outputs.hidden_states,)
- return output
-
- return BackboneOutput(
- feature_maps=feature_maps,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/dinov2/__init__.py b/transformers/models/dinov2/__init__.py
deleted file mode 100644
index 01d02a9e65fda02e543b116dc4bf7ccba6097c6e..0000000000000000000000000000000000000000
--- a/transformers/models/dinov2/__init__.py
+++ /dev/null
@@ -1,61 +0,0 @@
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_dinov2": ["DINOV2_PRETRAINED_CONFIG_ARCHIVE_MAP", "Dinov2Config", "Dinov2OnnxConfig"]
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_dinov2"] = [
- "DINOV2_PRETRAINED_MODEL_ARCHIVE_LIST",
- "Dinov2ForImageClassification",
- "Dinov2Model",
- "Dinov2PreTrainedModel",
- "Dinov2Backbone",
- ]
-
-if TYPE_CHECKING:
- from .configuration_dinov2 import DINOV2_PRETRAINED_CONFIG_ARCHIVE_MAP, Dinov2Config, Dinov2OnnxConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_dinov2 import (
- DINOV2_PRETRAINED_MODEL_ARCHIVE_LIST,
- Dinov2Backbone,
- Dinov2ForImageClassification,
- Dinov2Model,
- Dinov2PreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/dinov2/__pycache__/__init__.cpython-310.pyc b/transformers/models/dinov2/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index aeb1c4f1c31428a9201d3006d96d1d0194634355..0000000000000000000000000000000000000000
Binary files a/transformers/models/dinov2/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dinov2/__pycache__/configuration_dinov2.cpython-310.pyc b/transformers/models/dinov2/__pycache__/configuration_dinov2.cpython-310.pyc
deleted file mode 100644
index 564a7c6a88ae03d67c90f1bb1ef9a0807be64721..0000000000000000000000000000000000000000
Binary files a/transformers/models/dinov2/__pycache__/configuration_dinov2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dinov2/__pycache__/convert_dinov2_to_hf.cpython-310.pyc b/transformers/models/dinov2/__pycache__/convert_dinov2_to_hf.cpython-310.pyc
deleted file mode 100644
index 59eeb9ed1289031c3e7c7ef1c3ed63446a9ccf80..0000000000000000000000000000000000000000
Binary files a/transformers/models/dinov2/__pycache__/convert_dinov2_to_hf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dinov2/__pycache__/modeling_dinov2.cpython-310.pyc b/transformers/models/dinov2/__pycache__/modeling_dinov2.cpython-310.pyc
deleted file mode 100644
index dd859bd70df2cb3fe257dd13df294720bb373fda..0000000000000000000000000000000000000000
Binary files a/transformers/models/dinov2/__pycache__/modeling_dinov2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dinov2/configuration_dinov2.py b/transformers/models/dinov2/configuration_dinov2.py
deleted file mode 100644
index b5fe872a706fc71808562e8152db4eee4ca7218f..0000000000000000000000000000000000000000
--- a/transformers/models/dinov2/configuration_dinov2.py
+++ /dev/null
@@ -1,175 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DINOv2 model configuration"""
-
-from collections import OrderedDict
-from typing import Mapping
-
-from packaging import version
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DINOV2_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class Dinov2Config(BackboneConfigMixin, PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`Dinov2Model`]. It is used to instantiate an
- Dinov2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the Dinov2
- [google/dinov2-base-patch16-224](https://huggingface.co/google/dinov2-base-patch16-224) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- mlp_ratio (`int`, *optional*, defaults to 4):
- Ratio of the hidden size of the MLPs relative to the `hidden_size`.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-06):
- The epsilon used by the layer normalization layers.
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 16):
- The size (resolution) of each patch.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- qkv_bias (`bool`, *optional*, defaults to `True`):
- Whether to add a bias to the queries, keys and values.
- layerscale_value (`float`, *optional*, defaults to 1.0):
- Initial value to use for layer scale.
- drop_path_rate (`float`, *optional*, defaults to 0.0):
- Stochastic depth rate per sample (when applied in the main path of residual layers).
- use_swiglu_ffn (`bool`, *optional*, defaults to `False`):
- Whether to use the SwiGLU feedforward neural network.
- out_features (`List[str]`, *optional*):
- If used as backbone, list of features to output. Can be any of `"stem"`, `"stage1"`, `"stage2"`, etc.
- (depending on how many stages the model has). If unset and `out_indices` is set, will default to the
- corresponding stages. If unset and `out_indices` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
- out_indices (`List[int]`, *optional*):
- If used as backbone, list of indices of features to output. Can be any of 0, 1, 2, etc. (depending on how
- many stages the model has). If unset and `out_features` is set, will default to the corresponding stages.
- If unset and `out_features` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
- apply_layernorm (`bool`, *optional*, defaults to `True`):
- Whether to apply layer normalization to the feature maps in case the model is used as backbone.
- reshape_hidden_states (`bool`, *optional*, defaults to `True`):
- Whether to reshape the feature maps to 4D tensors of shape `(batch_size, hidden_size, height, width)` in
- case the model is used as backbone. If `False`, the feature maps will be 3D tensors of shape `(batch_size,
- seq_len, hidden_size)`.
-
- Example:
-
- ```python
- >>> from transformers import Dinov2Config, Dinov2Model
-
- >>> # Initializing a Dinov2 dinov2-base-patch16-224 style configuration
- >>> configuration = Dinov2Config()
-
- >>> # Initializing a model (with random weights) from the dinov2-base-patch16-224 style configuration
- >>> model = Dinov2Model(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "dinov2"
-
- def __init__(
- self,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- mlp_ratio=4,
- hidden_act="gelu",
- hidden_dropout_prob=0.0,
- attention_probs_dropout_prob=0.0,
- initializer_range=0.02,
- layer_norm_eps=1e-6,
- image_size=224,
- patch_size=16,
- num_channels=3,
- qkv_bias=True,
- layerscale_value=1.0,
- drop_path_rate=0.0,
- use_swiglu_ffn=False,
- out_features=None,
- out_indices=None,
- apply_layernorm=True,
- reshape_hidden_states=True,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.mlp_ratio = mlp_ratio
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.qkv_bias = qkv_bias
- self.layerscale_value = layerscale_value
- self.drop_path_rate = drop_path_rate
- self.use_swiglu_ffn = use_swiglu_ffn
- self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, num_hidden_layers + 1)]
- self._out_features, self._out_indices = get_aligned_output_features_output_indices(
- out_features=out_features, out_indices=out_indices, stage_names=self.stage_names
- )
- self.apply_layernorm = apply_layernorm
- self.reshape_hidden_states = reshape_hidden_states
-
-
-class Dinov2OnnxConfig(OnnxConfig):
- torch_onnx_minimum_version = version.parse("1.11")
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- return OrderedDict(
- [
- ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
- ]
- )
-
- @property
- def atol_for_validation(self) -> float:
- return 1e-4
diff --git a/transformers/models/dinov2/convert_dinov2_to_hf.py b/transformers/models/dinov2/convert_dinov2_to_hf.py
deleted file mode 100644
index dd5871e6c440661a7050bab7696db39e865b714d..0000000000000000000000000000000000000000
--- a/transformers/models/dinov2/convert_dinov2_to_hf.py
+++ /dev/null
@@ -1,287 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DINOv2 checkpoints from the original repository.
-
-URL: https://github.com/facebookresearch/dinov2/tree/main
-"""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-import torch.nn as nn
-from huggingface_hub import hf_hub_download
-from PIL import Image
-from torchvision import transforms
-
-from transformers import BitImageProcessor, Dinov2Config, Dinov2ForImageClassification, Dinov2Model
-from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_dinov2_config(model_name, image_classifier=False):
- config = Dinov2Config(image_size=518, patch_size=14)
-
- # size of the architecture
- if "vits" in model_name:
- config.hidden_size = 384
- config.num_attention_heads = 6
- elif "vitb" in model_name:
- pass
- elif "vitl" in model_name:
- config.hidden_size = 1024
- config.num_hidden_layers = 24
- config.num_attention_heads = 16
- elif "vitg" in model_name:
- config.use_swiglu_ffn = True
- config.hidden_size = 1536
- config.num_hidden_layers = 40
- config.num_attention_heads = 24
- else:
- raise ValueError("Model not supported")
-
- if image_classifier:
- repo_id = "huggingface/label-files"
- filename = "imagenet-1k-id2label.json"
- config.num_labels = 1000
- config.id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- config.id2label = {int(k): v for k, v in config.id2label.items()}
-
- return config
-
-
-def create_rename_keys(config):
- rename_keys = []
- # fmt: off
-
- # patch embedding layer
- rename_keys.append(("cls_token", "embeddings.cls_token"))
- rename_keys.append(("mask_token", "embeddings.mask_token"))
- rename_keys.append(("pos_embed", "embeddings.position_embeddings"))
- rename_keys.append(("patch_embed.proj.weight", "embeddings.patch_embeddings.projection.weight"))
- rename_keys.append(("patch_embed.proj.bias", "embeddings.patch_embeddings.projection.bias"))
-
- for i in range(config.num_hidden_layers):
- # layernorms
- rename_keys.append((f"blocks.{i}.norm1.weight", f"encoder.layer.{i}.norm1.weight"))
- rename_keys.append((f"blocks.{i}.norm1.bias", f"encoder.layer.{i}.norm1.bias"))
- rename_keys.append((f"blocks.{i}.norm2.weight", f"encoder.layer.{i}.norm2.weight"))
- rename_keys.append((f"blocks.{i}.norm2.bias", f"encoder.layer.{i}.norm2.bias"))
- # MLP
- if config.use_swiglu_ffn:
- rename_keys.append((f"blocks.{i}.mlp.w12.weight", f"encoder.layer.{i}.mlp.w12.weight"))
- rename_keys.append((f"blocks.{i}.mlp.w12.bias", f"encoder.layer.{i}.mlp.w12.bias"))
- rename_keys.append((f"blocks.{i}.mlp.w3.weight", f"encoder.layer.{i}.mlp.w3.weight"))
- rename_keys.append((f"blocks.{i}.mlp.w3.bias", f"encoder.layer.{i}.mlp.w3.bias"))
- else:
- rename_keys.append((f"blocks.{i}.mlp.fc1.weight", f"encoder.layer.{i}.mlp.fc1.weight"))
- rename_keys.append((f"blocks.{i}.mlp.fc1.bias", f"encoder.layer.{i}.mlp.fc1.bias"))
- rename_keys.append((f"blocks.{i}.mlp.fc2.weight", f"encoder.layer.{i}.mlp.fc2.weight"))
- rename_keys.append((f"blocks.{i}.mlp.fc2.bias", f"encoder.layer.{i}.mlp.fc2.bias"))
- # layerscale
- rename_keys.append((f"blocks.{i}.ls1.gamma", f"encoder.layer.{i}.layer_scale1.lambda1"))
- rename_keys.append((f"blocks.{i}.ls2.gamma", f"encoder.layer.{i}.layer_scale2.lambda1"))
- # attention projection layer
- rename_keys.append((f"blocks.{i}.attn.proj.weight", f"encoder.layer.{i}.attention.output.dense.weight"))
- rename_keys.append((f"blocks.{i}.attn.proj.bias", f"encoder.layer.{i}.attention.output.dense.bias"))
-
- # final layernorm
- rename_keys.append(("norm.weight", "layernorm.weight"))
- rename_keys.append(("norm.bias", "layernorm.bias"))
-
- # fmt: on
- return rename_keys
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config):
- for i in range(config.num_hidden_layers):
- # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"blocks.{i}.attn.qkv.weight")
- in_proj_bias = state_dict.pop(f"blocks.{i}.attn.qkv.bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[: config.hidden_size, :]
- state_dict[f"encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size]
- state_dict[f"encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- config.hidden_size : config.hidden_size * 2, :
- ]
- state_dict[f"encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
- config.hidden_size : config.hidden_size * 2
- ]
- state_dict[f"encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-config.hidden_size :, :]
- state_dict[f"encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- image = Image.open(requests.get(url, stream=True).raw)
- return image
-
-
-@torch.no_grad()
-def convert_dinov2_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub=False):
- """
- Copy/paste/tweak model's weights to our DINOv2 structure.
- """
-
- # define default Dinov2 configuration
- image_classifier = "1layer" in model_name
- config = get_dinov2_config(model_name, image_classifier=image_classifier)
-
- # load original model from torch hub
- original_model = torch.hub.load("facebookresearch/dinov2", model_name.replace("_1layer", ""))
- original_model.eval()
-
- # load state_dict of original model, remove and rename some keys
- state_dict = original_model.state_dict()
- rename_keys = create_rename_keys(config)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- read_in_q_k_v(state_dict, config)
-
- for key, val in state_dict.copy().items():
- val = state_dict.pop(key)
- if "w12" in key:
- key = key.replace("w12", "weights_in")
- if "w3" in key:
- key = key.replace("w3", "weights_out")
- state_dict[key] = val
-
- # load HuggingFace model
- if image_classifier:
- model = Dinov2ForImageClassification(config).eval()
- model.dinov2.load_state_dict(state_dict)
- model_name_to_classifier_dict_url = {
- "dinov2_vits14_1layer": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_linear_head.pth",
- "dinov2_vitb14_1layer": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_linear_head.pth",
- "dinov2_vitl14_1layer": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_linear_head.pth",
- "dinov2_vitg14_1layer": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_linear_head.pth",
- }
- url = model_name_to_classifier_dict_url[model_name]
- classifier_state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
- model.classifier.weight = nn.Parameter(classifier_state_dict["weight"])
- model.classifier.bias = nn.Parameter(classifier_state_dict["bias"])
- else:
- model = Dinov2Model(config).eval()
- model.load_state_dict(state_dict)
-
- # load image
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
-
- # preprocess image
- transformations = transforms.Compose(
- [
- transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
- transforms.CenterCrop(224),
- transforms.ToTensor(),
- transforms.Normalize(
- mean=IMAGENET_DEFAULT_MEAN, # these are RGB mean+std values
- std=IMAGENET_DEFAULT_STD, # across a large photo dataset.
- ),
- ]
- )
-
- original_pixel_values = transformations(image).unsqueeze(0) # insert batch dimension
-
- processor = BitImageProcessor(
- size={"shortest_edge": 256},
- resample=PILImageResampling.BICUBIC,
- image_mean=IMAGENET_DEFAULT_MEAN,
- image_std=IMAGENET_DEFAULT_STD,
- )
- pixel_values = processor(image, return_tensors="pt").pixel_values
-
- assert torch.allclose(original_pixel_values, pixel_values)
-
- with torch.no_grad():
- outputs = model(pixel_values, output_hidden_states=True)
- original_outputs = original_model(pixel_values)
-
- # assert values
- if image_classifier:
- print("Predicted class:")
- class_idx = outputs.logits.argmax(-1).item()
- print(model.config.id2label[class_idx])
- else:
- assert outputs.last_hidden_state[:, 0].shape == original_outputs.shape
- assert torch.allclose(outputs.last_hidden_state[:, 0], original_outputs, atol=1e-3)
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- print(f"Saving image processor to {pytorch_dump_folder_path}")
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- model_name_to_hf_name = {
- "dinov2_vits14": "dinov2-small",
- "dinov2_vitb14": "dinov2-base",
- "dinov2_vitl14": "dinov2-large",
- "dinov2_vitg14": "dinov2-giant",
- "dinov2_vits14_1layer": "dinov2-small-imagenet1k-1-layer",
- "dinov2_vitb14_1layer": "dinov2-base-imagenet1k-1-layer",
- "dinov2_vitl14_1layer": "dinov2-large-imagenet1k-1-layer",
- "dinov2_vitg14_1layer": "dinov2-giant-imagenet1k-1-layer",
- }
-
- name = model_name_to_hf_name[model_name]
- model.push_to_hub(f"facebook/{name}")
- processor.push_to_hub(f"facebook/{name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="dinov2_vitb14",
- type=str,
- choices=[
- "dinov2_vits14",
- "dinov2_vitb14",
- "dinov2_vitl14",
- "dinov2_vitg14",
- "dinov2_vits14_1layer",
- "dinov2_vitb14_1layer",
- "dinov2_vitl14_1layer",
- "dinov2_vitg14_1layer",
- ],
- help="Name of the model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
- )
- parser.add_argument(
- "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
- )
-
- args = parser.parse_args()
- convert_dinov2_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/dinov2/modeling_dinov2.py b/transformers/models/dinov2/modeling_dinov2.py
deleted file mode 100644
index c25022f6ec22d8d4afa1f926af9eb6e4d03adb35..0000000000000000000000000000000000000000
--- a/transformers/models/dinov2/modeling_dinov2.py
+++ /dev/null
@@ -1,856 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Meta AI and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DINOv2 model."""
-
-
-import collections.abc
-import math
-from typing import Dict, List, Optional, Set, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BackboneOutput,
- BaseModelOutput,
- BaseModelOutputWithPooling,
- ImageClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from ...utils.backbone_utils import BackboneMixin
-from .configuration_dinov2 import Dinov2Config
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "Dinov2Config"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "facebook/dinov2-base"
-_EXPECTED_OUTPUT_SHAPE = [1, 257, 768]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "facebook/dinov2-small-imagenet1k-1-layer"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import DINOV2_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class Dinov2Embeddings(nn.Module):
- """
- Construct the CLS token, mask token, position and patch embeddings.
- """
-
- def __init__(self, config: Dinov2Config) -> None:
- super().__init__()
-
- self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size))
- self.mask_token = nn.Parameter(torch.zeros(1, config.hidden_size))
- self.patch_embeddings = Dinov2PatchEmbeddings(config)
- num_patches = self.patch_embeddings.num_patches
- self.position_embeddings = nn.Parameter(torch.randn(1, num_patches + 1, config.hidden_size))
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.config = config
-
- def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
- """
- This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
- resolution images.
-
- Source:
- https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
- """
-
- num_patches = embeddings.shape[1] - 1
- num_positions = self.position_embeddings.shape[1] - 1
- if num_patches == num_positions and height == width:
- return self.position_embeddings
- class_pos_embed = self.position_embeddings[:, 0]
- patch_pos_embed = self.position_embeddings[:, 1:]
- dim = embeddings.shape[-1]
- height = height // self.config.patch_size
- width = width // self.config.patch_size
- # we add a small number to avoid floating point error in the interpolation
- # see discussion at https://github.com/facebookresearch/dino/issues/8
- height, width = height + 0.1, width + 0.1
- patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
- patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
- target_dtype = patch_pos_embed.dtype
- patch_pos_embed = nn.functional.interpolate(
- patch_pos_embed.to(dtype=torch.float32),
- scale_factor=(float(height / math.sqrt(num_positions)), float(width / math.sqrt(num_positions))),
- mode="bicubic",
- align_corners=False,
- ).to(dtype=target_dtype)
- if int(height) != patch_pos_embed.shape[-2] or int(width) != patch_pos_embed.shape[-1]:
- raise ValueError("Width or height does not match with the interpolated position embeddings")
- patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
- return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
-
- def forward(self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.Tensor] = None) -> torch.Tensor:
- batch_size, _, height, width = pixel_values.shape
- target_dtype = self.patch_embeddings.projection.weight.dtype
- embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype))
-
- if bool_masked_pos is not None:
- embeddings = torch.where(
- bool_masked_pos.unsqueeze(-1), self.mask_token.to(embeddings.dtype).unsqueeze(0), embeddings
- )
-
- # add the [CLS] token to the embedded patch tokens
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- embeddings = torch.cat((cls_tokens, embeddings), dim=1)
-
- # add positional encoding to each token
- embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
-
- embeddings = self.dropout(embeddings)
-
- return embeddings
-
-
-class Dinov2PatchEmbeddings(nn.Module):
- """
- This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
- `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
- Transformer.
- """
-
- def __init__(self, config):
- super().__init__()
- image_size, patch_size = config.image_size, config.patch_size
- num_channels, hidden_size = config.num_channels, config.hidden_size
-
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.num_patches = num_patches
-
- self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
-
- def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
- num_channels = pixel_values.shape[1]
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- f" Expected {self.num_channels} but got {num_channels}."
- )
- embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
- return embeddings
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->Dinov2
-class Dinov2SelfAttention(nn.Module):
- def __init__(self, config: Dinov2Config) -> None:
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
- f"heads {config.num_attention_heads}."
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
- self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
- self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self, hidden_states, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- mixed_query_layer = self.query(hidden_states)
-
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->Dinov2
-class Dinov2SelfOutput(nn.Module):
- """
- The residual connection is defined in Dinov2Layer instead of here (as is the case with other models), due to the
- layernorm applied before each block.
- """
-
- def __init__(self, config: Dinov2Config) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->Dinov2
-class Dinov2Attention(nn.Module):
- def __init__(self, config: Dinov2Config) -> None:
- super().__init__()
- self.attention = Dinov2SelfAttention(config)
- self.output = Dinov2SelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads: Set[int]) -> None:
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.attention.query = prune_linear_layer(self.attention.query, index)
- self.attention.key = prune_linear_layer(self.attention.key, index)
- self.attention.value = prune_linear_layer(self.attention.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
- self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- self_outputs = self.attention(hidden_states, head_mask, output_attentions)
-
- attention_output = self.output(self_outputs[0], hidden_states)
-
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-class Dinov2LayerScale(nn.Module):
- def __init__(self, config) -> None:
- super().__init__()
- self.lambda1 = nn.Parameter(config.layerscale_value * torch.ones(config.hidden_size))
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- return hidden_state * self.lambda1
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitDropPath
-class Dinov2DropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-class Dinov2MLP(nn.Module):
- def __init__(self, config) -> None:
- super().__init__()
- in_features = out_features = config.hidden_size
- hidden_features = int(config.hidden_size * config.mlp_ratio)
- self.fc1 = nn.Linear(in_features, hidden_features, bias=True)
- if isinstance(config.hidden_act, str):
- self.activation = ACT2FN[config.hidden_act]
- else:
- self.activation = config.hidden_act
- self.fc2 = nn.Linear(hidden_features, out_features, bias=True)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- hidden_state = self.fc1(hidden_state)
- hidden_state = self.activation(hidden_state)
- hidden_state = self.fc2(hidden_state)
- return hidden_state
-
-
-class Dinov2SwiGLUFFN(nn.Module):
- def __init__(self, config) -> None:
- super().__init__()
- in_features = out_features = config.hidden_size
- hidden_features = int(config.hidden_size * config.mlp_ratio)
- hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
-
- self.weights_in = nn.Linear(in_features, 2 * hidden_features, bias=True)
- self.weights_out = nn.Linear(hidden_features, out_features, bias=True)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- hidden_state = self.weights_in(hidden_state)
- x1, x2 = hidden_state.chunk(2, dim=-1)
- hidden = nn.functional.silu(x1) * x2
- return self.weights_out(hidden)
-
-
-class Dinov2Layer(nn.Module):
- """This corresponds to the Block class in the original implementation."""
-
- def __init__(self, config: Dinov2Config) -> None:
- super().__init__()
-
- self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.attention = Dinov2Attention(config)
- self.layer_scale1 = Dinov2LayerScale(config)
- self.drop_path = Dinov2DropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
-
- self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- if config.use_swiglu_ffn:
- self.mlp = Dinov2SwiGLUFFN(config)
- else:
- self.mlp = Dinov2MLP(config)
- self.layer_scale2 = Dinov2LayerScale(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- self_attention_outputs = self.attention(
- self.norm1(hidden_states), # in Dinov2, layernorm is applied before self-attention
- head_mask,
- output_attentions=output_attentions,
- )
- attention_output = self_attention_outputs[0]
-
- attention_output = self.layer_scale1(attention_output)
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- # first residual connection
- hidden_states = self.drop_path(attention_output) + hidden_states
-
- # in Dinov2, layernorm is also applied after self-attention
- layer_output = self.norm2(hidden_states)
- layer_output = self.mlp(layer_output)
- layer_output = self.layer_scale2(layer_output)
-
- # second residual connection
- layer_output = self.drop_path(layer_output) + hidden_states
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTEncoder with ViT->Dinov2
-class Dinov2Encoder(nn.Module):
- def __init__(self, config: Dinov2Config) -> None:
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([Dinov2Layer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[tuple, BaseModelOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- layer_head_mask,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions)
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-class Dinov2PreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = Dinov2Config
- base_model_prefix = "dinov2"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid
- # `trunc_normal_cpu` not implemented in `half` issues
- module.weight.data = nn.init.trunc_normal_(
- module.weight.data.to(torch.float32), mean=0.0, std=self.config.initializer_range
- ).to(module.weight.dtype)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
- elif isinstance(module, Dinov2Embeddings):
- module.position_embeddings.data = nn.init.trunc_normal_(
- module.position_embeddings.data.to(torch.float32),
- mean=0.0,
- std=self.config.initializer_range,
- ).to(module.position_embeddings.dtype)
-
- module.cls_token.data = nn.init.trunc_normal_(
- module.cls_token.data.to(torch.float32),
- mean=0.0,
- std=self.config.initializer_range,
- ).to(module.cls_token.dtype)
-
-
-DINOV2_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`Dinov2Config`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DINOV2_BASE_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`BitImageProcessor.preprocess`] for details.
-
- bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for
- pre-training.
-
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-DINOV2_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`BitImageProcessor.preprocess`] for details.
-
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DINOv2 Model transformer outputting raw hidden-states without any specific head on top.",
- DINOV2_START_DOCSTRING,
-)
-class Dinov2Model(Dinov2PreTrainedModel):
- def __init__(self, config: Dinov2Config):
- super().__init__(config)
- self.config = config
-
- self.embeddings = Dinov2Embeddings(config)
- self.encoder = Dinov2Encoder(config)
-
- self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self) -> Dinov2PatchEmbeddings:
- return self.embeddings.patch_embeddings
-
- def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(DINOV2_BASE_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPooling,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- bool_masked_pos: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPooling]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- embedding_output = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
-
- encoder_outputs = self.encoder(
- embedding_output,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
- pooled_output = sequence_output[:, 0, :]
-
- if not return_dict:
- head_outputs = (sequence_output, pooled_output)
- return head_outputs + encoder_outputs[1:]
-
- return BaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Dinov2 Model transformer with an image classification head on top (a linear layer on top of the final hidden state
- of the [CLS] token) e.g. for ImageNet.
- """,
- DINOV2_START_DOCSTRING,
-)
-class Dinov2ForImageClassification(Dinov2PreTrainedModel):
- def __init__(self, config: Dinov2Config) -> None:
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.dinov2 = Dinov2Model(config)
-
- # Classifier head
- self.classifier = (
- nn.Linear(config.hidden_size * 2, config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=ImageClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, ImageClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.dinov2(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0] # batch_size, sequence_length, hidden_size
-
- cls_token = sequence_output[:, 0]
- patch_tokens = sequence_output[:, 1:]
-
- linear_input = torch.cat([cls_token, patch_tokens.mean(dim=1)], dim=1)
-
- logits = self.classifier(linear_input)
-
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(logits.device)
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Dinov2 backbone, to be used with frameworks like DETR and MaskFormer.
- """,
- DINOV2_START_DOCSTRING,
-)
-class Dinov2Backbone(Dinov2PreTrainedModel, BackboneMixin):
- def __init__(self, config):
- super().__init__(config)
- super()._init_backbone(config)
-
- self.num_features = [config.hidden_size for _ in range(config.num_hidden_layers + 1)]
- self.embeddings = Dinov2Embeddings(config)
- self.encoder = Dinov2Encoder(config)
-
- self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self) -> Dinov2PatchEmbeddings:
- return self.embeddings.patch_embeddings
-
- @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.Tensor,
- output_hidden_states: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> BackboneOutput:
- """
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, AutoBackbone
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
- >>> model = AutoBackbone.from_pretrained(
- ... "facebook/dinov2-base", out_features=["stage2", "stage5", "stage8", "stage11"]
- ... )
-
- >>> inputs = processor(image, return_tensors="pt")
-
- >>> outputs = model(**inputs)
- >>> feature_maps = outputs.feature_maps
- >>> list(feature_maps[-1].shape)
- [1, 768, 16, 16]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-
- embedding_output = self.embeddings(pixel_values)
-
- outputs = self.encoder(
- embedding_output, output_hidden_states=True, output_attentions=output_attentions, return_dict=return_dict
- )
-
- hidden_states = outputs.hidden_states if return_dict else outputs[1]
-
- feature_maps = ()
- for stage, hidden_state in zip(self.stage_names, hidden_states):
- if stage in self.out_features:
- if self.config.apply_layernorm:
- hidden_state = self.layernorm(hidden_state)
- if self.config.reshape_hidden_states:
- hidden_state = hidden_state[:, 1:]
- # this was actually a bug in the original implementation that we copied here,
- # cause normally the order is height, width
- batch_size, _, height, width = pixel_values.shape
- patch_size = self.config.patch_size
- hidden_state = hidden_state.reshape(batch_size, height // patch_size, width // patch_size, -1)
- hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
- feature_maps += (hidden_state,)
-
- if not return_dict:
- if output_hidden_states:
- output = (feature_maps,) + outputs[1:]
- else:
- output = (feature_maps,) + outputs[2:]
- return output
-
- return BackboneOutput(
- feature_maps=feature_maps,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=outputs.attentions if output_attentions else None,
- )
diff --git a/transformers/models/distilbert/__init__.py b/transformers/models/distilbert/__init__.py
deleted file mode 100644
index 6a2756eb9d1c269e08446f9328120738196349d0..0000000000000000000000000000000000000000
--- a/transformers/models/distilbert/__init__.py
+++ /dev/null
@@ -1,166 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_flax_available,
- is_tf_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_distilbert": [
- "DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "DistilBertConfig",
- "DistilBertOnnxConfig",
- ],
- "tokenization_distilbert": ["DistilBertTokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_distilbert_fast"] = ["DistilBertTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_distilbert"] = [
- "DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DistilBertForMaskedLM",
- "DistilBertForMultipleChoice",
- "DistilBertForQuestionAnswering",
- "DistilBertForSequenceClassification",
- "DistilBertForTokenClassification",
- "DistilBertModel",
- "DistilBertPreTrainedModel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_distilbert"] = [
- "TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFDistilBertForMaskedLM",
- "TFDistilBertForMultipleChoice",
- "TFDistilBertForQuestionAnswering",
- "TFDistilBertForSequenceClassification",
- "TFDistilBertForTokenClassification",
- "TFDistilBertMainLayer",
- "TFDistilBertModel",
- "TFDistilBertPreTrainedModel",
- ]
-
-try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flax_distilbert"] = [
- "FlaxDistilBertForMaskedLM",
- "FlaxDistilBertForMultipleChoice",
- "FlaxDistilBertForQuestionAnswering",
- "FlaxDistilBertForSequenceClassification",
- "FlaxDistilBertForTokenClassification",
- "FlaxDistilBertModel",
- "FlaxDistilBertPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_distilbert import (
- DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
- DistilBertConfig,
- DistilBertOnnxConfig,
- )
- from .tokenization_distilbert import DistilBertTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_distilbert_fast import DistilBertTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_distilbert import (
- DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
- DistilBertForMaskedLM,
- DistilBertForMultipleChoice,
- DistilBertForQuestionAnswering,
- DistilBertForSequenceClassification,
- DistilBertForTokenClassification,
- DistilBertModel,
- DistilBertPreTrainedModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_distilbert import (
- TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFDistilBertForMaskedLM,
- TFDistilBertForMultipleChoice,
- TFDistilBertForQuestionAnswering,
- TFDistilBertForSequenceClassification,
- TFDistilBertForTokenClassification,
- TFDistilBertMainLayer,
- TFDistilBertModel,
- TFDistilBertPreTrainedModel,
- )
-
- try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flax_distilbert import (
- FlaxDistilBertForMaskedLM,
- FlaxDistilBertForMultipleChoice,
- FlaxDistilBertForQuestionAnswering,
- FlaxDistilBertForSequenceClassification,
- FlaxDistilBertForTokenClassification,
- FlaxDistilBertModel,
- FlaxDistilBertPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/distilbert/__pycache__/__init__.cpython-310.pyc b/transformers/models/distilbert/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index dd264e37128f4c4df3527ab7677d0e5ff3dc2f9b..0000000000000000000000000000000000000000
Binary files a/transformers/models/distilbert/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/distilbert/__pycache__/configuration_distilbert.cpython-310.pyc b/transformers/models/distilbert/__pycache__/configuration_distilbert.cpython-310.pyc
deleted file mode 100644
index d1ccc0e42b683eba379b2bacfcac042ed203acaf..0000000000000000000000000000000000000000
Binary files a/transformers/models/distilbert/__pycache__/configuration_distilbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/distilbert/__pycache__/modeling_distilbert.cpython-310.pyc b/transformers/models/distilbert/__pycache__/modeling_distilbert.cpython-310.pyc
deleted file mode 100644
index b45ab7044b6614bafc113f43f91c49a0744b1929..0000000000000000000000000000000000000000
Binary files a/transformers/models/distilbert/__pycache__/modeling_distilbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/distilbert/__pycache__/modeling_flax_distilbert.cpython-310.pyc b/transformers/models/distilbert/__pycache__/modeling_flax_distilbert.cpython-310.pyc
deleted file mode 100644
index 5ed6dec2ccdf764018752a94fdd16dbacf819e60..0000000000000000000000000000000000000000
Binary files a/transformers/models/distilbert/__pycache__/modeling_flax_distilbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/distilbert/__pycache__/modeling_tf_distilbert.cpython-310.pyc b/transformers/models/distilbert/__pycache__/modeling_tf_distilbert.cpython-310.pyc
deleted file mode 100644
index 73a767f9b4615a1c70a7e978be80fe26deb8543a..0000000000000000000000000000000000000000
Binary files a/transformers/models/distilbert/__pycache__/modeling_tf_distilbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/distilbert/__pycache__/tokenization_distilbert.cpython-310.pyc b/transformers/models/distilbert/__pycache__/tokenization_distilbert.cpython-310.pyc
deleted file mode 100644
index 09e446ebbd9bf15930cfa4c5032238e07fc991c6..0000000000000000000000000000000000000000
Binary files a/transformers/models/distilbert/__pycache__/tokenization_distilbert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/distilbert/__pycache__/tokenization_distilbert_fast.cpython-310.pyc b/transformers/models/distilbert/__pycache__/tokenization_distilbert_fast.cpython-310.pyc
deleted file mode 100644
index e111f16504410e07394bdc9dbe261ffcc8f96d57..0000000000000000000000000000000000000000
Binary files a/transformers/models/distilbert/__pycache__/tokenization_distilbert_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/distilbert/configuration_distilbert.py b/transformers/models/distilbert/configuration_distilbert.py
deleted file mode 100644
index 5f6b004dc0bbb978f7f4efea55ab4f643a0a9dc9..0000000000000000000000000000000000000000
--- a/transformers/models/distilbert/configuration_distilbert.py
+++ /dev/null
@@ -1,140 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DistilBERT model configuration"""
-from collections import OrderedDict
-from typing import Mapping
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DistilBertConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DistilBertModel`] or a [`TFDistilBertModel`]. It
- is used to instantiate a DistilBERT model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the DistilBERT
- [distilbert-base-uncased](https://huggingface.co/distilbert-base-uncased) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the DistilBERT model. Defines the number of different tokens that can be represented by
- the `inputs_ids` passed when calling [`DistilBertModel`] or [`TFDistilBertModel`].
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- sinusoidal_pos_embds (`boolean`, *optional*, defaults to `False`):
- Whether to use sinusoidal positional embeddings.
- n_layers (`int`, *optional*, defaults to 6):
- Number of hidden layers in the Transformer encoder.
- n_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- dim (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- hidden_dim (`int`, *optional*, defaults to 3072):
- The size of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
- dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- activation (`str` or `Callable`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- qa_dropout (`float`, *optional*, defaults to 0.1):
- The dropout probabilities used in the question answering model [`DistilBertForQuestionAnswering`].
- seq_classif_dropout (`float`, *optional*, defaults to 0.2):
- The dropout probabilities used in the sequence classification and the multiple choice model
- [`DistilBertForSequenceClassification`].
-
- Examples:
-
- ```python
- >>> from transformers import DistilBertConfig, DistilBertModel
-
- >>> # Initializing a DistilBERT configuration
- >>> configuration = DistilBertConfig()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = DistilBertModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "distilbert"
- attribute_map = {
- "hidden_size": "dim",
- "num_attention_heads": "n_heads",
- "num_hidden_layers": "n_layers",
- }
-
- def __init__(
- self,
- vocab_size=30522,
- max_position_embeddings=512,
- sinusoidal_pos_embds=False,
- n_layers=6,
- n_heads=12,
- dim=768,
- hidden_dim=4 * 768,
- dropout=0.1,
- attention_dropout=0.1,
- activation="gelu",
- initializer_range=0.02,
- qa_dropout=0.1,
- seq_classif_dropout=0.2,
- pad_token_id=0,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.sinusoidal_pos_embds = sinusoidal_pos_embds
- self.n_layers = n_layers
- self.n_heads = n_heads
- self.dim = dim
- self.hidden_dim = hidden_dim
- self.dropout = dropout
- self.attention_dropout = attention_dropout
- self.activation = activation
- self.initializer_range = initializer_range
- self.qa_dropout = qa_dropout
- self.seq_classif_dropout = seq_classif_dropout
- super().__init__(**kwargs, pad_token_id=pad_token_id)
-
-
-class DistilBertOnnxConfig(OnnxConfig):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- if self.task == "multiple-choice":
- dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
- else:
- dynamic_axis = {0: "batch", 1: "sequence"}
- return OrderedDict(
- [
- ("input_ids", dynamic_axis),
- ("attention_mask", dynamic_axis),
- ]
- )
diff --git a/transformers/models/distilbert/modeling_distilbert.py b/transformers/models/distilbert/modeling_distilbert.py
deleted file mode 100644
index 3a65e0296116dca625ef343ac027243cca3dc392..0000000000000000000000000000000000000000
--- a/transformers/models/distilbert/modeling_distilbert.py
+++ /dev/null
@@ -1,1384 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
- PyTorch DistilBERT model adapted in part from Facebook, Inc XLM model (https://github.com/facebookresearch/XLM) and in
- part from HuggingFace PyTorch version of Google AI Bert model (https://github.com/google-research/bert)
-"""
-
-
-import math
-from typing import Dict, List, Optional, Set, Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import get_activation
-from ...configuration_utils import PretrainedConfig
-from ...integrations.deepspeed import is_deepspeed_zero3_enabled
-from ...modeling_outputs import (
- BaseModelOutput,
- MaskedLMOutput,
- MultipleChoiceModelOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- logging,
- replace_return_docstrings,
-)
-from .configuration_distilbert import DistilBertConfig
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-
-logger = logging.get_logger(__name__)
-_CHECKPOINT_FOR_DOC = "distilbert-base-uncased"
-_CONFIG_FOR_DOC = "DistilBertConfig"
-
-
-from ..deprecated._archive_maps import DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# UTILS AND BUILDING BLOCKS OF THE ARCHITECTURE #
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-def create_sinusoidal_embeddings(n_pos: int, dim: int, out: torch.Tensor):
- if is_deepspeed_zero3_enabled():
- import deepspeed
-
- with deepspeed.zero.GatheredParameters(out, modifier_rank=0):
- if torch.distributed.get_rank() == 0:
- _create_sinusoidal_embeddings(n_pos=n_pos, dim=dim, out=out)
- else:
- _create_sinusoidal_embeddings(n_pos=n_pos, dim=dim, out=out)
-
-
-def _create_sinusoidal_embeddings(n_pos: int, dim: int, out: torch.Tensor):
- position_enc = np.array([[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)])
- out.requires_grad = False
- out[:, 0::2] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))
- out[:, 1::2] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))
- out.detach_()
-
-
-class Embeddings(nn.Module):
- def __init__(self, config: PretrainedConfig):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.dim, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.dim)
-
- self.LayerNorm = nn.LayerNorm(config.dim, eps=1e-12)
- self.dropout = nn.Dropout(config.dropout)
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
-
- def forward(self, input_ids: torch.Tensor, input_embeds: Optional[torch.Tensor] = None) -> torch.Tensor:
- """
- Parameters:
- input_ids (torch.Tensor):
- torch.tensor(bs, max_seq_length) The token ids to embed.
- input_embeds (*optional*, torch.Tensor):
- The pre-computed word embeddings. Can only be passed if the input ids are `None`.
-
-
- Returns: torch.tensor(bs, max_seq_length, dim) The embedded tokens (plus position embeddings, no token_type
- embeddings)
- """
- if input_ids is not None:
- input_embeds = self.word_embeddings(input_ids) # (bs, max_seq_length, dim)
-
- seq_length = input_embeds.size(1)
-
- # Setting the position-ids to the registered buffer in constructor, it helps
- # when tracing the model without passing position-ids, solves
- # isues similar to issue #5664
- if hasattr(self, "position_ids"):
- position_ids = self.position_ids[:, :seq_length]
- else:
- position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device) # (max_seq_length)
- position_ids = position_ids.unsqueeze(0).expand_as(input_ids) # (bs, max_seq_length)
-
- position_embeddings = self.position_embeddings(position_ids) # (bs, max_seq_length, dim)
-
- embeddings = input_embeds + position_embeddings # (bs, max_seq_length, dim)
- embeddings = self.LayerNorm(embeddings) # (bs, max_seq_length, dim)
- embeddings = self.dropout(embeddings) # (bs, max_seq_length, dim)
- return embeddings
-
-
-class MultiHeadSelfAttention(nn.Module):
- def __init__(self, config: PretrainedConfig):
- super().__init__()
- self.config = config
-
- self.n_heads = config.n_heads
- self.dim = config.dim
- self.dropout = nn.Dropout(p=config.attention_dropout)
- self.is_causal = False
-
- # Have an even number of multi heads that divide the dimensions
- if self.dim % self.n_heads != 0:
- # Raise value errors for even multi-head attention nodes
- raise ValueError(f"self.n_heads: {self.n_heads} must divide self.dim: {self.dim} evenly")
-
- self.q_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
- self.k_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
- self.v_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
- self.out_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
-
- self.pruned_heads: Set[int] = set()
- self.attention_head_size = self.dim // self.n_heads
-
- def prune_heads(self, heads: List[int]):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.n_heads, self.attention_head_size, self.pruned_heads
- )
- # Prune linear layers
- self.q_lin = prune_linear_layer(self.q_lin, index)
- self.k_lin = prune_linear_layer(self.k_lin, index)
- self.v_lin = prune_linear_layer(self.v_lin, index)
- self.out_lin = prune_linear_layer(self.out_lin, index, dim=1)
- # Update hyper params
- self.n_heads = self.n_heads - len(heads)
- self.dim = self.attention_head_size * self.n_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- query: torch.Tensor,
- key: torch.Tensor,
- value: torch.Tensor,
- mask: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, ...]:
- """
- Parameters:
- query: torch.tensor(bs, seq_length, dim)
- key: torch.tensor(bs, seq_length, dim)
- value: torch.tensor(bs, seq_length, dim)
- mask: torch.tensor(bs, seq_length)
-
- Returns:
- weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
- seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
- """
- bs, q_length, dim = query.size()
- k_length = key.size(1)
- # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
- # assert key.size() == value.size()
-
- dim_per_head = self.dim // self.n_heads
-
- mask_reshp = (bs, 1, 1, k_length)
-
- def shape(x: torch.Tensor) -> torch.Tensor:
- """separate heads"""
- return x.view(bs, -1, self.n_heads, dim_per_head).transpose(1, 2)
-
- def unshape(x: torch.Tensor) -> torch.Tensor:
- """group heads"""
- return x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head)
-
- q = shape(self.q_lin(query)) # (bs, n_heads, q_length, dim_per_head)
- k = shape(self.k_lin(key)) # (bs, n_heads, k_length, dim_per_head)
- v = shape(self.v_lin(value)) # (bs, n_heads, k_length, dim_per_head)
-
- q = q / math.sqrt(dim_per_head) # (bs, n_heads, q_length, dim_per_head)
- scores = torch.matmul(q, k.transpose(2, 3)) # (bs, n_heads, q_length, k_length)
- mask = (mask == 0).view(mask_reshp).expand_as(scores) # (bs, n_heads, q_length, k_length)
- scores = scores.masked_fill(
- mask, torch.tensor(torch.finfo(scores.dtype).min)
- ) # (bs, n_heads, q_length, k_length)
-
- weights = nn.functional.softmax(scores, dim=-1) # (bs, n_heads, q_length, k_length)
- weights = self.dropout(weights) # (bs, n_heads, q_length, k_length)
-
- # Mask heads if we want to
- if head_mask is not None:
- weights = weights * head_mask
-
- context = torch.matmul(weights, v) # (bs, n_heads, q_length, dim_per_head)
- context = unshape(context) # (bs, q_length, dim)
- context = self.out_lin(context) # (bs, q_length, dim)
-
- if output_attentions:
- return (context, weights)
- else:
- return (context,)
-
-
-class DistilBertFlashAttention2(MultiHeadSelfAttention):
- """
- DistilBert flash attention module. This module inherits from `MultiHeadSelfAttention` as the weights of the module
- stays untouched. The only required change would be on the forward pass where it needs to correctly call the public
- API of flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- query: torch.Tensor,
- key: torch.Tensor,
- value: torch.Tensor,
- mask: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, ...]:
- """
- Parameters:
- query: torch.tensor(bs, seq_length, dim)
- key: torch.tensor(bs, seq_length, dim)
- value: torch.tensor(bs, seq_length, dim)
- mask: torch.tensor(bs, seq_length)
-
- Returns:
- weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
- seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
- """
- batch_size, q_length, dim = query.size()
-
- dim_per_head = self.dim // self.n_heads
-
- def reshape(x: torch.Tensor) -> torch.Tensor:
- """separate heads"""
- return x.view(batch_size, -1, self.n_heads, dim_per_head)
-
- # Flash attention requires the input to have the shape
- # batch_size x seq_length x head_dim x hidden_dim
- query_states = reshape(self.q_lin(query))
- key_states = reshape(self.k_lin(key))
- value_states = reshape(self.v_lin(value))
-
- attn_dropout = self.config.attention_dropout if self.training else 0.0
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in the correct dtype just to be sure everything works as expected.
- # This might slowdown training & inference so it is recommended to not cast the LayerNorms
- # in fp32. (LlamaRMSNorm handles it correctly)
-
- if query_states.dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.q_lin.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
-
- query_states = query_states.to(target_dtype)
- key_states = key_states.to(target_dtype)
- value_states = value_states.to(target_dtype)
-
- attn_weights = self._flash_attention_forward(
- query_states, key_states, value_states, mask, q_length, dropout=attn_dropout
- )
-
- attn_weights_reshaped = attn_weights.reshape(batch_size, q_length, self.n_heads * dim_per_head)
- attn_output = self.out_lin(attn_weights_reshaped)
-
- if output_attentions:
- return (attn_output, attn_weights)
- else:
- return (attn_output,)
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward with causal=True->causal=False
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input with num_heads->n_heads
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.n_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-class FFN(nn.Module):
- def __init__(self, config: PretrainedConfig):
- super().__init__()
- self.dropout = nn.Dropout(p=config.dropout)
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.lin1 = nn.Linear(in_features=config.dim, out_features=config.hidden_dim)
- self.lin2 = nn.Linear(in_features=config.hidden_dim, out_features=config.dim)
- self.activation = get_activation(config.activation)
-
- def forward(self, input: torch.Tensor) -> torch.Tensor:
- return apply_chunking_to_forward(self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, input)
-
- def ff_chunk(self, input: torch.Tensor) -> torch.Tensor:
- x = self.lin1(input)
- x = self.activation(x)
- x = self.lin2(x)
- x = self.dropout(x)
- return x
-
-
-DISTILBERT_ATTENTION_CLASSES = {
- "eager": MultiHeadSelfAttention,
- "flash_attention_2": DistilBertFlashAttention2,
-}
-
-
-class TransformerBlock(nn.Module):
- def __init__(self, config: PretrainedConfig):
- super().__init__()
-
- # Have an even number of Configure multi-heads
- if config.dim % config.n_heads != 0:
- raise ValueError(f"config.n_heads {config.n_heads} must divide config.dim {config.dim} evenly")
-
- self.attention = DISTILBERT_ATTENTION_CLASSES[config._attn_implementation](config)
- self.sa_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12)
-
- self.ffn = FFN(config)
- self.output_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12)
-
- def forward(
- self,
- x: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, ...]:
- """
- Parameters:
- x: torch.tensor(bs, seq_length, dim)
- attn_mask: torch.tensor(bs, seq_length)
-
- Returns:
- sa_weights: torch.tensor(bs, n_heads, seq_length, seq_length) The attention weights ffn_output:
- torch.tensor(bs, seq_length, dim) The output of the transformer block contextualization.
- """
- # Self-Attention
- sa_output = self.attention(
- query=x,
- key=x,
- value=x,
- mask=attn_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- )
- if output_attentions:
- sa_output, sa_weights = sa_output # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)
- else: # To handle these `output_attentions` or `output_hidden_states` cases returning tuples
- if type(sa_output) != tuple:
- raise TypeError(f"sa_output must be a tuple but it is {type(sa_output)} type")
-
- sa_output = sa_output[0]
- sa_output = self.sa_layer_norm(sa_output + x) # (bs, seq_length, dim)
-
- # Feed Forward Network
- ffn_output = self.ffn(sa_output) # (bs, seq_length, dim)
- ffn_output: torch.Tensor = self.output_layer_norm(ffn_output + sa_output) # (bs, seq_length, dim)
-
- output = (ffn_output,)
- if output_attentions:
- output = (sa_weights,) + output
- return output
-
-
-class Transformer(nn.Module):
- def __init__(self, config: PretrainedConfig):
- super().__init__()
- self.n_layers = config.n_layers
- self.layer = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- x: torch.Tensor,
- attn_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: Optional[bool] = None,
- ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]: # docstyle-ignore
- """
- Parameters:
- x: torch.tensor(bs, seq_length, dim) Input sequence embedded.
- attn_mask: torch.tensor(bs, seq_length) Attention mask on the sequence.
-
- Returns:
- hidden_state: torch.tensor(bs, seq_length, dim) Sequence of hidden states in the last (top)
- layer all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)]
- Tuple of length n_layers with the hidden states from each layer.
- Optional: only if output_hidden_states=True
- all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)]
- Tuple of length n_layers with the attention weights from each layer
- Optional: only if output_attentions=True
- """
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- hidden_state = x
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_state,)
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_state,
- attn_mask,
- head_mask[i],
- output_attentions,
- )
- else:
- layer_outputs = layer_module(
- hidden_state,
- attn_mask,
- head_mask[i],
- output_attentions,
- )
-
- hidden_state = layer_outputs[-1]
-
- if output_attentions:
- if len(layer_outputs) != 2:
- raise ValueError(f"The length of the layer_outputs should be 2, but it is {len(layer_outputs)}")
-
- attentions = layer_outputs[0]
- all_attentions = all_attentions + (attentions,)
- else:
- if len(layer_outputs) != 1:
- raise ValueError(f"The length of the layer_outputs should be 1, but it is {len(layer_outputs)}")
-
- # Add last layer
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_state,)
-
- if not return_dict:
- return tuple(v for v in [hidden_state, all_hidden_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_state, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
-
-# INTERFACE FOR ENCODER AND TASK SPECIFIC MODEL #
-class DistilBertPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DistilBertConfig
- load_tf_weights = None
- base_model_prefix = "distilbert"
- supports_gradient_checkpointing = True
- _supports_flash_attn_2 = True
-
- def _init_weights(self, module: nn.Module):
- """Initialize the weights."""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
- elif isinstance(module, Embeddings) and self.config.sinusoidal_pos_embds:
- create_sinusoidal_embeddings(
- self.config.max_position_embeddings, self.config.dim, module.position_embeddings.weight
- )
-
-
-DISTILBERT_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`DistilBertConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DISTILBERT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DistilBERT encoder/transformer outputting raw hidden-states without any specific head on top.",
- DISTILBERT_START_DOCSTRING,
-)
-class DistilBertModel(DistilBertPreTrainedModel):
- def __init__(self, config: PretrainedConfig):
- super().__init__(config)
-
- self.embeddings = Embeddings(config) # Embeddings
- self.transformer = Transformer(config) # Encoder
- self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_position_embeddings(self) -> nn.Embedding:
- """
- Returns the position embeddings
- """
- return self.embeddings.position_embeddings
-
- def resize_position_embeddings(self, new_num_position_embeddings: int):
- """
- Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
-
- Arguments:
- new_num_position_embeddings (`int`):
- The number of new position embedding matrix. If position embeddings are learned, increasing the size
- will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
- end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
- size will add correct vectors at the end following the position encoding algorithm, whereas reducing
- the size will remove vectors from the end.
- """
- num_position_embeds_diff = new_num_position_embeddings - self.config.max_position_embeddings
-
- # no resizing needs to be done if the length stays the same
- if num_position_embeds_diff == 0:
- return
-
- logger.info(f"Setting `config.max_position_embeddings={new_num_position_embeddings}`...")
- self.config.max_position_embeddings = new_num_position_embeddings
-
- old_position_embeddings_weight = self.embeddings.position_embeddings.weight.clone()
-
- self.embeddings.position_embeddings = nn.Embedding(self.config.max_position_embeddings, self.config.dim)
-
- if self.config.sinusoidal_pos_embds:
- create_sinusoidal_embeddings(
- n_pos=self.config.max_position_embeddings, dim=self.config.dim, out=self.position_embeddings.weight
- )
- else:
- with torch.no_grad():
- if num_position_embeds_diff > 0:
- self.embeddings.position_embeddings.weight[:-num_position_embeds_diff] = nn.Parameter(
- old_position_embeddings_weight
- )
- else:
- self.embeddings.position_embeddings.weight = nn.Parameter(
- old_position_embeddings_weight[:num_position_embeds_diff]
- )
- # move position_embeddings to correct device
- self.embeddings.position_embeddings.to(self.device)
-
- def get_input_embeddings(self) -> nn.Embedding:
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, new_embeddings: nn.Embedding):
- self.embeddings.word_embeddings = new_embeddings
-
- def _prune_heads(self, heads_to_prune: Dict[int, List[List[int]]]):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.transformer.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- # Prepare head mask if needed
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- embeddings = self.embeddings(input_ids, inputs_embeds) # (bs, seq_length, dim)
-
- if self._use_flash_attention_2:
- attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
- else:
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device) # (bs, seq_length)
-
- return self.transformer(
- x=embeddings,
- attn_mask=attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
-
-@add_start_docstrings(
- """DistilBert Model with a `masked language modeling` head on top.""",
- DISTILBERT_START_DOCSTRING,
-)
-class DistilBertForMaskedLM(DistilBertPreTrainedModel):
- _tied_weights_keys = ["vocab_projector.weight"]
-
- def __init__(self, config: PretrainedConfig):
- super().__init__(config)
-
- self.activation = get_activation(config.activation)
-
- self.distilbert = DistilBertModel(config)
- self.vocab_transform = nn.Linear(config.dim, config.dim)
- self.vocab_layer_norm = nn.LayerNorm(config.dim, eps=1e-12)
- self.vocab_projector = nn.Linear(config.dim, config.vocab_size)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- self.mlm_loss_fct = nn.CrossEntropyLoss()
-
- def get_position_embeddings(self) -> nn.Embedding:
- """
- Returns the position embeddings
- """
- return self.distilbert.get_position_embeddings()
-
- def resize_position_embeddings(self, new_num_position_embeddings: int):
- """
- Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
-
- Arguments:
- new_num_position_embeddings (`int`):
- The number of new position embedding matrix. If position embeddings are learned, increasing the size
- will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
- end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
- size will add correct vectors at the end following the position encoding algorithm, whereas reducing
- the size will remove vectors from the end.
- """
- self.distilbert.resize_position_embeddings(new_num_position_embeddings)
-
- def get_output_embeddings(self) -> nn.Module:
- return self.vocab_projector
-
- def set_output_embeddings(self, new_embeddings: nn.Module):
- self.vocab_projector = new_embeddings
-
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[MaskedLMOutput, Tuple[torch.Tensor, ...]]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- dlbrt_output = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = dlbrt_output[0] # (bs, seq_length, dim)
- prediction_logits = self.vocab_transform(hidden_states) # (bs, seq_length, dim)
- prediction_logits = self.activation(prediction_logits) # (bs, seq_length, dim)
- prediction_logits = self.vocab_layer_norm(prediction_logits) # (bs, seq_length, dim)
- prediction_logits = self.vocab_projector(prediction_logits) # (bs, seq_length, vocab_size)
-
- mlm_loss = None
- if labels is not None:
- mlm_loss = self.mlm_loss_fct(prediction_logits.view(-1, prediction_logits.size(-1)), labels.view(-1))
-
- if not return_dict:
- output = (prediction_logits,) + dlbrt_output[1:]
- return ((mlm_loss,) + output) if mlm_loss is not None else output
-
- return MaskedLMOutput(
- loss=mlm_loss,
- logits=prediction_logits,
- hidden_states=dlbrt_output.hidden_states,
- attentions=dlbrt_output.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DistilBert Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- DISTILBERT_START_DOCSTRING,
-)
-class DistilBertForSequenceClassification(DistilBertPreTrainedModel):
- def __init__(self, config: PretrainedConfig):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
-
- self.distilbert = DistilBertModel(config)
- self.pre_classifier = nn.Linear(config.dim, config.dim)
- self.classifier = nn.Linear(config.dim, config.num_labels)
- self.dropout = nn.Dropout(config.seq_classif_dropout)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_position_embeddings(self) -> nn.Embedding:
- """
- Returns the position embeddings
- """
- return self.distilbert.get_position_embeddings()
-
- def resize_position_embeddings(self, new_num_position_embeddings: int):
- """
- Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
-
- Arguments:
- new_num_position_embeddings (`int`):
- The number of new position embedding matrix. If position embeddings are learned, increasing the size
- will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
- end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
- size will add correct vectors at the end following the position encoding algorithm, whereas reducing
- the size will remove vectors from the end.
- """
- self.distilbert.resize_position_embeddings(new_num_position_embeddings)
-
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[SequenceClassifierOutput, Tuple[torch.Tensor, ...]]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- distilbert_output = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_state = distilbert_output[0] # (bs, seq_len, dim)
- pooled_output = hidden_state[:, 0] # (bs, dim)
- pooled_output = self.pre_classifier(pooled_output) # (bs, dim)
- pooled_output = nn.ReLU()(pooled_output) # (bs, dim)
- pooled_output = self.dropout(pooled_output) # (bs, dim)
- logits = self.classifier(pooled_output) # (bs, num_labels)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + distilbert_output[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=distilbert_output.hidden_states,
- attentions=distilbert_output.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a
- linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- DISTILBERT_START_DOCSTRING,
-)
-class DistilBertForQuestionAnswering(DistilBertPreTrainedModel):
- def __init__(self, config: PretrainedConfig):
- super().__init__(config)
-
- self.distilbert = DistilBertModel(config)
- self.qa_outputs = nn.Linear(config.dim, config.num_labels)
- if config.num_labels != 2:
- raise ValueError(f"config.num_labels should be 2, but it is {config.num_labels}")
-
- self.dropout = nn.Dropout(config.qa_dropout)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_position_embeddings(self) -> nn.Embedding:
- """
- Returns the position embeddings
- """
- return self.distilbert.get_position_embeddings()
-
- def resize_position_embeddings(self, new_num_position_embeddings: int):
- """
- Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
-
- Arguments:
- new_num_position_embeddings (`int`):
- The number of new position embedding matrix. If position embeddings are learned, increasing the size
- will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
- end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
- size will add correct vectors at the end following the position encoding algorithm, whereas reducing
- the size will remove vectors from the end.
- """
- self.distilbert.resize_position_embeddings(new_num_position_embeddings)
-
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[QuestionAnsweringModelOutput, Tuple[torch.Tensor, ...]]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- distilbert_output = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = distilbert_output[0] # (bs, max_query_len, dim)
-
- hidden_states = self.dropout(hidden_states) # (bs, max_query_len, dim)
- logits = self.qa_outputs(hidden_states) # (bs, max_query_len, 2)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous() # (bs, max_query_len)
- end_logits = end_logits.squeeze(-1).contiguous() # (bs, max_query_len)
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + distilbert_output[1:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=distilbert_output.hidden_states,
- attentions=distilbert_output.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g.
- for Named-Entity-Recognition (NER) tasks.
- """,
- DISTILBERT_START_DOCSTRING,
-)
-class DistilBertForTokenClassification(DistilBertPreTrainedModel):
- def __init__(self, config: PretrainedConfig):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.distilbert = DistilBertModel(config)
- self.dropout = nn.Dropout(config.dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_position_embeddings(self) -> nn.Embedding:
- """
- Returns the position embeddings
- """
- return self.distilbert.get_position_embeddings()
-
- def resize_position_embeddings(self, new_num_position_embeddings: int):
- """
- Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
-
- Arguments:
- new_num_position_embeddings (`int`):
- The number of new position embedding matrix. If position embeddings are learned, increasing the size
- will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
- end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
- size will add correct vectors at the end following the position encoding algorithm, whereas reducing
- the size will remove vectors from the end.
- """
- self.distilbert.resize_position_embeddings(new_num_position_embeddings)
-
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[TokenClassifierOutput, Tuple[torch.Tensor, ...]]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.distilbert(
- input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and
- a softmax) e.g. for RocStories/SWAG tasks.
- """,
- DISTILBERT_START_DOCSTRING,
-)
-class DistilBertForMultipleChoice(DistilBertPreTrainedModel):
- def __init__(self, config: PretrainedConfig):
- super().__init__(config)
-
- self.distilbert = DistilBertModel(config)
- self.pre_classifier = nn.Linear(config.dim, config.dim)
- self.classifier = nn.Linear(config.dim, 1)
- self.dropout = nn.Dropout(config.seq_classif_dropout)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_position_embeddings(self) -> nn.Embedding:
- """
- Returns the position embeddings
- """
- return self.distilbert.get_position_embeddings()
-
- def resize_position_embeddings(self, new_num_position_embeddings: int):
- """
- Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
-
- Arguments:
- new_num_position_embeddings (`int`)
- The number of new position embeddings. If position embeddings are learned, increasing the size will add
- newly initialized vectors at the end, whereas reducing the size will remove vectors from the end. If
- position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the size will
- add correct vectors at the end following the position encoding algorithm, whereas reducing the size
- will remove vectors from the end.
- """
- self.distilbert.resize_position_embeddings(new_num_position_embeddings)
-
- @add_start_docstrings_to_model_forward(
- DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
- )
- @replace_return_docstrings(output_type=MultipleChoiceModelOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[MultipleChoiceModelOutput, Tuple[torch.Tensor, ...]]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoTokenizer, DistilBertForMultipleChoice
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased")
- >>> model = DistilBertForMultipleChoice.from_pretrained("distilbert-base-cased")
-
- >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
- >>> choice0 = "It is eaten with a fork and a knife."
- >>> choice1 = "It is eaten while held in the hand."
- >>> labels = torch.tensor(0).unsqueeze(0) # choice0 is correct (according to Wikipedia ;)), batch size 1
-
- >>> encoding = tokenizer([[prompt, choice0], [prompt, choice1]], return_tensors="pt", padding=True)
- >>> outputs = model(**{k: v.unsqueeze(0) for k, v in encoding.items()}, labels=labels) # batch size is 1
-
- >>> # the linear classifier still needs to be trained
- >>> loss = outputs.loss
- >>> logits = outputs.logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.distilbert(
- input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_state = outputs[0] # (bs * num_choices, seq_len, dim)
- pooled_output = hidden_state[:, 0] # (bs * num_choices, dim)
- pooled_output = self.pre_classifier(pooled_output) # (bs * num_choices, dim)
- pooled_output = nn.ReLU()(pooled_output) # (bs * num_choices, dim)
- pooled_output = self.dropout(pooled_output) # (bs * num_choices, dim)
- logits = self.classifier(pooled_output) # (bs * num_choices, 1)
-
- reshaped_logits = logits.view(-1, num_choices) # (bs, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/distilbert/modeling_flax_distilbert.py b/transformers/models/distilbert/modeling_flax_distilbert.py
deleted file mode 100644
index d3c48c077adc529a1e942fcbce1999c2d0f8d524..0000000000000000000000000000000000000000
--- a/transformers/models/distilbert/modeling_flax_distilbert.py
+++ /dev/null
@@ -1,895 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import math
-from typing import Callable, Optional, Tuple
-
-import flax.linen as nn
-import jax
-import jax.numpy as jnp
-import numpy as np
-from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
-from flax.traverse_util import flatten_dict, unflatten_dict
-from jax import lax
-
-from ...modeling_flax_outputs import (
- FlaxBaseModelOutput,
- FlaxMaskedLMOutput,
- FlaxMultipleChoiceModelOutput,
- FlaxQuestionAnsweringModelOutput,
- FlaxSequenceClassifierOutput,
- FlaxTokenClassifierOutput,
-)
-from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring, overwrite_call_docstring
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_distilbert import DistilBertConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "distilbert-base-uncased"
-_CONFIG_FOR_DOC = "DistilBertConfig"
-
-
-FLAX_DISTILBERT_START_DOCSTRING = r"""
-
- This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading, saving and converting weights from PyTorch models)
-
- This model is also a
- [flax.linen.Module](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html) subclass. Use it as
- a regular Flax linen Module and refer to the Flax documentation for all matter related to general usage and
- behavior.
-
- Finally, this model supports inherent JAX features such as:
-
- - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
-
- Parameters:
- config ([`DistilBertConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DISTILBERT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`numpy.ndarray` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`numpy.ndarray` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-def get_angles(pos, i, d_model):
- angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
- return pos * angle_rates
-
-
-def positional_encoding(position, d_model):
- # create the sinusoidal pattern for the positional encoding
- angle_rads = get_angles(np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :], d_model)
-
- # apply sin to even indices in the array; 2i
- angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
-
- # apply cos to odd indices in the array; 2i+1
- angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
-
- pos_encoding = angle_rads[np.newaxis, ...]
-
- return jnp.array(pos_encoding)
-
-
-class FlaxEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.word_embeddings = nn.Embed(
- self.config.vocab_size,
- self.config.dim,
- embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- if not self.config.sinusoidal_pos_embds:
- self.position_embeddings = nn.Embed(
- self.config.max_position_embeddings,
- self.config.dim,
- embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- else:
- self.pos_encoding = positional_encoding(self.config.max_position_embeddings, self.config.dim)
- self.LayerNorm = nn.LayerNorm(epsilon=1e-12, dtype=self.dtype)
- self.dropout = nn.Dropout(rate=self.config.dropout)
-
- def __call__(self, input_ids, deterministic: bool = True):
- # Embed
- batch_size, seq_length = input_ids.shape
- inputs_embeds = self.word_embeddings(input_ids.astype("i4"))
- if not self.config.sinusoidal_pos_embds:
- position_ids = jnp.arange(seq_length).astype("i4")
- position_ids = jnp.broadcast_to(position_ids, shape=(batch_size, seq_length))
- position_embeds = self.position_embeddings(position_ids.astype("i4"))
- else:
- position_embeds = self.pos_encoding[:, :seq_length, :]
- # explicitly cast the positions here, since self.embed_positions are not registered as parameters
- position_embeds = position_embeds.astype(inputs_embeds.dtype)
-
- # Sum all embeddings
- hidden_states = inputs_embeds + position_embeds
-
- # Layer Norm
- hidden_states = self.LayerNorm(hidden_states)
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- return hidden_states
-
-
-class FlaxMultiHeadSelfAttention(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.n_heads = self.config.n_heads
- self.dim = self.config.dim
- self.dropout = nn.Dropout(rate=self.config.attention_dropout)
-
- if not (self.dim % self.n_heads == 0):
- raise ValueError(f"Hidden size {self.dim} not dividable by number of heads {self.n_heads}")
-
- self.q_lin = nn.Dense(
- self.dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.k_lin = nn.Dense(
- self.dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.v_lin = nn.Dense(
- self.dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.out_lin = nn.Dense(
- self.dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
-
- def __call__(
- self,
- query,
- key,
- value,
- mask,
- deterministic: bool = True,
- output_attentions: bool = False,
- ):
- bs, q_len, dim = query.shape
- k_len = key.shape[1]
- # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
- # assert key.size() == value.size()
-
- dim_per_head = self.dim // self.n_heads
-
- mask_reshp = (bs, 1, 1, k_len)
-
- def shape(x):
- """separate heads"""
- return x.reshape(bs, -1, self.n_heads, dim_per_head).transpose(0, 2, 1, 3)
-
- def unshape(x):
- """group heads"""
- return x.transpose(0, 2, 1, 3).reshape(bs, -1, self.n_heads * dim_per_head)
-
- q = shape(self.q_lin(query)) # (bs, n_heads, q_len, dim_per_head)
- k = shape(self.k_lin(key)) # (bs, n_heads, k_len, dim_per_head)
- v = shape(self.v_lin(value)) # (bs, n_heads, k_len, dim_per_head)
-
- q = q / math.sqrt(dim_per_head) # (bs, n_heads, q_len, dim_per_head)
- scores = jnp.matmul(q, k.transpose(0, 1, 3, 2)) # (bs, n_heads, q_len, k_len)
- mask = jnp.reshape(mask, mask_reshp)
-
- mask = mask.astype(scores.dtype)
- scores = scores - 1e30 * (1.0 - mask)
-
- weights = nn.softmax(scores, axis=-1) # (bs, n_heads, q_len, k_len)
- weights = self.dropout(weights, deterministic=deterministic)
-
- context = jnp.matmul(weights, v) # (bs, n_heads, q_len, dim_per_head)
- context = unshape(context) # (bs, q_len, dim)
- context = self.out_lin(context) # (bs, q_len, dim)
-
- if output_attentions:
- return (context, weights)
- else:
- return (context,)
-
-
-class FlaxFFN(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.dropout = nn.Dropout(rate=self.config.dropout)
- self.chunk_size_feed_forward = self.config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.lin1 = nn.Dense(
- self.config.hidden_dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.lin2 = nn.Dense(
- self.config.dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
-
- self.activation = ACT2FN[self.config.activation]
-
- def __call__(self, hidden_states, deterministic: bool = True):
- hidden_states = self.lin1(hidden_states)
- hidden_states = self.activation(hidden_states)
- hidden_states = self.lin2(hidden_states)
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- return hidden_states
-
-
-class FlaxTransformerBlock(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- assert (
- self.config.dim % self.config.n_heads == 0
- ), f"Hidden size {self.config.dim} not dividable by number of heads {self.config.n_heads}"
-
- self.attention = FlaxMultiHeadSelfAttention(self.config, dtype=self.dtype)
- self.sa_layer_norm = nn.LayerNorm(epsilon=1e-12, dtype=self.dtype)
-
- self.ffn = FlaxFFN(self.config, dtype=self.dtype)
- self.output_layer_norm = nn.LayerNorm(epsilon=1e-12, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attn_mask,
- output_attentions: bool = False,
- deterministic: bool = True,
- ):
- # Self-Attention
- sa_output = self.attention(
- query=hidden_states,
- key=hidden_states,
- value=hidden_states,
- mask=attn_mask,
- output_attentions=output_attentions,
- deterministic=deterministic,
- )
- if output_attentions:
- sa_output, sa_weights = sa_output
- else:
- assert type(sa_output) == tuple
- sa_output = sa_output[0]
- sa_output = self.sa_layer_norm(sa_output + hidden_states)
-
- # Feed Forward Network
- ffn_output = self.ffn(sa_output, deterministic=deterministic)
- ffn_output = self.output_layer_norm(ffn_output + sa_output)
- output = (ffn_output,)
- if output_attentions:
- output = (sa_weights,) + output
- return output
-
-
-class FlaxTransformer(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.layers = [
- FlaxTransformerBlock(self.config, name=str(i), dtype=self.dtype) for i in range(self.config.n_layers)
- ]
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- deterministic: bool = True,
- return_dict: bool = False,
- ):
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- for layer_module in self.layers:
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_outputs = layer_module(
- hidden_states=hidden_states,
- attn_mask=attention_mask,
- output_attentions=output_attentions,
- deterministic=deterministic,
- )
- hidden_states = layer_outputs[-1]
-
- if output_attentions:
- assert len(layer_outputs) == 2
- attentions = layer_outputs[0]
- all_attentions = all_attentions + (attentions,)
- else:
- assert len(layer_outputs) == 1
-
- # Add last layer
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_attentions, all_hidden_states] if v is not None)
- return FlaxBaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
-
-class FlaxTransformerEncoder(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.layer = FlaxTransformer(self.config, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- deterministic: bool = True,
- return_dict: bool = False,
- ):
- return self.layer(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- deterministic=deterministic,
- return_dict=return_dict,
- )
-
-
-class FlaxDistilBertLMDecoder(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
- bias_init: Callable[..., np.ndarray] = jax.nn.initializers.zeros
-
- def setup(self):
- self.bias = self.param("bias", self.bias_init, (self.config.vocab_size,))
-
- def __call__(self, inputs, kernel):
- inputs = jnp.asarray(inputs, self.dtype)
- kernel = jnp.asarray(kernel, self.dtype)
- y = lax.dot_general(inputs, kernel, (((inputs.ndim - 1,), (0,)), ((), ())))
- bias = jnp.asarray(self.bias, self.dtype)
- y = y + bias
- return y
-
-
-class FlaxDistilBertPreTrainedModel(FlaxPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DistilBertConfig
- base_model_prefix = "distilbert"
- module_class: nn.Module = None
-
- def __init__(
- self,
- config: DistilBertConfig,
- input_shape: Tuple = (1, 1),
- seed: int = 0,
- dtype: jnp.dtype = jnp.float32,
- _do_init: bool = True,
- **kwargs,
- ):
- module = self.module_class(config=config, dtype=dtype, **kwargs)
- super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
-
- def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
- # init input tensors
- input_ids = jnp.zeros(input_shape, dtype="i4")
- attention_mask = jnp.ones_like(input_ids)
-
- params_rng, dropout_rng = jax.random.split(rng)
- rngs = {"params": params_rng, "dropout": dropout_rng}
-
- random_params = self.module.init(rngs, input_ids, attention_mask, return_dict=False)["params"]
-
- if params is not None:
- random_params = flatten_dict(unfreeze(random_params))
- params = flatten_dict(unfreeze(params))
- for missing_key in self._missing_keys:
- params[missing_key] = random_params[missing_key]
- self._missing_keys = set()
- return freeze(unflatten_dict(params))
- else:
- return random_params
-
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- head_mask=None,
- params: dict = None,
- dropout_rng: jax.random.PRNGKey = None,
- train: bool = False,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ):
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- if attention_mask is None:
- attention_mask = jnp.ones_like(input_ids)
-
- # Handle any PRNG if needed
- rngs = {}
- if dropout_rng is not None:
- rngs["dropout"] = dropout_rng
-
- return self.module.apply(
- {"params": params or self.params},
- jnp.array(input_ids, dtype="i4"),
- jnp.array(attention_mask, dtype="i4"),
- not train,
- output_attentions,
- output_hidden_states,
- return_dict,
- rngs=rngs,
- )
-
-
-class FlaxDistilBertModule(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.embeddings = FlaxEmbeddings(self.config, dtype=self.dtype)
- self.transformer = FlaxTransformerEncoder(self.config, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- input_embeds = self.embeddings(input_ids, deterministic=deterministic)
- return self.transformer(
- hidden_states=input_embeds,
- attention_mask=attention_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
-
-@add_start_docstrings(
- "The bare DistilBert Model transformer outputting raw hidden-states without any specific head on top.",
- FLAX_DISTILBERT_START_DOCSTRING,
-)
-class FlaxDistilBertModel(FlaxDistilBertPreTrainedModel):
- module_class = FlaxDistilBertModule
-
-
-append_call_sample_docstring(FlaxDistilBertModel, _CHECKPOINT_FOR_DOC, None, _CONFIG_FOR_DOC)
-
-
-class FlaxDistilBertForMaskedLMModule(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.distilbert = FlaxDistilBertModule(self.config, dtype=self.dtype)
- self.vocab_transform = nn.Dense(
- self.config.dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.vocab_layer_norm = nn.LayerNorm(epsilon=1e-12, dtype=self.dtype)
- if self.config.tie_word_embeddings:
- self.vocab_projector = FlaxDistilBertLMDecoder(
- self.config,
- dtype=self.dtype,
- )
- else:
- self.vocab_projector = nn.Dense(
- self.config.vocab_size,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- dlbrt_output = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- deterministic=deterministic,
- return_dict=return_dict,
- )
- hidden_states = dlbrt_output[0]
- prediction_logits = self.vocab_transform(hidden_states)
- prediction_logits = ACT2FN[self.config.activation](prediction_logits)
- prediction_logits = self.vocab_layer_norm(prediction_logits)
-
- if self.config.tie_word_embeddings:
- shared_embedding = self.distilbert.variables["params"]["embeddings"]["word_embeddings"]["embedding"]
- prediction_logits = self.vocab_projector(prediction_logits, shared_embedding.T)
- else:
- prediction_logits = self.vocab_projector(prediction_logits)
-
- if not return_dict:
- output = (prediction_logits,) + dlbrt_output[1:]
- return output
-
- return FlaxMaskedLMOutput(
- logits=prediction_logits,
- hidden_states=dlbrt_output.hidden_states,
- attentions=dlbrt_output.attentions,
- )
-
-
-@add_start_docstrings("""DistilBert Model with a `language modeling` head on top.""", FLAX_DISTILBERT_START_DOCSTRING)
-class FlaxDistilBertForMaskedLM(FlaxDistilBertPreTrainedModel):
- module_class = FlaxDistilBertForMaskedLMModule
-
-
-append_call_sample_docstring(FlaxDistilBertForMaskedLM, _CHECKPOINT_FOR_DOC, FlaxMaskedLMOutput, _CONFIG_FOR_DOC)
-
-
-class FlaxDistilBertForSequenceClassificationModule(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.distilbert = FlaxDistilBertModule(config=self.config, dtype=self.dtype)
- self.pre_classifier = nn.Dense(
- self.config.dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.dropout = nn.Dropout(rate=self.config.seq_classif_dropout)
- self.classifier = nn.Dense(
- self.config.num_labels,
- dtype=self.dtype,
- )
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- # Model
- distilbert_output = self.distilbert(
- input_ids,
- attention_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_state = distilbert_output[0] # (bs, seq_len, dim)
- pooled_output = hidden_state[:, 0] # (bs, dim)
- pooled_output = self.pre_classifier(pooled_output) # (bs, dim)
- pooled_output = ACT2FN["relu"](pooled_output)
- pooled_output = self.dropout(pooled_output, deterministic=deterministic)
- logits = self.classifier(pooled_output) # (bs, dim)
-
- if not return_dict:
- return (logits,) + distilbert_output[1:]
-
- return FlaxSequenceClassifierOutput(
- logits=logits,
- hidden_states=distilbert_output.hidden_states,
- attentions=distilbert_output.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DistilBert Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- FLAX_DISTILBERT_START_DOCSTRING,
-)
-class FlaxDistilBertForSequenceClassification(FlaxDistilBertPreTrainedModel):
- module_class = FlaxDistilBertForSequenceClassificationModule
-
-
-append_call_sample_docstring(
- FlaxDistilBertForSequenceClassification,
- _CHECKPOINT_FOR_DOC,
- FlaxSequenceClassifierOutput,
- _CONFIG_FOR_DOC,
-)
-
-
-class FlaxDistilBertForMultipleChoiceModule(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.distilbert = FlaxDistilBertModule(config=self.config, dtype=self.dtype)
- self.pre_classifier = nn.Dense(
- self.config.dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.dropout = nn.Dropout(rate=self.config.seq_classif_dropout)
- self.classifier = nn.Dense(
- 1,
- dtype=self.dtype,
- )
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1]
- input_ids = input_ids.reshape(-1, input_ids.shape[-1]) if input_ids is not None else None
- attention_mask = attention_mask.reshape(-1, attention_mask.shape[-1]) if attention_mask is not None else None
-
- # Model
- outputs = self.distilbert(
- input_ids,
- attention_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_state = outputs[0]
- pooled_output = hidden_state[:, 0]
- pooled_output = self.pre_classifier(pooled_output)
- pooled_output = ACT2FN["relu"](pooled_output)
- pooled_output = self.dropout(pooled_output, deterministic=deterministic)
- logits = self.classifier(pooled_output)
-
- reshaped_logits = logits.reshape(-1, num_choices)
-
- if not return_dict:
- return (reshaped_logits,) + outputs[2:]
-
- return FlaxMultipleChoiceModelOutput(
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and
- a softmax) e.g. for RocStories/SWAG tasks.
- """,
- FLAX_DISTILBERT_START_DOCSTRING,
-)
-class FlaxDistilBertForMultipleChoice(FlaxDistilBertPreTrainedModel):
- module_class = FlaxDistilBertForMultipleChoiceModule
-
-
-overwrite_call_docstring(
- FlaxDistilBertForMultipleChoice, DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
-)
-append_call_sample_docstring(
- FlaxDistilBertForMultipleChoice,
- _CHECKPOINT_FOR_DOC,
- FlaxMultipleChoiceModelOutput,
- _CONFIG_FOR_DOC,
-)
-
-
-class FlaxDistilBertForTokenClassificationModule(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.distilbert = FlaxDistilBertModule(config=self.config, dtype=self.dtype)
- self.dropout = nn.Dropout(rate=self.config.dropout)
- self.classifier = nn.Dense(self.config.num_labels, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- # Model
- outputs = self.distilbert(
- input_ids,
- attention_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- logits = self.classifier(hidden_states)
-
- if not return_dict:
- return (logits,) + outputs[1:]
-
- return FlaxTokenClassifierOutput(
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g.
- for Named-Entity-Recognition (NER) tasks.
- """,
- FLAX_DISTILBERT_START_DOCSTRING,
-)
-class FlaxDistilBertForTokenClassification(FlaxDistilBertPreTrainedModel):
- module_class = FlaxDistilBertForTokenClassificationModule
-
-
-append_call_sample_docstring(
- FlaxDistilBertForTokenClassification,
- _CHECKPOINT_FOR_DOC,
- FlaxTokenClassifierOutput,
- _CONFIG_FOR_DOC,
-)
-
-
-class FlaxDistilBertForQuestionAnsweringModule(nn.Module):
- config: DistilBertConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.distilbert = FlaxDistilBertModule(config=self.config, dtype=self.dtype)
- self.qa_outputs = nn.Dense(self.config.num_labels, dtype=self.dtype)
- assert self.config.num_labels == 2
- self.dropout = nn.Dropout(rate=self.config.qa_dropout)
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # Model
- distilbert_output = self.distilbert(
- input_ids,
- attention_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = distilbert_output[0]
-
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- logits = self.qa_outputs(hidden_states)
- start_logits, end_logits = logits.split(self.config.num_labels, axis=-1)
- start_logits = start_logits.squeeze(-1)
- end_logits = end_logits.squeeze(-1)
-
- if not return_dict:
- return (start_logits, end_logits) + distilbert_output[1:]
-
- return FlaxQuestionAnsweringModelOutput(
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=distilbert_output.hidden_states,
- attentions=distilbert_output.attentions,
- )
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a
- linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- FLAX_DISTILBERT_START_DOCSTRING,
-)
-class FlaxDistilBertForQuestionAnswering(FlaxDistilBertPreTrainedModel):
- module_class = FlaxDistilBertForQuestionAnsweringModule
-
-
-append_call_sample_docstring(
- FlaxDistilBertForQuestionAnswering,
- _CHECKPOINT_FOR_DOC,
- FlaxQuestionAnsweringModelOutput,
- _CONFIG_FOR_DOC,
-)
diff --git a/transformers/models/distilbert/modeling_tf_distilbert.py b/transformers/models/distilbert/modeling_tf_distilbert.py
deleted file mode 100644
index c41deac3f2e57e53572e99fefaa4b7e26eb4309f..0000000000000000000000000000000000000000
--- a/transformers/models/distilbert/modeling_tf_distilbert.py
+++ /dev/null
@@ -1,1139 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
- TF 2.0 DistilBERT model
-"""
-
-
-from __future__ import annotations
-
-import warnings
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFMaskedLMOutput,
- TFMultipleChoiceModelOutput,
- TFQuestionAnsweringModelOutput,
- TFSequenceClassifierOutput,
- TFTokenClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFMaskedLanguageModelingLoss,
- TFModelInputType,
- TFMultipleChoiceLoss,
- TFPreTrainedModel,
- TFQuestionAnsweringLoss,
- TFSequenceClassificationLoss,
- TFTokenClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_distilbert import DistilBertConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "distilbert-base-uncased"
-_CONFIG_FOR_DOC = "DistilBertConfig"
-
-
-from ..deprecated._archive_maps import TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class TFEmbeddings(keras.layers.Layer):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- self.config = config
- self.dim = config.dim
- self.initializer_range = config.initializer_range
- self.max_position_embeddings = config.max_position_embeddings
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=1e-12, name="LayerNorm")
- self.dropout = keras.layers.Dropout(rate=config.dropout)
-
- def build(self, input_shape=None):
- with tf.name_scope("word_embeddings"):
- self.weight = self.add_weight(
- name="weight",
- shape=[self.config.vocab_size, self.dim],
- initializer=get_initializer(initializer_range=self.initializer_range),
- )
-
- with tf.name_scope("position_embeddings"):
- self.position_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.max_position_embeddings, self.dim],
- initializer=get_initializer(initializer_range=self.initializer_range),
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.dim])
-
- def call(self, input_ids=None, position_ids=None, inputs_embeds=None, training=False):
- """
- Applies embedding based on inputs tensor.
-
- Returns:
- final_embeddings (`tf.Tensor`): output embedding tensor.
- """
- assert not (input_ids is None and inputs_embeds is None)
-
- if input_ids is not None:
- check_embeddings_within_bounds(input_ids, self.config.vocab_size)
- inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
-
- input_shape = shape_list(inputs_embeds)[:-1]
-
- if position_ids is None:
- position_ids = tf.expand_dims(tf.range(start=0, limit=input_shape[-1]), axis=0)
-
- position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
- final_embeddings = inputs_embeds + position_embeds
- final_embeddings = self.LayerNorm(inputs=final_embeddings)
- final_embeddings = self.dropout(inputs=final_embeddings, training=training)
-
- return final_embeddings
-
-
-class TFMultiHeadSelfAttention(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.n_heads = config.n_heads
- self.dim = config.dim
- self.dropout = keras.layers.Dropout(config.attention_dropout)
- self.output_attentions = config.output_attentions
-
- assert self.dim % self.n_heads == 0, f"Hidden size {self.dim} not dividable by number of heads {self.n_heads}"
-
- self.q_lin = keras.layers.Dense(
- config.dim, kernel_initializer=get_initializer(config.initializer_range), name="q_lin"
- )
- self.k_lin = keras.layers.Dense(
- config.dim, kernel_initializer=get_initializer(config.initializer_range), name="k_lin"
- )
- self.v_lin = keras.layers.Dense(
- config.dim, kernel_initializer=get_initializer(config.initializer_range), name="v_lin"
- )
- self.out_lin = keras.layers.Dense(
- config.dim, kernel_initializer=get_initializer(config.initializer_range), name="out_lin"
- )
-
- self.pruned_heads = set()
- self.config = config
-
- def prune_heads(self, heads):
- raise NotImplementedError
-
- def call(self, query, key, value, mask, head_mask, output_attentions, training=False):
- """
- Parameters:
- query: tf.Tensor(bs, seq_length, dim)
- key: tf.Tensor(bs, seq_length, dim)
- value: tf.Tensor(bs, seq_length, dim)
- mask: tf.Tensor(bs, seq_length)
-
- Returns:
- weights: tf.Tensor(bs, n_heads, seq_length, seq_length) Attention weights context: tf.Tensor(bs,
- seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
- """
- bs, q_length, dim = shape_list(query)
- k_length = shape_list(key)[1]
- # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
- # assert key.size() == value.size()
- dim_per_head = int(self.dim / self.n_heads)
- dim_per_head = tf.cast(dim_per_head, dtype=tf.int32)
- mask_reshape = [bs, 1, 1, k_length]
-
- def shape(x):
- """separate heads"""
- return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))
-
- def unshape(x):
- """group heads"""
- return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))
-
- q = shape(self.q_lin(query)) # (bs, n_heads, q_length, dim_per_head)
- k = shape(self.k_lin(key)) # (bs, n_heads, k_length, dim_per_head)
- v = shape(self.v_lin(value)) # (bs, n_heads, k_length, dim_per_head)
- q = tf.cast(q, dtype=tf.float32)
- q = tf.multiply(q, tf.math.rsqrt(tf.cast(dim_per_head, dtype=tf.float32)))
- k = tf.cast(k, dtype=q.dtype)
- scores = tf.matmul(q, k, transpose_b=True) # (bs, n_heads, q_length, k_length)
- mask = tf.reshape(mask, mask_reshape) # (bs, n_heads, qlen, klen)
- # scores.masked_fill_(mask, -float('inf')) # (bs, n_heads, q_length, k_length)
-
- mask = tf.cast(mask, dtype=scores.dtype)
- scores = scores - 1e30 * (1.0 - mask)
- weights = stable_softmax(scores, axis=-1) # (bs, n_heads, qlen, klen)
- weights = self.dropout(weights, training=training) # (bs, n_heads, qlen, klen)
-
- # Mask heads if we want to
- if head_mask is not None:
- weights = weights * head_mask
-
- context = tf.matmul(weights, v) # (bs, n_heads, qlen, dim_per_head)
- context = unshape(context) # (bs, q_length, dim)
- context = self.out_lin(context) # (bs, q_length, dim)
-
- if output_attentions:
- return (context, weights)
- else:
- return (context,)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "q_lin", None) is not None:
- with tf.name_scope(self.q_lin.name):
- self.q_lin.build([None, None, self.config.dim])
- if getattr(self, "k_lin", None) is not None:
- with tf.name_scope(self.k_lin.name):
- self.k_lin.build([None, None, self.config.dim])
- if getattr(self, "v_lin", None) is not None:
- with tf.name_scope(self.v_lin.name):
- self.v_lin.build([None, None, self.config.dim])
- if getattr(self, "out_lin", None) is not None:
- with tf.name_scope(self.out_lin.name):
- self.out_lin.build([None, None, self.config.dim])
-
-
-class TFFFN(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- self.dropout = keras.layers.Dropout(config.dropout)
- self.lin1 = keras.layers.Dense(
- config.hidden_dim, kernel_initializer=get_initializer(config.initializer_range), name="lin1"
- )
- self.lin2 = keras.layers.Dense(
- config.dim, kernel_initializer=get_initializer(config.initializer_range), name="lin2"
- )
- self.activation = get_tf_activation(config.activation)
- self.config = config
-
- def call(self, input, training=False):
- x = self.lin1(input)
- x = self.activation(x)
- x = self.lin2(x)
- x = self.dropout(x, training=training)
- return x
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "lin1", None) is not None:
- with tf.name_scope(self.lin1.name):
- self.lin1.build([None, None, self.config.dim])
- if getattr(self, "lin2", None) is not None:
- with tf.name_scope(self.lin2.name):
- self.lin2.build([None, None, self.config.hidden_dim])
-
-
-class TFTransformerBlock(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.n_heads = config.n_heads
- self.dim = config.dim
- self.hidden_dim = config.hidden_dim
- self.dropout = keras.layers.Dropout(config.dropout)
- self.activation = config.activation
- self.output_attentions = config.output_attentions
-
- assert (
- config.dim % config.n_heads == 0
- ), f"Hidden size {config.dim} not dividable by number of heads {config.n_heads}"
-
- self.attention = TFMultiHeadSelfAttention(config, name="attention")
- self.sa_layer_norm = keras.layers.LayerNormalization(epsilon=1e-12, name="sa_layer_norm")
-
- self.ffn = TFFFN(config, name="ffn")
- self.output_layer_norm = keras.layers.LayerNormalization(epsilon=1e-12, name="output_layer_norm")
- self.config = config
-
- def call(self, x, attn_mask, head_mask, output_attentions, training=False): # removed: src_enc=None, src_len=None
- """
- Parameters:
- x: tf.Tensor(bs, seq_length, dim)
- attn_mask: tf.Tensor(bs, seq_length)
-
- Outputs: sa_weights: tf.Tensor(bs, n_heads, seq_length, seq_length) The attention weights ffn_output:
- tf.Tensor(bs, seq_length, dim) The output of the transformer block contextualization.
- """
- # Self-Attention
- sa_output = self.attention(x, x, x, attn_mask, head_mask, output_attentions, training=training)
- if output_attentions:
- sa_output, sa_weights = sa_output # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)
- else: # To handle these `output_attentions` or `output_hidden_states` cases returning tuples
- # assert type(sa_output) == tuple
- sa_output = sa_output[0]
- sa_output = self.sa_layer_norm(sa_output + x) # (bs, seq_length, dim)
-
- # Feed Forward Network
- ffn_output = self.ffn(sa_output, training=training) # (bs, seq_length, dim)
- ffn_output = self.output_layer_norm(ffn_output + sa_output) # (bs, seq_length, dim)
-
- output = (ffn_output,)
- if output_attentions:
- output = (sa_weights,) + output
- return output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "sa_layer_norm", None) is not None:
- with tf.name_scope(self.sa_layer_norm.name):
- self.sa_layer_norm.build([None, None, self.config.dim])
- if getattr(self, "ffn", None) is not None:
- with tf.name_scope(self.ffn.name):
- self.ffn.build(None)
- if getattr(self, "output_layer_norm", None) is not None:
- with tf.name_scope(self.output_layer_norm.name):
- self.output_layer_norm.build([None, None, self.config.dim])
-
-
-class TFTransformer(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- self.n_layers = config.n_layers
- self.output_hidden_states = config.output_hidden_states
- self.output_attentions = config.output_attentions
-
- self.layer = [TFTransformerBlock(config, name=f"layer_._{i}") for i in range(config.n_layers)]
-
- def call(self, x, attn_mask, head_mask, output_attentions, output_hidden_states, return_dict, training=False):
- # docstyle-ignore
- """
- Parameters:
- x: tf.Tensor(bs, seq_length, dim) Input sequence embedded.
- attn_mask: tf.Tensor(bs, seq_length) Attention mask on the sequence.
-
- Returns:
- hidden_state: tf.Tensor(bs, seq_length, dim)
- Sequence of hidden states in the last (top) layer
- all_hidden_states: Tuple[tf.Tensor(bs, seq_length, dim)]
- Tuple of length n_layers with the hidden states from each layer.
- Optional: only if output_hidden_states=True
- all_attentions: Tuple[tf.Tensor(bs, n_heads, seq_length, seq_length)]
- Tuple of length n_layers with the attention weights from each layer
- Optional: only if output_attentions=True
- """
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- hidden_state = x
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_state,)
-
- layer_outputs = layer_module(hidden_state, attn_mask, head_mask[i], output_attentions, training=training)
- hidden_state = layer_outputs[-1]
-
- if output_attentions:
- assert len(layer_outputs) == 2
- attentions = layer_outputs[0]
- all_attentions = all_attentions + (attentions,)
- else:
- assert len(layer_outputs) == 1, f"Incorrect number of outputs {len(layer_outputs)} instead of 1"
-
- # Add last layer
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_state,)
-
- if not return_dict:
- return tuple(v for v in [hidden_state, all_hidden_states, all_attentions] if v is not None)
- return TFBaseModelOutput(
- last_hidden_state=hidden_state, hidden_states=all_hidden_states, attentions=all_attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "layer", None) is not None:
- for layer in self.layer:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-@keras_serializable
-class TFDistilBertMainLayer(keras.layers.Layer):
- config_class = DistilBertConfig
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.num_hidden_layers = config.num_hidden_layers
- self.output_attentions = config.output_attentions
- self.output_hidden_states = config.output_hidden_states
- self.return_dict = config.use_return_dict
-
- self.embeddings = TFEmbeddings(config, name="embeddings") # Embeddings
- self.transformer = TFTransformer(config, name="transformer") # Encoder
-
- def get_input_embeddings(self):
- return self.embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.weight = value
- self.embeddings.vocab_size = value.shape[0]
-
- def _prune_heads(self, heads_to_prune):
- raise NotImplementedError
-
- @unpack_inputs
- def call(
- self,
- input_ids=None,
- attention_mask=None,
- head_mask=None,
- inputs_embeds=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- training=False,
- ):
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = tf.ones(input_shape) # (bs, seq_length)
-
- attention_mask = tf.cast(attention_mask, dtype=tf.float32)
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.num_hidden_layers
-
- embedding_output = self.embeddings(input_ids, inputs_embeds=inputs_embeds) # (bs, seq_length, dim)
- tfmr_output = self.transformer(
- embedding_output,
- attention_mask,
- head_mask,
- output_attentions,
- output_hidden_states,
- return_dict,
- training=training,
- )
-
- return tfmr_output # last-layer hidden-state, (all hidden_states), (all attentions)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
-
-
-# INTERFACE FOR ENCODER AND TASK SPECIFIC MODEL #
-class TFDistilBertPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DistilBertConfig
- base_model_prefix = "distilbert"
-
-
-DISTILBERT_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`DistilBertConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DISTILBERT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- head_mask (`Numpy array` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare DistilBERT encoder/transformer outputting raw hidden-states without any specific head on top.",
- DISTILBERT_START_DOCSTRING,
-)
-class TFDistilBertModel(TFDistilBertPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.distilbert = TFDistilBertMainLayer(config, name="distilbert") # Embeddings
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
- outputs = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "distilbert", None) is not None:
- with tf.name_scope(self.distilbert.name):
- self.distilbert.build(None)
-
-
-class TFDistilBertLMHead(keras.layers.Layer):
- def __init__(self, config, input_embeddings, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.dim = config.dim
-
- # The output weights are the same as the input embeddings, but there is
- # an output-only bias for each token.
- self.input_embeddings = input_embeddings
-
- def build(self, input_shape):
- self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")
-
- super().build(input_shape)
-
- def get_output_embeddings(self):
- return self.input_embeddings
-
- def set_output_embeddings(self, value):
- self.input_embeddings.weight = value
- self.input_embeddings.vocab_size = shape_list(value)[0]
-
- def get_bias(self):
- return {"bias": self.bias}
-
- def set_bias(self, value):
- self.bias = value["bias"]
- self.config.vocab_size = shape_list(value["bias"])[0]
-
- def call(self, hidden_states):
- seq_length = shape_list(tensor=hidden_states)[1]
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.dim])
- hidden_states = tf.matmul(a=hidden_states, b=self.input_embeddings.weight, transpose_b=True)
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])
- hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)
-
- return hidden_states
-
-
-@add_start_docstrings(
- """DistilBert Model with a `masked language modeling` head on top.""",
- DISTILBERT_START_DOCSTRING,
-)
-class TFDistilBertForMaskedLM(TFDistilBertPreTrainedModel, TFMaskedLanguageModelingLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.config = config
-
- self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
- self.vocab_transform = keras.layers.Dense(
- config.dim, kernel_initializer=get_initializer(config.initializer_range), name="vocab_transform"
- )
- self.act = get_tf_activation(config.activation)
- self.vocab_layer_norm = keras.layers.LayerNormalization(epsilon=1e-12, name="vocab_layer_norm")
- self.vocab_projector = TFDistilBertLMHead(config, self.distilbert.embeddings, name="vocab_projector")
-
- def get_lm_head(self):
- return self.vocab_projector
-
- def get_prefix_bias_name(self):
- warnings.warn("The method get_prefix_bias_name is deprecated. Please use `get_bias` instead.", FutureWarning)
- return self.name + "/" + self.vocab_projector.name
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFMaskedLMOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- distilbert_output = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_states = distilbert_output[0] # (bs, seq_length, dim)
- prediction_logits = self.vocab_transform(hidden_states) # (bs, seq_length, dim)
- prediction_logits = self.act(prediction_logits) # (bs, seq_length, dim)
- prediction_logits = self.vocab_layer_norm(prediction_logits) # (bs, seq_length, dim)
- prediction_logits = self.vocab_projector(prediction_logits)
-
- loss = None if labels is None else self.hf_compute_loss(labels, prediction_logits)
-
- if not return_dict:
- output = (prediction_logits,) + distilbert_output[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFMaskedLMOutput(
- loss=loss,
- logits=prediction_logits,
- hidden_states=distilbert_output.hidden_states,
- attentions=distilbert_output.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "distilbert", None) is not None:
- with tf.name_scope(self.distilbert.name):
- self.distilbert.build(None)
- if getattr(self, "vocab_transform", None) is not None:
- with tf.name_scope(self.vocab_transform.name):
- self.vocab_transform.build([None, None, self.config.dim])
- if getattr(self, "vocab_layer_norm", None) is not None:
- with tf.name_scope(self.vocab_layer_norm.name):
- self.vocab_layer_norm.build([None, None, self.config.dim])
- if getattr(self, "vocab_projector", None) is not None:
- with tf.name_scope(self.vocab_projector.name):
- self.vocab_projector.build(None)
-
-
-@add_start_docstrings(
- """
- DistilBert Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- DISTILBERT_START_DOCSTRING,
-)
-class TFDistilBertForSequenceClassification(TFDistilBertPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
-
- self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
- self.pre_classifier = keras.layers.Dense(
- config.dim,
- kernel_initializer=get_initializer(config.initializer_range),
- activation="relu",
- name="pre_classifier",
- )
- self.classifier = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.dropout = keras.layers.Dropout(config.seq_classif_dropout)
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- distilbert_output = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_state = distilbert_output[0] # (bs, seq_len, dim)
- pooled_output = hidden_state[:, 0] # (bs, dim)
- pooled_output = self.pre_classifier(pooled_output) # (bs, dim)
- pooled_output = self.dropout(pooled_output, training=training) # (bs, dim)
- logits = self.classifier(pooled_output) # (bs, dim)
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + distilbert_output[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=distilbert_output.hidden_states,
- attentions=distilbert_output.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "distilbert", None) is not None:
- with tf.name_scope(self.distilbert.name):
- self.distilbert.build(None)
- if getattr(self, "pre_classifier", None) is not None:
- with tf.name_scope(self.pre_classifier.name):
- self.pre_classifier.build([None, None, self.config.dim])
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.dim])
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g.
- for Named-Entity-Recognition (NER) tasks.
- """,
- DISTILBERT_START_DOCSTRING,
-)
-class TFDistilBertForTokenClassification(TFDistilBertPreTrainedModel, TFTokenClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
-
- self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
- self.dropout = keras.layers.Dropout(config.dropout)
- self.classifier = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFTokenClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- outputs = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- sequence_output = self.dropout(sequence_output, training=training)
- logits = self.classifier(sequence_output)
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFTokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "distilbert", None) is not None:
- with tf.name_scope(self.distilbert.name):
- self.distilbert.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and
- a softmax) e.g. for RocStories/SWAG tasks.
- """,
- DISTILBERT_START_DOCSTRING,
-)
-class TFDistilBertForMultipleChoice(TFDistilBertPreTrainedModel, TFMultipleChoiceLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
- self.dropout = keras.layers.Dropout(config.seq_classif_dropout)
- self.pre_classifier = keras.layers.Dense(
- config.dim,
- kernel_initializer=get_initializer(config.initializer_range),
- activation="relu",
- name="pre_classifier",
- )
- self.classifier = keras.layers.Dense(
- 1, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(
- DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
- )
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFMultipleChoiceModelOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]`
- where `num_choices` is the size of the second dimension of the input tensors. (See `input_ids` above)
- """
- if input_ids is not None:
- num_choices = shape_list(input_ids)[1]
- seq_length = shape_list(input_ids)[2]
- else:
- num_choices = shape_list(inputs_embeds)[1]
- seq_length = shape_list(inputs_embeds)[2]
-
- flat_input_ids = tf.reshape(input_ids, (-1, seq_length)) if input_ids is not None else None
- flat_attention_mask = tf.reshape(attention_mask, (-1, seq_length)) if attention_mask is not None else None
- flat_inputs_embeds = (
- tf.reshape(inputs_embeds, (-1, seq_length, shape_list(inputs_embeds)[3]))
- if inputs_embeds is not None
- else None
- )
- distilbert_output = self.distilbert(
- flat_input_ids,
- flat_attention_mask,
- head_mask,
- flat_inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_state = distilbert_output[0] # (bs, seq_len, dim)
- pooled_output = hidden_state[:, 0] # (bs, dim)
- pooled_output = self.pre_classifier(pooled_output) # (bs, dim)
- pooled_output = self.dropout(pooled_output, training=training) # (bs, dim)
- logits = self.classifier(pooled_output)
- reshaped_logits = tf.reshape(logits, (-1, num_choices))
-
- loss = None if labels is None else self.hf_compute_loss(labels, reshaped_logits)
-
- if not return_dict:
- output = (reshaped_logits,) + distilbert_output[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFMultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=distilbert_output.hidden_states,
- attentions=distilbert_output.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "distilbert", None) is not None:
- with tf.name_scope(self.distilbert.name):
- self.distilbert.build(None)
- if getattr(self, "pre_classifier", None) is not None:
- with tf.name_scope(self.pre_classifier.name):
- self.pre_classifier.build([None, None, self.config.dim])
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.dim])
-
-
-@add_start_docstrings(
- """
- DistilBert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a
- linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- DISTILBERT_START_DOCSTRING,
-)
-class TFDistilBertForQuestionAnswering(TFDistilBertPreTrainedModel, TFQuestionAnsweringLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
- self.qa_outputs = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
- )
- assert config.num_labels == 2, f"Incorrect number of labels {config.num_labels} instead of 2"
- self.dropout = keras.layers.Dropout(config.qa_dropout)
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFQuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- start_positions: np.ndarray | tf.Tensor | None = None,
- end_positions: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFQuestionAnsweringModelOutput, Tuple[tf.Tensor]]:
- r"""
- start_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- distilbert_output = self.distilbert(
- input_ids=input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_states = distilbert_output[0] # (bs, max_query_len, dim)
- hidden_states = self.dropout(hidden_states, training=training) # (bs, max_query_len, dim)
- logits = self.qa_outputs(hidden_states) # (bs, max_query_len, 2)
- start_logits, end_logits = tf.split(logits, 2, axis=-1)
- start_logits = tf.squeeze(start_logits, axis=-1)
- end_logits = tf.squeeze(end_logits, axis=-1)
-
- loss = None
- if start_positions is not None and end_positions is not None:
- labels = {"start_position": start_positions}
- labels["end_position"] = end_positions
- loss = self.hf_compute_loss(labels, (start_logits, end_logits))
-
- if not return_dict:
- output = (start_logits, end_logits) + distilbert_output[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFQuestionAnsweringModelOutput(
- loss=loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=distilbert_output.hidden_states,
- attentions=distilbert_output.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "distilbert", None) is not None:
- with tf.name_scope(self.distilbert.name):
- self.distilbert.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.config.dim])
diff --git a/transformers/models/distilbert/tokenization_distilbert.py b/transformers/models/distilbert/tokenization_distilbert.py
deleted file mode 100644
index ff8854ba3dcf893a1a38d9491b2aa148a64057ca..0000000000000000000000000000000000000000
--- a/transformers/models/distilbert/tokenization_distilbert.py
+++ /dev/null
@@ -1,514 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for DistilBERT."""
-
-import collections
-import os
-import unicodedata
-from typing import List, Optional, Tuple
-
-from ...tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
-
-
-# Copied from transformers.models.bert.tokenization_bert.load_vocab
-def load_vocab(vocab_file):
- """Loads a vocabulary file into a dictionary."""
- vocab = collections.OrderedDict()
- with open(vocab_file, "r", encoding="utf-8") as reader:
- tokens = reader.readlines()
- for index, token in enumerate(tokens):
- token = token.rstrip("\n")
- vocab[token] = index
- return vocab
-
-
-# Copied from transformers.models.bert.tokenization_bert.whitespace_tokenize
-def whitespace_tokenize(text):
- """Runs basic whitespace cleaning and splitting on a piece of text."""
- text = text.strip()
- if not text:
- return []
- tokens = text.split()
- return tokens
-
-
-class DistilBertTokenizer(PreTrainedTokenizer):
- r"""
- Construct a DistilBERT tokenizer. Based on WordPiece.
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- do_basic_tokenize (`bool`, *optional*, defaults to `True`):
- Whether or not to do basic tokenization before WordPiece.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- do_lower_case=True,
- do_basic_tokenize=True,
- never_split=None,
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- if not os.path.isfile(vocab_file):
- raise ValueError(
- f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
- " model use `tokenizer = DistilBertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- self.vocab = load_vocab(vocab_file)
- self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
- self.do_basic_tokenize = do_basic_tokenize
- if do_basic_tokenize:
- self.basic_tokenizer = BasicTokenizer(
- do_lower_case=do_lower_case,
- never_split=never_split,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- )
- self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token))
-
- super().__init__(
- do_lower_case=do_lower_case,
- do_basic_tokenize=do_basic_tokenize,
- never_split=never_split,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- @property
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.do_lower_case
- def do_lower_case(self):
- return self.basic_tokenizer.do_lower_case
-
- @property
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.vocab_size
- def vocab_size(self):
- return len(self.vocab)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_vocab
- def get_vocab(self):
- return dict(self.vocab, **self.added_tokens_encoder)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._tokenize
- def _tokenize(self, text, split_special_tokens=False):
- split_tokens = []
- if self.do_basic_tokenize:
- for token in self.basic_tokenizer.tokenize(
- text, never_split=self.all_special_tokens if not split_special_tokens else None
- ):
- # If the token is part of the never_split set
- if token in self.basic_tokenizer.never_split:
- split_tokens.append(token)
- else:
- split_tokens += self.wordpiece_tokenizer.tokenize(token)
- else:
- split_tokens = self.wordpiece_tokenizer.tokenize(text)
- return split_tokens
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.vocab.get(token, self.vocab.get(self.unk_token))
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.ids_to_tokens.get(index, self.unk_token)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.convert_tokens_to_string
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- out_string = " ".join(tokens).replace(" ##", "").strip()
- return out_string
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A BERT sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_special_tokens_mask
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- index = 0
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- else:
- vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(token + "\n")
- index += 1
- return (vocab_file,)
-
-
-# Copied from transformers.models.bert.tokenization_bert.BasicTokenizer
-class BasicTokenizer(object):
- """
- Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
-
- Args:
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- do_split_on_punc (`bool`, *optional*, defaults to `True`):
- In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
- the full context of the words, such as contractions.
- """
-
- def __init__(
- self,
- do_lower_case=True,
- never_split=None,
- tokenize_chinese_chars=True,
- strip_accents=None,
- do_split_on_punc=True,
- ):
- if never_split is None:
- never_split = []
- self.do_lower_case = do_lower_case
- self.never_split = set(never_split)
- self.tokenize_chinese_chars = tokenize_chinese_chars
- self.strip_accents = strip_accents
- self.do_split_on_punc = do_split_on_punc
-
- def tokenize(self, text, never_split=None):
- """
- Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
-
- Args:
- never_split (`List[str]`, *optional*)
- Kept for backward compatibility purposes. Now implemented directly at the base class level (see
- [`PreTrainedTokenizer.tokenize`]) List of token not to split.
- """
- # union() returns a new set by concatenating the two sets.
- never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
- text = self._clean_text(text)
-
- # This was added on November 1st, 2018 for the multilingual and Chinese
- # models. This is also applied to the English models now, but it doesn't
- # matter since the English models were not trained on any Chinese data
- # and generally don't have any Chinese data in them (there are Chinese
- # characters in the vocabulary because Wikipedia does have some Chinese
- # words in the English Wikipedia.).
- if self.tokenize_chinese_chars:
- text = self._tokenize_chinese_chars(text)
- # prevents treating the same character with different unicode codepoints as different characters
- unicode_normalized_text = unicodedata.normalize("NFC", text)
- orig_tokens = whitespace_tokenize(unicode_normalized_text)
- split_tokens = []
- for token in orig_tokens:
- if token not in never_split:
- if self.do_lower_case:
- token = token.lower()
- if self.strip_accents is not False:
- token = self._run_strip_accents(token)
- elif self.strip_accents:
- token = self._run_strip_accents(token)
- split_tokens.extend(self._run_split_on_punc(token, never_split))
-
- output_tokens = whitespace_tokenize(" ".join(split_tokens))
- return output_tokens
-
- def _run_strip_accents(self, text):
- """Strips accents from a piece of text."""
- text = unicodedata.normalize("NFD", text)
- output = []
- for char in text:
- cat = unicodedata.category(char)
- if cat == "Mn":
- continue
- output.append(char)
- return "".join(output)
-
- def _run_split_on_punc(self, text, never_split=None):
- """Splits punctuation on a piece of text."""
- if not self.do_split_on_punc or (never_split is not None and text in never_split):
- return [text]
- chars = list(text)
- i = 0
- start_new_word = True
- output = []
- while i < len(chars):
- char = chars[i]
- if _is_punctuation(char):
- output.append([char])
- start_new_word = True
- else:
- if start_new_word:
- output.append([])
- start_new_word = False
- output[-1].append(char)
- i += 1
-
- return ["".join(x) for x in output]
-
- def _tokenize_chinese_chars(self, text):
- """Adds whitespace around any CJK character."""
- output = []
- for char in text:
- cp = ord(char)
- if self._is_chinese_char(cp):
- output.append(" ")
- output.append(char)
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
- def _is_chinese_char(self, cp):
- """Checks whether CP is the codepoint of a CJK character."""
- # This defines a "chinese character" as anything in the CJK Unicode block:
- # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
- #
- # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
- # despite its name. The modern Korean Hangul alphabet is a different block,
- # as is Japanese Hiragana and Katakana. Those alphabets are used to write
- # space-separated words, so they are not treated specially and handled
- # like the all of the other languages.
- if (
- (cp >= 0x4E00 and cp <= 0x9FFF)
- or (cp >= 0x3400 and cp <= 0x4DBF) #
- or (cp >= 0x20000 and cp <= 0x2A6DF) #
- or (cp >= 0x2A700 and cp <= 0x2B73F) #
- or (cp >= 0x2B740 and cp <= 0x2B81F) #
- or (cp >= 0x2B820 and cp <= 0x2CEAF) #
- or (cp >= 0xF900 and cp <= 0xFAFF)
- or (cp >= 0x2F800 and cp <= 0x2FA1F) #
- ): #
- return True
-
- return False
-
- def _clean_text(self, text):
- """Performs invalid character removal and whitespace cleanup on text."""
- output = []
- for char in text:
- cp = ord(char)
- if cp == 0 or cp == 0xFFFD or _is_control(char):
- continue
- if _is_whitespace(char):
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
-
-# Copied from transformers.models.bert.tokenization_bert.WordpieceTokenizer
-class WordpieceTokenizer(object):
- """Runs WordPiece tokenization."""
-
- def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
- self.vocab = vocab
- self.unk_token = unk_token
- self.max_input_chars_per_word = max_input_chars_per_word
-
- def tokenize(self, text):
- """
- Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
- tokenization using the given vocabulary.
-
- For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
-
- Args:
- text: A single token or whitespace separated tokens. This should have
- already been passed through *BasicTokenizer*.
-
- Returns:
- A list of wordpiece tokens.
- """
-
- output_tokens = []
- for token in whitespace_tokenize(text):
- chars = list(token)
- if len(chars) > self.max_input_chars_per_word:
- output_tokens.append(self.unk_token)
- continue
-
- is_bad = False
- start = 0
- sub_tokens = []
- while start < len(chars):
- end = len(chars)
- cur_substr = None
- while start < end:
- substr = "".join(chars[start:end])
- if start > 0:
- substr = "##" + substr
- if substr in self.vocab:
- cur_substr = substr
- break
- end -= 1
- if cur_substr is None:
- is_bad = True
- break
- sub_tokens.append(cur_substr)
- start = end
-
- if is_bad:
- output_tokens.append(self.unk_token)
- else:
- output_tokens.extend(sub_tokens)
- return output_tokens
diff --git a/transformers/models/distilbert/tokenization_distilbert_fast.py b/transformers/models/distilbert/tokenization_distilbert_fast.py
deleted file mode 100644
index f1d69a27d67c081301adb22b263928eb02f4dd84..0000000000000000000000000000000000000000
--- a/transformers/models/distilbert/tokenization_distilbert_fast.py
+++ /dev/null
@@ -1,176 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for DistilBERT."""
-
-import json
-from typing import List, Optional, Tuple
-
-from tokenizers import normalizers
-
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import logging
-from .tokenization_distilbert import DistilBertTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
-
-
-class DistilBertTokenizerFast(PreTrainedTokenizerFast):
- r"""
- Construct a "fast" DistilBERT tokenizer (backed by HuggingFace's *tokenizers* library). Based on WordPiece.
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- clean_text (`bool`, *optional*, defaults to `True`):
- Whether or not to clean the text before tokenization by removing any control characters and replacing all
- whitespaces by the classic one.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see [this
- issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- wordpieces_prefix (`str`, *optional*, defaults to `"##"`):
- The prefix for subwords.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
- slow_tokenizer_class = DistilBertTokenizer
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- do_lower_case=True,
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- super().__init__(
- vocab_file,
- tokenizer_file=tokenizer_file,
- do_lower_case=do_lower_case,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- normalizer_state = json.loads(self.backend_tokenizer.normalizer.__getstate__())
- if (
- normalizer_state.get("lowercase", do_lower_case) != do_lower_case
- or normalizer_state.get("strip_accents", strip_accents) != strip_accents
- or normalizer_state.get("handle_chinese_chars", tokenize_chinese_chars) != tokenize_chinese_chars
- ):
- normalizer_class = getattr(normalizers, normalizer_state.pop("type"))
- normalizer_state["lowercase"] = do_lower_case
- normalizer_state["strip_accents"] = strip_accents
- normalizer_state["handle_chinese_chars"] = tokenize_chinese_chars
- self.backend_tokenizer.normalizer = normalizer_class(**normalizer_state)
-
- self.do_lower_case = do_lower_case
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A BERT sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- output = [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
-
- if token_ids_1 is not None:
- output += token_ids_1 + [self.sep_token_id]
-
- return output
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
diff --git a/transformers/models/dit/__init__.py b/transformers/models/dit/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/transformers/models/dit/__pycache__/__init__.cpython-310.pyc b/transformers/models/dit/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index f6cd861c548df54637feb3ab0dd997eb49250ea0..0000000000000000000000000000000000000000
Binary files a/transformers/models/dit/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dit/__pycache__/convert_dit_unilm_to_pytorch.cpython-310.pyc b/transformers/models/dit/__pycache__/convert_dit_unilm_to_pytorch.cpython-310.pyc
deleted file mode 100644
index ea9588df4d492d40ce9803c6403b58db3b6d548d..0000000000000000000000000000000000000000
Binary files a/transformers/models/dit/__pycache__/convert_dit_unilm_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dit/convert_dit_unilm_to_pytorch.py b/transformers/models/dit/convert_dit_unilm_to_pytorch.py
deleted file mode 100644
index c754b9bbf3eac7b6c5d50aa546383334c5adbf54..0000000000000000000000000000000000000000
--- a/transformers/models/dit/convert_dit_unilm_to_pytorch.py
+++ /dev/null
@@ -1,231 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DiT checkpoints from the unilm repository."""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from transformers import BeitConfig, BeitForImageClassification, BeitForMaskedImageModeling, BeitImageProcessor
-from transformers.image_utils import PILImageResampling
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys(config, has_lm_head=False, is_semantic=False):
- prefix = "backbone." if is_semantic else ""
-
- rename_keys = []
- for i in range(config.num_hidden_layers):
- # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
- rename_keys.append((f"{prefix}blocks.{i}.norm1.weight", f"beit.encoder.layer.{i}.layernorm_before.weight"))
- rename_keys.append((f"{prefix}blocks.{i}.norm1.bias", f"beit.encoder.layer.{i}.layernorm_before.bias"))
- rename_keys.append(
- (f"{prefix}blocks.{i}.attn.proj.weight", f"beit.encoder.layer.{i}.attention.output.dense.weight")
- )
- rename_keys.append(
- (f"{prefix}blocks.{i}.attn.proj.bias", f"beit.encoder.layer.{i}.attention.output.dense.bias")
- )
- rename_keys.append((f"{prefix}blocks.{i}.norm2.weight", f"beit.encoder.layer.{i}.layernorm_after.weight"))
- rename_keys.append((f"{prefix}blocks.{i}.norm2.bias", f"beit.encoder.layer.{i}.layernorm_after.bias"))
- rename_keys.append((f"{prefix}blocks.{i}.mlp.fc1.weight", f"beit.encoder.layer.{i}.intermediate.dense.weight"))
- rename_keys.append((f"{prefix}blocks.{i}.mlp.fc1.bias", f"beit.encoder.layer.{i}.intermediate.dense.bias"))
- rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.weight", f"beit.encoder.layer.{i}.output.dense.weight"))
- rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.bias", f"beit.encoder.layer.{i}.output.dense.bias"))
-
- # projection layer + position embeddings
- rename_keys.extend(
- [
- (f"{prefix}cls_token", "beit.embeddings.cls_token"),
- (f"{prefix}patch_embed.proj.weight", "beit.embeddings.patch_embeddings.projection.weight"),
- (f"{prefix}patch_embed.proj.bias", "beit.embeddings.patch_embeddings.projection.bias"),
- (f"{prefix}pos_embed", "beit.embeddings.position_embeddings"),
- ]
- )
-
- if has_lm_head:
- # mask token + layernorm
- rename_keys.extend(
- [
- ("mask_token", "beit.embeddings.mask_token"),
- ("norm.weight", "layernorm.weight"),
- ("norm.bias", "layernorm.bias"),
- ]
- )
- else:
- # layernorm + classification head
- rename_keys.extend(
- [
- ("fc_norm.weight", "beit.pooler.layernorm.weight"),
- ("fc_norm.bias", "beit.pooler.layernorm.bias"),
- ("head.weight", "classifier.weight"),
- ("head.bias", "classifier.bias"),
- ]
- )
-
- return rename_keys
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config, has_lm_head=False, is_semantic=False):
- for i in range(config.num_hidden_layers):
- prefix = "backbone." if is_semantic else ""
- # queries, keys and values
- in_proj_weight = state_dict.pop(f"{prefix}blocks.{i}.attn.qkv.weight")
- q_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.q_bias")
- v_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.v_bias")
-
- state_dict[f"beit.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[
- : config.hidden_size, :
- ]
- state_dict[f"beit.encoder.layer.{i}.attention.attention.query.bias"] = q_bias
- state_dict[f"beit.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- config.hidden_size : config.hidden_size * 2, :
- ]
- state_dict[f"beit.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
- -config.hidden_size :, :
- ]
- state_dict[f"beit.encoder.layer.{i}.attention.attention.value.bias"] = v_bias
-
- # gamma_1 and gamma_2
- # we call them lambda because otherwise they are renamed when using .from_pretrained
- gamma_1 = state_dict.pop(f"{prefix}blocks.{i}.gamma_1")
- gamma_2 = state_dict.pop(f"{prefix}blocks.{i}.gamma_2")
-
- state_dict[f"beit.encoder.layer.{i}.lambda_1"] = gamma_1
- state_dict[f"beit.encoder.layer.{i}.lambda_2"] = gamma_2
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-@torch.no_grad()
-def convert_dit_checkpoint(checkpoint_url, pytorch_dump_folder_path, push_to_hub=False):
- """
- Copy/paste/tweak model's weights to our BEiT structure.
- """
-
- # define default BEiT configuration
- has_lm_head = False if "rvlcdip" in checkpoint_url else True
- config = BeitConfig(use_absolute_position_embeddings=True, use_mask_token=has_lm_head)
-
- # size of the architecture
- if "large" in checkpoint_url or "dit-l" in checkpoint_url:
- config.hidden_size = 1024
- config.intermediate_size = 4096
- config.num_hidden_layers = 24
- config.num_attention_heads = 16
-
- # labels
- if "rvlcdip" in checkpoint_url:
- config.num_labels = 16
- repo_id = "huggingface/label-files"
- filename = "rvlcdip-id2label.json"
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
-
- # load state_dict of original model, remove and rename some keys
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["model"]
-
- rename_keys = create_rename_keys(config, has_lm_head=has_lm_head)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- read_in_q_k_v(state_dict, config, has_lm_head=has_lm_head)
-
- # load HuggingFace model
- model = BeitForMaskedImageModeling(config) if has_lm_head else BeitForImageClassification(config)
- model.eval()
- model.load_state_dict(state_dict)
-
- # Check outputs on an image
- image_processor = BeitImageProcessor(
- size=config.image_size, resample=PILImageResampling.BILINEAR, do_center_crop=False
- )
- image = prepare_img()
-
- encoding = image_processor(images=image, return_tensors="pt")
- pixel_values = encoding["pixel_values"]
-
- outputs = model(pixel_values)
- logits = outputs.logits
-
- # verify logits
- expected_shape = [1, 16] if "rvlcdip" in checkpoint_url else [1, 196, 8192]
- assert logits.shape == torch.Size(expected_shape), "Shape of logits not as expected"
-
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- print(f"Saving image processor to {pytorch_dump_folder_path}")
- image_processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- if has_lm_head:
- model_name = "dit-base" if "base" in checkpoint_url else "dit-large"
- else:
- model_name = "dit-base-finetuned-rvlcdip" if "dit-b" in checkpoint_url else "dit-large-finetuned-rvlcdip"
- image_processor.push_to_hub(
- repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
- organization="nielsr",
- commit_message="Add image processor",
- use_temp_dir=True,
- )
- model.push_to_hub(
- repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
- organization="nielsr",
- commit_message="Add model",
- use_temp_dir=True,
- )
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- "--checkpoint_url",
- default="https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth",
- type=str,
- help="URL to the original PyTorch checkpoint (.pth file).",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- )
- args = parser.parse_args()
- convert_dit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/donut/__init__.py b/transformers/models/donut/__init__.py
deleted file mode 100644
index c548a181a3bf3023fd64defca5a3748624db6b7c..0000000000000000000000000000000000000000
--- a/transformers/models/donut/__init__.py
+++ /dev/null
@@ -1,74 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {
- "configuration_donut_swin": ["DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP", "DonutSwinConfig"],
- "processing_donut": ["DonutProcessor"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_donut_swin"] = [
- "DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DonutSwinModel",
- "DonutSwinPreTrainedModel",
- ]
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_donut"] = ["DonutFeatureExtractor"]
- _import_structure["image_processing_donut"] = ["DonutImageProcessor"]
-
-
-if TYPE_CHECKING:
- from .configuration_donut_swin import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, DonutSwinConfig
- from .processing_donut import DonutProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_donut_swin import (
- DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
- DonutSwinModel,
- DonutSwinPreTrainedModel,
- )
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_donut import DonutFeatureExtractor
- from .image_processing_donut import DonutImageProcessor
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/donut/__pycache__/__init__.cpython-310.pyc b/transformers/models/donut/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 0d619155cbbd722b3d20cf1c7d837aa1880bb772..0000000000000000000000000000000000000000
Binary files a/transformers/models/donut/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/donut/__pycache__/configuration_donut_swin.cpython-310.pyc b/transformers/models/donut/__pycache__/configuration_donut_swin.cpython-310.pyc
deleted file mode 100644
index d32768bb8e9d284d47a71b183fb0dbed795c9032..0000000000000000000000000000000000000000
Binary files a/transformers/models/donut/__pycache__/configuration_donut_swin.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/donut/__pycache__/convert_donut_to_pytorch.cpython-310.pyc b/transformers/models/donut/__pycache__/convert_donut_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 539f8cb11fc68815bbf54129edff6aeac2c42e19..0000000000000000000000000000000000000000
Binary files a/transformers/models/donut/__pycache__/convert_donut_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/donut/__pycache__/feature_extraction_donut.cpython-310.pyc b/transformers/models/donut/__pycache__/feature_extraction_donut.cpython-310.pyc
deleted file mode 100644
index 26224e9ac16d081fe86c69ef7eb9df3cbc259141..0000000000000000000000000000000000000000
Binary files a/transformers/models/donut/__pycache__/feature_extraction_donut.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/donut/__pycache__/image_processing_donut.cpython-310.pyc b/transformers/models/donut/__pycache__/image_processing_donut.cpython-310.pyc
deleted file mode 100644
index 6c91f48f07238b131a65ad2f7bc7746c3e7e39dd..0000000000000000000000000000000000000000
Binary files a/transformers/models/donut/__pycache__/image_processing_donut.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/donut/__pycache__/modeling_donut_swin.cpython-310.pyc b/transformers/models/donut/__pycache__/modeling_donut_swin.cpython-310.pyc
deleted file mode 100644
index a3ed1f3f74380ee7de28a0c97bb6b251af3f3f2d..0000000000000000000000000000000000000000
Binary files a/transformers/models/donut/__pycache__/modeling_donut_swin.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/donut/__pycache__/processing_donut.cpython-310.pyc b/transformers/models/donut/__pycache__/processing_donut.cpython-310.pyc
deleted file mode 100644
index 61394bac8c8cf9119de039b06aadd38ee1f37c61..0000000000000000000000000000000000000000
Binary files a/transformers/models/donut/__pycache__/processing_donut.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/donut/configuration_donut_swin.py b/transformers/models/donut/configuration_donut_swin.py
deleted file mode 100644
index e57ddb255a71185cec8567db362def58e16e5fc1..0000000000000000000000000000000000000000
--- a/transformers/models/donut/configuration_donut_swin.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Donut Swin Transformer model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DonutSwinConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DonutSwinModel`]. It is used to instantiate a
- Donut model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the Donut
- [naver-clova-ix/donut-base](https://huggingface.co/naver-clova-ix/donut-base) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 4):
- The size (resolution) of each patch.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- embed_dim (`int`, *optional*, defaults to 96):
- Dimensionality of patch embedding.
- depths (`list(int)`, *optional*, defaults to `[2, 2, 6, 2]`):
- Depth of each layer in the Transformer encoder.
- num_heads (`list(int)`, *optional*, defaults to `[3, 6, 12, 24]`):
- Number of attention heads in each layer of the Transformer encoder.
- window_size (`int`, *optional*, defaults to 7):
- Size of windows.
- mlp_ratio (`float`, *optional*, defaults to 4.0):
- Ratio of MLP hidden dimensionality to embedding dimensionality.
- qkv_bias (`bool`, *optional*, defaults to `True`):
- Whether or not a learnable bias should be added to the queries, keys and values.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings and encoder.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- drop_path_rate (`float`, *optional*, defaults to 0.1):
- Stochastic depth rate.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder. If string, `"gelu"`, `"relu"`,
- `"selu"` and `"gelu_new"` are supported.
- use_absolute_embeddings (`bool`, *optional*, defaults to `False`):
- Whether or not to add absolute position embeddings to the patch embeddings.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the layer normalization layers.
-
- Example:
-
- ```python
- >>> from transformers import DonutSwinConfig, DonutSwinModel
-
- >>> # Initializing a Donut naver-clova-ix/donut-base style configuration
- >>> configuration = DonutSwinConfig()
-
- >>> # Randomly initializing a model from the naver-clova-ix/donut-base style configuration
- >>> model = DonutSwinModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "donut-swin"
-
- attribute_map = {
- "num_attention_heads": "num_heads",
- "num_hidden_layers": "num_layers",
- }
-
- def __init__(
- self,
- image_size=224,
- patch_size=4,
- num_channels=3,
- embed_dim=96,
- depths=[2, 2, 6, 2],
- num_heads=[3, 6, 12, 24],
- window_size=7,
- mlp_ratio=4.0,
- qkv_bias=True,
- hidden_dropout_prob=0.0,
- attention_probs_dropout_prob=0.0,
- drop_path_rate=0.1,
- hidden_act="gelu",
- use_absolute_embeddings=False,
- initializer_range=0.02,
- layer_norm_eps=1e-5,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.embed_dim = embed_dim
- self.depths = depths
- self.num_layers = len(depths)
- self.num_heads = num_heads
- self.window_size = window_size
- self.mlp_ratio = mlp_ratio
- self.qkv_bias = qkv_bias
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.drop_path_rate = drop_path_rate
- self.hidden_act = hidden_act
- self.use_absolute_embeddings = use_absolute_embeddings
- self.layer_norm_eps = layer_norm_eps
- self.initializer_range = initializer_range
- # we set the hidden_size attribute in order to make Swin work with VisionEncoderDecoderModel
- # this indicates the channel dimension after the last stage of the model
- self.hidden_size = int(embed_dim * 2 ** (len(depths) - 1))
diff --git a/transformers/models/donut/convert_donut_to_pytorch.py b/transformers/models/donut/convert_donut_to_pytorch.py
deleted file mode 100644
index 13f669ad97fdcc5bbfcbb2a92536fcca491253a5..0000000000000000000000000000000000000000
--- a/transformers/models/donut/convert_donut_to_pytorch.py
+++ /dev/null
@@ -1,234 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Donut checkpoints using the original `donut-python` library. URL: https://github.com/clovaai/donut"""
-
-import argparse
-
-import torch
-from datasets import load_dataset
-from donut import DonutModel
-
-from transformers import (
- DonutImageProcessor,
- DonutProcessor,
- DonutSwinConfig,
- DonutSwinModel,
- MBartConfig,
- MBartForCausalLM,
- VisionEncoderDecoderModel,
- XLMRobertaTokenizerFast,
-)
-
-
-def get_configs(model):
- original_config = model.config
-
- encoder_config = DonutSwinConfig(
- image_size=original_config.input_size,
- patch_size=4,
- depths=original_config.encoder_layer,
- num_heads=[4, 8, 16, 32],
- window_size=original_config.window_size,
- embed_dim=128,
- )
- decoder_config = MBartConfig(
- is_decoder=True,
- is_encoder_decoder=False,
- add_cross_attention=True,
- decoder_layers=original_config.decoder_layer,
- max_position_embeddings=original_config.max_position_embeddings,
- vocab_size=len(
- model.decoder.tokenizer
- ), # several special tokens are added to the vocab of XLMRobertaTokenizer, see repo on the hub (added_tokens.json)
- scale_embedding=True,
- add_final_layer_norm=True,
- )
-
- return encoder_config, decoder_config
-
-
-def rename_key(name):
- if "encoder.model" in name:
- name = name.replace("encoder.model", "encoder")
- if "decoder.model" in name:
- name = name.replace("decoder.model", "decoder")
- if "patch_embed.proj" in name:
- name = name.replace("patch_embed.proj", "embeddings.patch_embeddings.projection")
- if "patch_embed.norm" in name:
- name = name.replace("patch_embed.norm", "embeddings.norm")
- if name.startswith("encoder"):
- if "layers" in name:
- name = "encoder." + name
- if "attn.proj" in name:
- name = name.replace("attn.proj", "attention.output.dense")
- if "attn" in name and "mask" not in name:
- name = name.replace("attn", "attention.self")
- if "norm1" in name:
- name = name.replace("norm1", "layernorm_before")
- if "norm2" in name:
- name = name.replace("norm2", "layernorm_after")
- if "mlp.fc1" in name:
- name = name.replace("mlp.fc1", "intermediate.dense")
- if "mlp.fc2" in name:
- name = name.replace("mlp.fc2", "output.dense")
-
- if name == "encoder.norm.weight":
- name = "encoder.layernorm.weight"
- if name == "encoder.norm.bias":
- name = "encoder.layernorm.bias"
-
- return name
-
-
-def convert_state_dict(orig_state_dict, model):
- for key in orig_state_dict.copy().keys():
- val = orig_state_dict.pop(key)
-
- if "qkv" in key:
- key_split = key.split(".")
- layer_num = int(key_split[3])
- block_num = int(key_split[5])
- dim = model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size
-
- if "weight" in key:
- orig_state_dict[
- f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.query.weight"
- ] = val[:dim, :]
- orig_state_dict[
- f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.key.weight"
- ] = val[dim : dim * 2, :]
- orig_state_dict[
- f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.value.weight"
- ] = val[-dim:, :]
- else:
- orig_state_dict[
- f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.query.bias"
- ] = val[:dim]
- orig_state_dict[
- f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.key.bias"
- ] = val[dim : dim * 2]
- orig_state_dict[
- f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.value.bias"
- ] = val[-dim:]
- elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]:
- # HuggingFace implementation doesn't use attn_mask buffer
- # and model doesn't use final LayerNorms for the encoder
- pass
- else:
- orig_state_dict[rename_key(key)] = val
-
- return orig_state_dict
-
-
-def convert_donut_checkpoint(model_name, pytorch_dump_folder_path=None, push_to_hub=False):
- # load original model
- original_model = DonutModel.from_pretrained(model_name).eval()
-
- # load HuggingFace model
- encoder_config, decoder_config = get_configs(original_model)
- encoder = DonutSwinModel(encoder_config)
- decoder = MBartForCausalLM(decoder_config)
- model = VisionEncoderDecoderModel(encoder=encoder, decoder=decoder)
- model.eval()
-
- state_dict = original_model.state_dict()
- new_state_dict = convert_state_dict(state_dict, model)
- model.load_state_dict(new_state_dict)
-
- # verify results on scanned document
- dataset = load_dataset("hf-internal-testing/example-documents")
- image = dataset["test"][0]["image"].convert("RGB")
-
- tokenizer = XLMRobertaTokenizerFast.from_pretrained(model_name, from_slow=True)
- image_processor = DonutImageProcessor(
- do_align_long_axis=original_model.config.align_long_axis, size=original_model.config.input_size[::-1]
- )
- processor = DonutProcessor(image_processor, tokenizer)
- pixel_values = processor(image, return_tensors="pt").pixel_values
-
- if model_name == "naver-clova-ix/donut-base-finetuned-docvqa":
- task_prompt = "{user_input}"
- question = "When is the coffee break?"
- task_prompt = task_prompt.replace("{user_input}", question)
- elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip":
- task_prompt = ""
- elif model_name in [
- "naver-clova-ix/donut-base-finetuned-cord-v1",
- "naver-clova-ix/donut-base-finetuned-cord-v1-2560",
- ]:
- task_prompt = ""
- elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2":
- task_prompt = "s_cord-v2>"
- elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket":
- task_prompt = ""
- elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]:
- # use a random prompt
- task_prompt = "hello world"
- else:
- raise ValueError("Model name not supported")
- prompt_tensors = original_model.decoder.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt")[
- "input_ids"
- ]
-
- original_patch_embed = original_model.encoder.model.patch_embed(pixel_values)
- patch_embeddings, _ = model.encoder.embeddings(pixel_values)
- assert torch.allclose(original_patch_embed, patch_embeddings, atol=1e-3)
-
- # verify encoder hidden states
- original_last_hidden_state = original_model.encoder(pixel_values)
- last_hidden_state = model.encoder(pixel_values).last_hidden_state
- assert torch.allclose(original_last_hidden_state, last_hidden_state, atol=1e-2)
-
- # verify decoder hidden states
- original_logits = original_model(pixel_values, prompt_tensors, None).logits
- logits = model(pixel_values, decoder_input_ids=prompt_tensors).logits
- assert torch.allclose(original_logits, logits, atol=1e-3)
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- print(f"Saving model and processor to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- model.push_to_hub("nielsr/" + model_name.split("/")[-1], commit_message="Update model")
- processor.push_to_hub("nielsr/" + model_name.split("/")[-1], commit_message="Update model")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="naver-clova-ix/donut-base-finetuned-docvqa",
- required=False,
- type=str,
- help="Name of the original model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- required=False,
- type=str,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- help="Whether or not to push the converted model and processor to the 🤗 hub.",
- )
-
- args = parser.parse_args()
- convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/donut/feature_extraction_donut.py b/transformers/models/donut/feature_extraction_donut.py
deleted file mode 100644
index e6ca078c0e8ac4939514dcb297f5d2c63de032f7..0000000000000000000000000000000000000000
--- a/transformers/models/donut/feature_extraction_donut.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for Donut."""
-
-import warnings
-
-from ...utils import logging
-from .image_processing_donut import DonutImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-class DonutFeatureExtractor(DonutImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class DonutFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
- " use DonutImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/donut/image_processing_donut.py b/transformers/models/donut/image_processing_donut.py
deleted file mode 100644
index 1c6e4723139046ae4c479690c5242e35ef5e604d..0000000000000000000000000000000000000000
--- a/transformers/models/donut/image_processing_donut.py
+++ /dev/null
@@ -1,480 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for Donut."""
-
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
-from ...image_transforms import (
- get_resize_output_image_size,
- pad,
- resize,
- to_channel_dimension_format,
-)
-from ...image_utils import (
- IMAGENET_STANDARD_MEAN,
- IMAGENET_STANDARD_STD,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- get_image_size,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import TensorType, logging
-from ...utils.import_utils import is_vision_available
-
-
-logger = logging.get_logger(__name__)
-
-
-if is_vision_available():
- import PIL
-
-
-class DonutImageProcessor(BaseImageProcessor):
- r"""
- Constructs a Donut image processor.
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by
- `do_resize` in the `preprocess` method.
- size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 224}`):
- Size of the image after resizing. The shortest edge of the image is resized to size["shortest_edge"], with
- the longest edge resized to keep the input aspect ratio. Can be overridden by `size` in the `preprocess`
- method.
- resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
- Resampling filter to use if resizing the image. Can be overridden by `resample` in the `preprocess` method.
- do_thumbnail (`bool`, *optional*, defaults to `True`):
- Whether to resize the image using thumbnail method.
- do_align_long_axis (`bool`, *optional*, defaults to `False`):
- Whether to align the long axis of the image with the long axis of `size` by rotating by 90 degrees.
- do_pad (`bool`, *optional*, defaults to `True`):
- Whether to pad the image. If `random_padding` is set to `True` in `preprocess`, each image is padded with a
- random amont of padding on each size, up to the largest image size in the batch. Otherwise, all images are
- padded to the largest image size in the batch.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by `do_rescale` in
- the `preprocess` method.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by `rescale_factor` in the `preprocess`
- method.
- do_normalize (`bool`, *optional*, defaults to `True`):
- Whether to normalize the image. Can be overridden by `do_normalize` in the `preprocess` method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
- Mean to use if normalizing the image. This is a float or list of floats the length of the number of
- channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
- Image standard deviation.
- """
-
- model_input_names = ["pixel_values"]
-
- def __init__(
- self,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- do_thumbnail: bool = True,
- do_align_long_axis: bool = False,
- do_pad: bool = True,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- do_normalize: bool = True,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
-
- size = size if size is not None else {"height": 2560, "width": 1920}
- if isinstance(size, (tuple, list)):
- # The previous feature extractor size parameter was in (width, height) format
- size = size[::-1]
- size = get_size_dict(size)
-
- self.do_resize = do_resize
- self.size = size
- self.resample = resample
- self.do_thumbnail = do_thumbnail
- self.do_align_long_axis = do_align_long_axis
- self.do_pad = do_pad
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_normalize = do_normalize
- self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size",
- "resample",
- "do_thumbnail",
- "do_align_long_axis",
- "do_pad",
- "random_padding",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "image_mean",
- "image_std",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- def align_long_axis(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.ndarray:
- """
- Align the long axis of the image to the longest axis of the specified size.
-
- Args:
- image (`np.ndarray`):
- The image to be aligned.
- size (`Dict[str, int]`):
- The size `{"height": h, "width": w}` to align the long axis to.
- data_format (`str` or `ChannelDimension`, *optional*):
- The data format of the output image. If unset, the same format as the input image is used.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
-
- Returns:
- `np.ndarray`: The aligned image.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- output_height, output_width = size["height"], size["width"]
-
- if (output_width < output_height and input_width > input_height) or (
- output_width > output_height and input_width < input_height
- ):
- image = np.rot90(image, 3)
-
- if data_format is not None:
- image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
-
- return image
-
- def pad_image(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- random_padding: bool = False,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.ndarray:
- """
- Pad the image to the specified size.
-
- Args:
- image (`np.ndarray`):
- The image to be padded.
- size (`Dict[str, int]`):
- The size `{"height": h, "width": w}` to pad the image to.
- random_padding (`bool`, *optional*, defaults to `False`):
- Whether to use random padding or not.
- data_format (`str` or `ChannelDimension`, *optional*):
- The data format of the output image. If unset, the same format as the input image is used.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- output_height, output_width = size["height"], size["width"]
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
-
- delta_width = output_width - input_width
- delta_height = output_height - input_height
-
- if random_padding:
- pad_top = np.random.randint(low=0, high=delta_height + 1)
- pad_left = np.random.randint(low=0, high=delta_width + 1)
- else:
- pad_top = delta_height // 2
- pad_left = delta_width // 2
-
- pad_bottom = delta_height - pad_top
- pad_right = delta_width - pad_left
-
- padding = ((pad_top, pad_bottom), (pad_left, pad_right))
- return pad(image, padding, data_format=data_format, input_data_format=input_data_format)
-
- def pad(self, *args, **kwargs):
- logger.info("pad is deprecated and will be removed in version 4.27. Please use pad_image instead.")
- return self.pad_image(*args, **kwargs)
-
- def thumbnail(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize the image to make a thumbnail. The image is resized so that no dimension is larger than any
- corresponding dimension of the specified size.
-
- Args:
- image (`np.ndarray`):
- The image to be resized.
- size (`Dict[str, int]`):
- The size `{"height": h, "width": w}` to resize the image to.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
- The resampling filter to use.
- data_format (`Optional[Union[str, ChannelDimension]]`, *optional*):
- The data format of the output image. If unset, the same format as the input image is used.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- input_height, input_width = get_image_size(image, channel_dim=input_data_format)
- output_height, output_width = size["height"], size["width"]
-
- # We always resize to the smallest of either the input or output size.
- height = min(input_height, output_height)
- width = min(input_width, output_width)
-
- if height == input_height and width == input_width:
- return image
-
- if input_height > input_width:
- width = int(input_width * height / input_height)
- elif input_width > input_height:
- height = int(input_height * width / input_width)
-
- return resize(
- image,
- size=(height, width),
- resample=resample,
- reducing_gap=2.0,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
-
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resizes `image` to `(height, width)` specified by `size` using the PIL library.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Size of the output image.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
- Resampling filter to use when resiizing the image.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the image. If not provided, it will be the same as the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- size = get_size_dict(size)
- shortest_edge = min(size["height"], size["width"])
- output_size = get_resize_output_image_size(
- image, size=shortest_edge, default_to_square=False, input_data_format=input_data_format
- )
- resized_image = resize(
- image,
- size=output_size,
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
- return resized_image
-
- def preprocess(
- self,
- images: ImageInput,
- do_resize: bool = None,
- size: Dict[str, int] = None,
- resample: PILImageResampling = None,
- do_thumbnail: bool = None,
- do_align_long_axis: bool = None,
- do_pad: bool = None,
- random_padding: bool = False,
- do_rescale: bool = None,
- rescale_factor: float = None,
- do_normalize: bool = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> PIL.Image.Image:
- """
- Preprocess an image or batch of images.
-
- Args:
- images (`ImageInput`):
- Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
- passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
- Size of the image after resizing. Shortest edge of the image is resized to min(size["height"],
- size["width"]) with the longest edge resized to keep the input aspect ratio.
- resample (`int`, *optional*, defaults to `self.resample`):
- Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
- has an effect if `do_resize` is set to `True`.
- do_thumbnail (`bool`, *optional*, defaults to `self.do_thumbnail`):
- Whether to resize the image using thumbnail method.
- do_align_long_axis (`bool`, *optional*, defaults to `self.do_align_long_axis`):
- Whether to align the long axis of the image with the long axis of `size` by rotating by 90 degrees.
- do_pad (`bool`, *optional*, defaults to `self.do_pad`):
- Whether to pad the image. If `random_padding` is set to `True`, each image is padded with a random
- amont of padding on each size, up to the largest image size in the batch. Otherwise, all images are
- padded to the largest image size in the batch.
- random_padding (`bool`, *optional*, defaults to `self.random_padding`):
- Whether to use random padding when padding the image. If `True`, each image in the batch with be padded
- with a random amount of padding on each side up to the size of the largest image in the batch.
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether to rescale the image pixel values.
- rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
- Rescale factor to rescale the image by if `do_rescale` is set to `True`.
- do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
- Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
- Image mean to use for normalization.
- image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
- Image standard deviation to use for normalization.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - Unset: defaults to the channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- do_resize = do_resize if do_resize is not None else self.do_resize
- size = size if size is not None else self.size
- if isinstance(size, (tuple, list)):
- # Previous feature extractor had size in (width, height) format
- size = size[::-1]
- size = get_size_dict(size)
- resample = resample if resample is not None else self.resample
- do_thumbnail = do_thumbnail if do_thumbnail is not None else self.do_thumbnail
- do_align_long_axis = do_align_long_axis if do_align_long_axis is not None else self.do_align_long_axis
- do_pad = do_pad if do_pad is not None else self.do_pad
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- do_normalize = do_normalize if do_normalize is not None else self.do_normalize
- image_mean = image_mean if image_mean is not None else self.image_mean
- image_std = image_std if image_std is not None else self.image_std
-
- images = make_list_of_images(images)
-
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_pad=do_pad,
- size_divisibility=size, # There is no pad divisibility in this processor, but pad requires the size arg.
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
-
- # All transformations expect numpy arrays.
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- if do_align_long_axis:
- images = [self.align_long_axis(image, size=size, input_data_format=input_data_format) for image in images]
-
- if do_resize:
- images = [
- self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_thumbnail:
- images = [self.thumbnail(image=image, size=size, input_data_format=input_data_format) for image in images]
-
- if do_pad:
- images = [
- self.pad_image(
- image=image, size=size, random_padding=random_padding, input_data_format=input_data_format
- )
- for image in images
- ]
-
- if do_rescale:
- images = [
- self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_normalize:
- images = [
- self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
- for image in images
- ]
-
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
- ]
-
- data = {"pixel_values": images}
- return BatchFeature(data=data, tensor_type=return_tensors)
diff --git a/transformers/models/donut/modeling_donut_swin.py b/transformers/models/donut/modeling_donut_swin.py
deleted file mode 100644
index b2aa8d61b1d8d102fa4ddbc5171c77069c265b4b..0000000000000000000000000000000000000000
--- a/transformers/models/donut/modeling_donut_swin.py
+++ /dev/null
@@ -1,955 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Donut Swin Transformer model.
-
-This implementation is identical to a regular Swin Transformer, without final layer norm on top of the final hidden
-states."""
-
-import collections.abc
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from ...activations import ACT2FN
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import find_pruneable_heads_and_indices, meshgrid, prune_linear_layer
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_donut_swin import DonutSwinConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "DonutSwinConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "https://huggingface.co/naver-clova-ix/donut-base"
-_EXPECTED_OUTPUT_SHAPE = [1, 49, 768]
-
-
-from ..deprecated._archive_maps import DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-# Copied from transformers.models.swin.modeling_swin.SwinEncoderOutput with Swin->DonutSwin
-class DonutSwinEncoderOutput(ModelOutput):
- """
- DonutSwin encoder's outputs, with potential hidden states and attentions.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each stage) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- last_hidden_state: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-@dataclass
-# Copied from transformers.models.swin.modeling_swin.SwinModelOutput with Swin->DonutSwin
-class DonutSwinModelOutput(ModelOutput):
- """
- DonutSwin model's outputs that also contains a pooling of the last hidden states.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*, returned when `add_pooling_layer=True` is passed):
- Average pooling of the last layer hidden-state.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each stage) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- last_hidden_state: torch.FloatTensor = None
- pooler_output: Optional[torch.FloatTensor] = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-# Copied from transformers.models.swin.modeling_swin.window_partition
-def window_partition(input_feature, window_size):
- """
- Partitions the given input into windows.
- """
- batch_size, height, width, num_channels = input_feature.shape
- input_feature = input_feature.view(
- batch_size, height // window_size, window_size, width // window_size, window_size, num_channels
- )
- windows = input_feature.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)
- return windows
-
-
-# Copied from transformers.models.swin.modeling_swin.window_reverse
-def window_reverse(windows, window_size, height, width):
- """
- Merges windows to produce higher resolution features.
- """
- num_channels = windows.shape[-1]
- windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels)
- windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, num_channels)
- return windows
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinEmbeddings with Swin->DonutSwin
-class DonutSwinEmbeddings(nn.Module):
- """
- Construct the patch and position embeddings. Optionally, also the mask token.
- """
-
- def __init__(self, config, use_mask_token=False):
- super().__init__()
-
- self.patch_embeddings = DonutSwinPatchEmbeddings(config)
- num_patches = self.patch_embeddings.num_patches
- self.patch_grid = self.patch_embeddings.grid_size
- self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim)) if use_mask_token else None
-
- if config.use_absolute_embeddings:
- self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.embed_dim))
- else:
- self.position_embeddings = None
-
- self.norm = nn.LayerNorm(config.embed_dim)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(
- self, pixel_values: Optional[torch.FloatTensor], bool_masked_pos: Optional[torch.BoolTensor] = None
- ) -> Tuple[torch.Tensor]:
- embeddings, output_dimensions = self.patch_embeddings(pixel_values)
- embeddings = self.norm(embeddings)
- batch_size, seq_len, _ = embeddings.size()
-
- if bool_masked_pos is not None:
- mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
- # replace the masked visual tokens by mask_tokens
- mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
- embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
-
- if self.position_embeddings is not None:
- embeddings = embeddings + self.position_embeddings
-
- embeddings = self.dropout(embeddings)
-
- return embeddings, output_dimensions
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinPatchEmbeddings
-class DonutSwinPatchEmbeddings(nn.Module):
- """
- This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
- `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
- Transformer.
- """
-
- def __init__(self, config):
- super().__init__()
- image_size, patch_size = config.image_size, config.patch_size
- num_channels, hidden_size = config.num_channels, config.embed_dim
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.num_patches = num_patches
- self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])
-
- self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
-
- def maybe_pad(self, pixel_values, height, width):
- if width % self.patch_size[1] != 0:
- pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
- pixel_values = nn.functional.pad(pixel_values, pad_values)
- if height % self.patch_size[0] != 0:
- pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
- pixel_values = nn.functional.pad(pixel_values, pad_values)
- return pixel_values
-
- def forward(self, pixel_values: Optional[torch.FloatTensor]) -> Tuple[torch.Tensor, Tuple[int]]:
- _, num_channels, height, width = pixel_values.shape
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- # pad the input to be divisible by self.patch_size, if needed
- pixel_values = self.maybe_pad(pixel_values, height, width)
- embeddings = self.projection(pixel_values)
- _, _, height, width = embeddings.shape
- output_dimensions = (height, width)
- embeddings = embeddings.flatten(2).transpose(1, 2)
-
- return embeddings, output_dimensions
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinPatchMerging
-class DonutSwinPatchMerging(nn.Module):
- """
- Patch Merging Layer.
-
- Args:
- input_resolution (`Tuple[int]`):
- Resolution of input feature.
- dim (`int`):
- Number of input channels.
- norm_layer (`nn.Module`, *optional*, defaults to `nn.LayerNorm`):
- Normalization layer class.
- """
-
- def __init__(self, input_resolution: Tuple[int], dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None:
- super().__init__()
- self.input_resolution = input_resolution
- self.dim = dim
- self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
- self.norm = norm_layer(4 * dim)
-
- def maybe_pad(self, input_feature, height, width):
- should_pad = (height % 2 == 1) or (width % 2 == 1)
- if should_pad:
- pad_values = (0, 0, 0, width % 2, 0, height % 2)
- input_feature = nn.functional.pad(input_feature, pad_values)
-
- return input_feature
-
- def forward(self, input_feature: torch.Tensor, input_dimensions: Tuple[int, int]) -> torch.Tensor:
- height, width = input_dimensions
- # `dim` is height * width
- batch_size, dim, num_channels = input_feature.shape
-
- input_feature = input_feature.view(batch_size, height, width, num_channels)
- # pad input to be disible by width and height, if needed
- input_feature = self.maybe_pad(input_feature, height, width)
- # [batch_size, height/2, width/2, num_channels]
- input_feature_0 = input_feature[:, 0::2, 0::2, :]
- # [batch_size, height/2, width/2, num_channels]
- input_feature_1 = input_feature[:, 1::2, 0::2, :]
- # [batch_size, height/2, width/2, num_channels]
- input_feature_2 = input_feature[:, 0::2, 1::2, :]
- # [batch_size, height/2, width/2, num_channels]
- input_feature_3 = input_feature[:, 1::2, 1::2, :]
- # batch_size height/2 width/2 4*num_channels
- input_feature = torch.cat([input_feature_0, input_feature_1, input_feature_2, input_feature_3], -1)
- input_feature = input_feature.view(batch_size, -1, 4 * num_channels) # batch_size height/2*width/2 4*C
-
- input_feature = self.norm(input_feature)
- input_feature = self.reduction(input_feature)
-
- return input_feature
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinDropPath
-class DonutSwinDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinSelfAttention with Swin->DonutSwin
-class DonutSwinSelfAttention(nn.Module):
- def __init__(self, config, dim, num_heads, window_size):
- super().__init__()
- if dim % num_heads != 0:
- raise ValueError(
- f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})"
- )
-
- self.num_attention_heads = num_heads
- self.attention_head_size = int(dim / num_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.window_size = (
- window_size if isinstance(window_size, collections.abc.Iterable) else (window_size, window_size)
- )
-
- self.relative_position_bias_table = nn.Parameter(
- torch.zeros((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), num_heads)
- )
-
- # get pair-wise relative position index for each token inside the window
- coords_h = torch.arange(self.window_size[0])
- coords_w = torch.arange(self.window_size[1])
- coords = torch.stack(meshgrid([coords_h, coords_w], indexing="ij"))
- coords_flatten = torch.flatten(coords, 1)
- relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
- relative_coords = relative_coords.permute(1, 2, 0).contiguous()
- relative_coords[:, :, 0] += self.window_size[0] - 1
- relative_coords[:, :, 1] += self.window_size[1] - 1
- relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
- relative_position_index = relative_coords.sum(-1)
- self.register_buffer("relative_position_index", relative_position_index)
-
- self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
- self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
- self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- def transpose_for_scores(self, x):
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- batch_size, dim, num_channels = hidden_states.shape
- mixed_query_layer = self.query(hidden_states)
-
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
-
- relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)]
- relative_position_bias = relative_position_bias.view(
- self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
- )
-
- relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
- attention_scores = attention_scores + relative_position_bias.unsqueeze(0)
-
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in DonutSwinModel forward() function)
- mask_shape = attention_mask.shape[0]
- attention_scores = attention_scores.view(
- batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim
- )
- attention_scores = attention_scores + attention_mask.unsqueeze(1).unsqueeze(0)
- attention_scores = attention_scores.view(-1, self.num_attention_heads, dim, dim)
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinSelfOutput
-class DonutSwinSelfOutput(nn.Module):
- def __init__(self, config, dim):
- super().__init__()
- self.dense = nn.Linear(dim, dim)
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinAttention with Swin->DonutSwin
-class DonutSwinAttention(nn.Module):
- def __init__(self, config, dim, num_heads, window_size):
- super().__init__()
- self.self = DonutSwinSelfAttention(config, dim, num_heads, window_size)
- self.output = DonutSwinSelfOutput(config, dim)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- self_outputs = self.self(hidden_states, attention_mask, head_mask, output_attentions)
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinIntermediate
-class DonutSwinIntermediate(nn.Module):
- def __init__(self, config, dim):
- super().__init__()
- self.dense = nn.Linear(dim, int(config.mlp_ratio * dim))
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinOutput
-class DonutSwinOutput(nn.Module):
- def __init__(self, config, dim):
- super().__init__()
- self.dense = nn.Linear(int(config.mlp_ratio * dim), dim)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinLayer with Swin->DonutSwin
-class DonutSwinLayer(nn.Module):
- def __init__(self, config, dim, input_resolution, num_heads, shift_size=0):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.shift_size = shift_size
- self.window_size = config.window_size
- self.input_resolution = input_resolution
- self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
- self.attention = DonutSwinAttention(config, dim, num_heads, window_size=self.window_size)
- self.drop_path = DonutSwinDropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
- self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
- self.intermediate = DonutSwinIntermediate(config, dim)
- self.output = DonutSwinOutput(config, dim)
-
- def set_shift_and_window_size(self, input_resolution):
- if min(input_resolution) <= self.window_size:
- # if window size is larger than input resolution, we don't partition windows
- self.shift_size = 0
- self.window_size = min(input_resolution)
-
- def get_attn_mask(self, height, width, dtype):
- if self.shift_size > 0:
- # calculate attention mask for SW-MSA
- img_mask = torch.zeros((1, height, width, 1), dtype=dtype)
- height_slices = (
- slice(0, -self.window_size),
- slice(-self.window_size, -self.shift_size),
- slice(-self.shift_size, None),
- )
- width_slices = (
- slice(0, -self.window_size),
- slice(-self.window_size, -self.shift_size),
- slice(-self.shift_size, None),
- )
- count = 0
- for height_slice in height_slices:
- for width_slice in width_slices:
- img_mask[:, height_slice, width_slice, :] = count
- count += 1
-
- mask_windows = window_partition(img_mask, self.window_size)
- mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
- attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
- attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
- else:
- attn_mask = None
- return attn_mask
-
- def maybe_pad(self, hidden_states, height, width):
- pad_right = (self.window_size - width % self.window_size) % self.window_size
- pad_bottom = (self.window_size - height % self.window_size) % self.window_size
- pad_values = (0, 0, 0, pad_right, 0, pad_bottom)
- hidden_states = nn.functional.pad(hidden_states, pad_values)
- return hidden_states, pad_values
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- input_dimensions: Tuple[int, int],
- head_mask: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = False,
- always_partition: Optional[bool] = False,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- if not always_partition:
- self.set_shift_and_window_size(input_dimensions)
- else:
- pass
- height, width = input_dimensions
- batch_size, _, channels = hidden_states.size()
- shortcut = hidden_states
-
- hidden_states = self.layernorm_before(hidden_states)
-
- hidden_states = hidden_states.view(batch_size, height, width, channels)
-
- # pad hidden_states to multiples of window size
- hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)
-
- _, height_pad, width_pad, _ = hidden_states.shape
- # cyclic shift
- if self.shift_size > 0:
- shifted_hidden_states = torch.roll(hidden_states, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
- else:
- shifted_hidden_states = hidden_states
-
- # partition windows
- hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)
- hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
- attn_mask = self.get_attn_mask(height_pad, width_pad, dtype=hidden_states.dtype)
- if attn_mask is not None:
- attn_mask = attn_mask.to(hidden_states_windows.device)
-
- attention_outputs = self.attention(
- hidden_states_windows, attn_mask, head_mask, output_attentions=output_attentions
- )
-
- attention_output = attention_outputs[0]
-
- attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels)
- shifted_windows = window_reverse(attention_windows, self.window_size, height_pad, width_pad)
-
- # reverse cyclic shift
- if self.shift_size > 0:
- attention_windows = torch.roll(shifted_windows, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
- else:
- attention_windows = shifted_windows
-
- was_padded = pad_values[3] > 0 or pad_values[5] > 0
- if was_padded:
- attention_windows = attention_windows[:, :height, :width, :].contiguous()
-
- attention_windows = attention_windows.view(batch_size, height * width, channels)
-
- hidden_states = shortcut + self.drop_path(attention_windows)
-
- layer_output = self.layernorm_after(hidden_states)
- layer_output = self.intermediate(layer_output)
- layer_output = hidden_states + self.output(layer_output)
-
- layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)
- return layer_outputs
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinStage with Swin->DonutSwin
-class DonutSwinStage(nn.Module):
- def __init__(self, config, dim, input_resolution, depth, num_heads, drop_path, downsample):
- super().__init__()
- self.config = config
- self.dim = dim
- self.blocks = nn.ModuleList(
- [
- DonutSwinLayer(
- config=config,
- dim=dim,
- input_resolution=input_resolution,
- num_heads=num_heads,
- shift_size=0 if (i % 2 == 0) else config.window_size // 2,
- )
- for i in range(depth)
- ]
- )
-
- # patch merging layer
- if downsample is not None:
- self.downsample = downsample(input_resolution, dim=dim, norm_layer=nn.LayerNorm)
- else:
- self.downsample = None
-
- self.pointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- input_dimensions: Tuple[int, int],
- head_mask: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = False,
- always_partition: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- height, width = input_dimensions
- for i, layer_module in enumerate(self.blocks):
- layer_head_mask = head_mask[i] if head_mask is not None else None
-
- layer_outputs = layer_module(
- hidden_states, input_dimensions, layer_head_mask, output_attentions, always_partition
- )
-
- hidden_states = layer_outputs[0]
-
- hidden_states_before_downsampling = hidden_states
- if self.downsample is not None:
- height_downsampled, width_downsampled = (height + 1) // 2, (width + 1) // 2
- output_dimensions = (height, width, height_downsampled, width_downsampled)
- hidden_states = self.downsample(hidden_states_before_downsampling, input_dimensions)
- else:
- output_dimensions = (height, width, height, width)
-
- stage_outputs = (hidden_states, hidden_states_before_downsampling, output_dimensions)
-
- if output_attentions:
- stage_outputs += layer_outputs[1:]
- return stage_outputs
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinEncoder with Swin->DonutSwin
-class DonutSwinEncoder(nn.Module):
- def __init__(self, config, grid_size):
- super().__init__()
- self.num_layers = len(config.depths)
- self.config = config
- dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
- self.layers = nn.ModuleList(
- [
- DonutSwinStage(
- config=config,
- dim=int(config.embed_dim * 2**i_layer),
- input_resolution=(grid_size[0] // (2**i_layer), grid_size[1] // (2**i_layer)),
- depth=config.depths[i_layer],
- num_heads=config.num_heads[i_layer],
- drop_path=dpr[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])],
- downsample=DonutSwinPatchMerging if (i_layer < self.num_layers - 1) else None,
- )
- for i_layer in range(self.num_layers)
- ]
- )
-
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- input_dimensions: Tuple[int, int],
- head_mask: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = False,
- output_hidden_states: Optional[bool] = False,
- output_hidden_states_before_downsampling: Optional[bool] = False,
- always_partition: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple, DonutSwinEncoderOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_reshaped_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- if output_hidden_states:
- batch_size, _, hidden_size = hidden_states.shape
- # rearrange b (h w) c -> b c h w
- reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
- reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
-
- for i, layer_module in enumerate(self.layers):
- layer_head_mask = head_mask[i] if head_mask is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- input_dimensions,
- layer_head_mask,
- output_attentions,
- always_partition,
- )
- else:
- layer_outputs = layer_module(
- hidden_states, input_dimensions, layer_head_mask, output_attentions, always_partition
- )
-
- hidden_states = layer_outputs[0]
- hidden_states_before_downsampling = layer_outputs[1]
- output_dimensions = layer_outputs[2]
-
- input_dimensions = (output_dimensions[-2], output_dimensions[-1])
-
- if output_hidden_states and output_hidden_states_before_downsampling:
- batch_size, _, hidden_size = hidden_states_before_downsampling.shape
- # rearrange b (h w) c -> b c h w
- # here we use the original (not downsampled) height and width
- reshaped_hidden_state = hidden_states_before_downsampling.view(
- batch_size, *(output_dimensions[0], output_dimensions[1]), hidden_size
- )
- reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states_before_downsampling,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
- elif output_hidden_states and not output_hidden_states_before_downsampling:
- batch_size, _, hidden_size = hidden_states.shape
- # rearrange b (h w) c -> b c h w
- reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
- reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
-
- if output_attentions:
- all_self_attentions += layer_outputs[3:]
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
-
- return DonutSwinEncoderOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- reshaped_hidden_states=all_reshaped_hidden_states,
- )
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinPreTrainedModel with Swin->DonutSwin
-class DonutSwinPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DonutSwinConfig
- base_model_prefix = "swin"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-SWIN_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`DonutSwinConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-SWIN_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`DonutImageProcessor.__call__`] for details.
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Donut Swin Model transformer outputting raw hidden-states without any specific head on top.",
- SWIN_START_DOCSTRING,
-)
-class DonutSwinModel(DonutSwinPreTrainedModel):
- def __init__(self, config, add_pooling_layer=True, use_mask_token=False):
- super().__init__(config)
- self.config = config
- self.num_layers = len(config.depths)
- self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1))
-
- self.embeddings = DonutSwinEmbeddings(config, use_mask_token=use_mask_token)
- self.encoder = DonutSwinEncoder(config, self.embeddings.patch_grid)
-
- self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.patch_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(SWIN_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=DonutSwinModelOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, DonutSwinModelOutput]:
- r"""
- bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, len(self.config.depths))
-
- embedding_output, input_dimensions = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
-
- encoder_outputs = self.encoder(
- embedding_output,
- input_dimensions,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = encoder_outputs[0]
-
- pooled_output = None
- if self.pooler is not None:
- pooled_output = self.pooler(sequence_output.transpose(1, 2))
- pooled_output = torch.flatten(pooled_output, 1)
-
- if not return_dict:
- output = (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return output
-
- return DonutSwinModelOutput(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- reshaped_hidden_states=encoder_outputs.reshaped_hidden_states,
- )
diff --git a/transformers/models/donut/processing_donut.py b/transformers/models/donut/processing_donut.py
deleted file mode 100644
index 1f03fd6306fc0a9940fe1bf2b497d705fc4e22ea..0000000000000000000000000000000000000000
--- a/transformers/models/donut/processing_donut.py
+++ /dev/null
@@ -1,196 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Processor class for Donut.
-"""
-import re
-import warnings
-from contextlib import contextmanager
-
-from ...processing_utils import ProcessorMixin
-
-
-class DonutProcessor(ProcessorMixin):
- r"""
- Constructs a Donut processor which wraps a Donut image processor and an XLMRoBERTa tokenizer into a single
- processor.
-
- [`DonutProcessor`] offers all the functionalities of [`DonutImageProcessor`] and
- [`XLMRobertaTokenizer`/`XLMRobertaTokenizerFast`]. See the [`~DonutProcessor.__call__`] and
- [`~DonutProcessor.decode`] for more information.
-
- Args:
- image_processor ([`DonutImageProcessor`], *optional*):
- An instance of [`DonutImageProcessor`]. The image processor is a required input.
- tokenizer ([`XLMRobertaTokenizer`/`XLMRobertaTokenizerFast`], *optional*):
- An instance of [`XLMRobertaTokenizer`/`XLMRobertaTokenizerFast`]. The tokenizer is a required input.
- """
-
- attributes = ["image_processor", "tokenizer"]
- image_processor_class = "AutoImageProcessor"
- tokenizer_class = "AutoTokenizer"
-
- def __init__(self, image_processor=None, tokenizer=None, **kwargs):
- feature_extractor = None
- if "feature_extractor" in kwargs:
- warnings.warn(
- "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
- " instead.",
- FutureWarning,
- )
- feature_extractor = kwargs.pop("feature_extractor")
-
- image_processor = image_processor if image_processor is not None else feature_extractor
- if image_processor is None:
- raise ValueError("You need to specify an `image_processor`.")
- if tokenizer is None:
- raise ValueError("You need to specify a `tokenizer`.")
-
- super().__init__(image_processor, tokenizer)
- self.current_processor = self.image_processor
- self._in_target_context_manager = False
-
- def __call__(self, *args, **kwargs):
- """
- When used in normal mode, this method forwards all its arguments to AutoImageProcessor's
- [`~AutoImageProcessor.__call__`] and returns its output. If used in the context
- [`~DonutProcessor.as_target_processor`] this method forwards all its arguments to DonutTokenizer's
- [`~DonutTokenizer.__call__`]. Please refer to the doctsring of the above two methods for more information.
- """
- # For backward compatibility
- if self._in_target_context_manager:
- return self.current_processor(*args, **kwargs)
-
- images = kwargs.pop("images", None)
- text = kwargs.pop("text", None)
- if len(args) > 0:
- images = args[0]
- args = args[1:]
-
- if images is None and text is None:
- raise ValueError("You need to specify either an `images` or `text` input to process.")
-
- if images is not None:
- inputs = self.image_processor(images, *args, **kwargs)
- if text is not None:
- encodings = self.tokenizer(text, **kwargs)
-
- if text is None:
- return inputs
- elif images is None:
- return encodings
- else:
- inputs["labels"] = encodings["input_ids"]
- return inputs
-
- def batch_decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to DonutTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please refer
- to the docstring of this method for more information.
- """
- return self.tokenizer.batch_decode(*args, **kwargs)
-
- def decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to DonutTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to the
- docstring of this method for more information.
- """
- return self.tokenizer.decode(*args, **kwargs)
-
- @contextmanager
- def as_target_processor(self):
- """
- Temporarily sets the tokenizer for processing the input. Useful for encoding the labels when fine-tuning TrOCR.
- """
- warnings.warn(
- "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
- "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
- "your images inputs, or in a separate call."
- )
- self._in_target_context_manager = True
- self.current_processor = self.tokenizer
- yield
- self.current_processor = self.image_processor
- self._in_target_context_manager = False
-
- def token2json(self, tokens, is_inner_value=False, added_vocab=None):
- """
- Convert a (generated) token sequence into an ordered JSON format.
- """
- if added_vocab is None:
- added_vocab = self.tokenizer.get_added_vocab()
-
- output = {}
-
- while tokens:
- start_token = re.search(r"", tokens, re.IGNORECASE)
- if start_token is None:
- break
- key = start_token.group(1)
- key_escaped = re.escape(key)
-
- end_token = re.search(rf"", tokens, re.IGNORECASE)
- start_token = start_token.group()
- if end_token is None:
- tokens = tokens.replace(start_token, "")
- else:
- end_token = end_token.group()
- start_token_escaped = re.escape(start_token)
- end_token_escaped = re.escape(end_token)
- content = re.search(
- f"{start_token_escaped}(.*?){end_token_escaped}", tokens, re.IGNORECASE | re.DOTALL
- )
- if content is not None:
- content = content.group(1).strip()
- if r""):
- leaf = leaf.strip()
- if leaf in added_vocab and leaf[0] == "<" and leaf[-2:] == "/>":
- leaf = leaf[1:-2] # for categorical special tokens
- output[key].append(leaf)
- if len(output[key]) == 1:
- output[key] = output[key][0]
-
- tokens = tokens[tokens.find(end_token) + len(end_token) :].strip()
- if tokens[:6] == r"": # non-leaf nodes
- return [output] + self.token2json(tokens[6:], is_inner_value=True, added_vocab=added_vocab)
-
- if len(output):
- return [output] if is_inner_value else output
- else:
- return [] if is_inner_value else {"text_sequence": tokens}
-
- @property
- def feature_extractor_class(self):
- warnings.warn(
- "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.",
- FutureWarning,
- )
- return self.image_processor_class
-
- @property
- def feature_extractor(self):
- warnings.warn(
- "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.",
- FutureWarning,
- )
- return self.image_processor
diff --git a/transformers/models/dpr/__init__.py b/transformers/models/dpr/__init__.py
deleted file mode 100644
index 6ea8b78e503739e91991ff14b23d8abb0cbdb975..0000000000000000000000000000000000000000
--- a/transformers/models/dpr/__init__.py
+++ /dev/null
@@ -1,148 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_tf_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_dpr": ["DPR_PRETRAINED_CONFIG_ARCHIVE_MAP", "DPRConfig"],
- "tokenization_dpr": [
- "DPRContextEncoderTokenizer",
- "DPRQuestionEncoderTokenizer",
- "DPRReaderOutput",
- "DPRReaderTokenizer",
- ],
-}
-
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_dpr_fast"] = [
- "DPRContextEncoderTokenizerFast",
- "DPRQuestionEncoderTokenizerFast",
- "DPRReaderTokenizerFast",
- ]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_dpr"] = [
- "DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DPRContextEncoder",
- "DPRPretrainedContextEncoder",
- "DPRPreTrainedModel",
- "DPRPretrainedQuestionEncoder",
- "DPRPretrainedReader",
- "DPRQuestionEncoder",
- "DPRReader",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_dpr"] = [
- "TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFDPRContextEncoder",
- "TFDPRPretrainedContextEncoder",
- "TFDPRPretrainedQuestionEncoder",
- "TFDPRPretrainedReader",
- "TFDPRQuestionEncoder",
- "TFDPRReader",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
- from .tokenization_dpr import (
- DPRContextEncoderTokenizer,
- DPRQuestionEncoderTokenizer,
- DPRReaderOutput,
- DPRReaderTokenizer,
- )
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_dpr_fast import (
- DPRContextEncoderTokenizerFast,
- DPRQuestionEncoderTokenizerFast,
- DPRReaderTokenizerFast,
- )
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_dpr import (
- DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
- DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
- DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
- DPRContextEncoder,
- DPRPretrainedContextEncoder,
- DPRPreTrainedModel,
- DPRPretrainedQuestionEncoder,
- DPRPretrainedReader,
- DPRQuestionEncoder,
- DPRReader,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_dpr import (
- TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
- TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
- TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFDPRContextEncoder,
- TFDPRPretrainedContextEncoder,
- TFDPRPretrainedQuestionEncoder,
- TFDPRPretrainedReader,
- TFDPRQuestionEncoder,
- TFDPRReader,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/dpr/__pycache__/__init__.cpython-310.pyc b/transformers/models/dpr/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index fe5fc80023cee5b76beb8c3b16146cafa388adec..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpr/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpr/__pycache__/configuration_dpr.cpython-310.pyc b/transformers/models/dpr/__pycache__/configuration_dpr.cpython-310.pyc
deleted file mode 100644
index d83c1cc2b9c6b7b959656d0e8ade21549c94b718..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpr/__pycache__/configuration_dpr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpr/__pycache__/convert_dpr_original_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/dpr/__pycache__/convert_dpr_original_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index d29c1a6c96bd860916d7915eadeff97ea3fe6ee2..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpr/__pycache__/convert_dpr_original_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpr/__pycache__/modeling_dpr.cpython-310.pyc b/transformers/models/dpr/__pycache__/modeling_dpr.cpython-310.pyc
deleted file mode 100644
index 6276c10e2e10a9726c7e150dd978295443bcdbb6..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpr/__pycache__/modeling_dpr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpr/__pycache__/modeling_tf_dpr.cpython-310.pyc b/transformers/models/dpr/__pycache__/modeling_tf_dpr.cpython-310.pyc
deleted file mode 100644
index 9b2d3b050fc9f08ee96a175242c90e310d18571e..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpr/__pycache__/modeling_tf_dpr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpr/__pycache__/tokenization_dpr.cpython-310.pyc b/transformers/models/dpr/__pycache__/tokenization_dpr.cpython-310.pyc
deleted file mode 100644
index 0d00e17dcccf06e397eddcfa15dd7e3f1c5e3e97..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpr/__pycache__/tokenization_dpr.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpr/__pycache__/tokenization_dpr_fast.cpython-310.pyc b/transformers/models/dpr/__pycache__/tokenization_dpr_fast.cpython-310.pyc
deleted file mode 100644
index ceea0339abab559c3c0de4d6babe9b1e527fa343..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpr/__pycache__/tokenization_dpr_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpr/configuration_dpr.py b/transformers/models/dpr/configuration_dpr.py
deleted file mode 100644
index 74ac90a4beb50812eef881e67218812834d7bf93..0000000000000000000000000000000000000000
--- a/transformers/models/dpr/configuration_dpr.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# coding=utf-8
-# Copyright 2010, DPR authors, The Hugging Face Team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DPR model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DPRConfig(PretrainedConfig):
- r"""
- [`DPRConfig`] is the configuration class to store the configuration of a *DPRModel*.
-
- This is the configuration class to store the configuration of a [`DPRContextEncoder`], [`DPRQuestionEncoder`], or a
- [`DPRReader`]. It is used to instantiate the components of the DPR model according to the specified arguments,
- defining the model component architectures. Instantiating a configuration with the defaults will yield a similar
- configuration to that of the DPRContextEncoder
- [facebook/dpr-ctx_encoder-single-nq-base](https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base)
- architecture.
-
- This class is a subclass of [`BertConfig`]. Please check the superclass for the documentation of all kwargs.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the DPR model. Defines the different tokens that can be represented by the *inputs_ids*
- passed to the forward method of [`BertModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 2):
- The vocabulary size of the *token_type_ids* passed into [`BertModel`].
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- pad_token_id (`int`, *optional*, defaults to 0):
- Padding token id.
- position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
- Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
- positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
- [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
- For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
- with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
- projection_dim (`int`, *optional*, defaults to 0):
- Dimension of the projection for the context and question encoders. If it is set to zero (default), then no
- projection is done.
-
- Example:
-
- ```python
- >>> from transformers import DPRConfig, DPRContextEncoder
-
- >>> # Initializing a DPR facebook/dpr-ctx_encoder-single-nq-base style configuration
- >>> configuration = DPRConfig()
-
- >>> # Initializing a model (with random weights) from the facebook/dpr-ctx_encoder-single-nq-base style configuration
- >>> model = DPRContextEncoder(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "dpr"
-
- def __init__(
- self,
- vocab_size=30522,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=2,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- pad_token_id=0,
- position_embedding_type="absolute",
- projection_dim: int = 0,
- **kwargs,
- ):
- super().__init__(pad_token_id=pad_token_id, **kwargs)
-
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.hidden_act = hidden_act
- self.intermediate_size = intermediate_size
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.projection_dim = projection_dim
- self.position_embedding_type = position_embedding_type
diff --git a/transformers/models/dpr/convert_dpr_original_checkpoint_to_pytorch.py b/transformers/models/dpr/convert_dpr_original_checkpoint_to_pytorch.py
deleted file mode 100644
index c11345d1eb4e466004c77743884ec78d4f54f97b..0000000000000000000000000000000000000000
--- a/transformers/models/dpr/convert_dpr_original_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,143 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import collections
-from pathlib import Path
-
-import torch
-from torch.serialization import default_restore_location
-
-from transformers import BertConfig, DPRConfig, DPRContextEncoder, DPRQuestionEncoder, DPRReader
-
-
-CheckpointState = collections.namedtuple(
- "CheckpointState", ["model_dict", "optimizer_dict", "scheduler_dict", "offset", "epoch", "encoder_params"]
-)
-
-
-def load_states_from_checkpoint(model_file: str) -> CheckpointState:
- print(f"Reading saved model from {model_file}")
- state_dict = torch.load(model_file, map_location=lambda s, l: default_restore_location(s, "cpu"))
- return CheckpointState(**state_dict)
-
-
-class DPRState:
- def __init__(self, src_file: Path):
- self.src_file = src_file
-
- def load_dpr_model(self):
- raise NotImplementedError
-
- @staticmethod
- def from_type(comp_type: str, *args, **kwargs) -> "DPRState":
- if comp_type.startswith("c"):
- return DPRContextEncoderState(*args, **kwargs)
- if comp_type.startswith("q"):
- return DPRQuestionEncoderState(*args, **kwargs)
- if comp_type.startswith("r"):
- return DPRReaderState(*args, **kwargs)
- else:
- raise ValueError("Component type must be either 'ctx_encoder', 'question_encoder' or 'reader'.")
-
-
-class DPRContextEncoderState(DPRState):
- def load_dpr_model(self):
- model = DPRContextEncoder(DPRConfig(**BertConfig.get_config_dict("google-bert/bert-base-uncased")[0]))
- print(f"Loading DPR biencoder from {self.src_file}")
- saved_state = load_states_from_checkpoint(self.src_file)
- encoder, prefix = model.ctx_encoder, "ctx_model."
- # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
- state_dict = {"bert_model.embeddings.position_ids": model.ctx_encoder.bert_model.embeddings.position_ids}
- for key, value in saved_state.model_dict.items():
- if key.startswith(prefix):
- key = key[len(prefix) :]
- if not key.startswith("encode_proj."):
- key = "bert_model." + key
- state_dict[key] = value
- encoder.load_state_dict(state_dict)
- return model
-
-
-class DPRQuestionEncoderState(DPRState):
- def load_dpr_model(self):
- model = DPRQuestionEncoder(DPRConfig(**BertConfig.get_config_dict("google-bert/bert-base-uncased")[0]))
- print(f"Loading DPR biencoder from {self.src_file}")
- saved_state = load_states_from_checkpoint(self.src_file)
- encoder, prefix = model.question_encoder, "question_model."
- # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
- state_dict = {"bert_model.embeddings.position_ids": model.question_encoder.bert_model.embeddings.position_ids}
- for key, value in saved_state.model_dict.items():
- if key.startswith(prefix):
- key = key[len(prefix) :]
- if not key.startswith("encode_proj."):
- key = "bert_model." + key
- state_dict[key] = value
- encoder.load_state_dict(state_dict)
- return model
-
-
-class DPRReaderState(DPRState):
- def load_dpr_model(self):
- model = DPRReader(DPRConfig(**BertConfig.get_config_dict("google-bert/bert-base-uncased")[0]))
- print(f"Loading DPR reader from {self.src_file}")
- saved_state = load_states_from_checkpoint(self.src_file)
- # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
- state_dict = {
- "encoder.bert_model.embeddings.position_ids": model.span_predictor.encoder.bert_model.embeddings.position_ids
- }
- for key, value in saved_state.model_dict.items():
- if key.startswith("encoder.") and not key.startswith("encoder.encode_proj"):
- key = "encoder.bert_model." + key[len("encoder.") :]
- state_dict[key] = value
- model.span_predictor.load_state_dict(state_dict)
- return model
-
-
-def convert(comp_type: str, src_file: Path, dest_dir: Path):
- dest_dir = Path(dest_dir)
- dest_dir.mkdir(exist_ok=True)
-
- dpr_state = DPRState.from_type(comp_type, src_file=src_file)
- model = dpr_state.load_dpr_model()
- model.save_pretrained(dest_dir)
- model.from_pretrained(dest_dir) # sanity check
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--type", type=str, help="Type of the component to convert: 'ctx_encoder', 'question_encoder' or 'reader'."
- )
- parser.add_argument(
- "--src",
- type=str,
- help=(
- "Path to the dpr checkpoint file. They can be downloaded from the official DPR repo"
- " https://github.com/facebookresearch/DPR. Note that in the official repo, both encoders are stored in the"
- " 'retriever' checkpoints."
- ),
- )
- parser.add_argument("--dest", type=str, default=None, help="Path to the output PyTorch model directory.")
- args = parser.parse_args()
-
- src_file = Path(args.src)
- dest_dir = f"converted-{src_file.name}" if args.dest is None else args.dest
- dest_dir = Path(dest_dir)
- assert src_file.exists()
- assert (
- args.type is not None
- ), "Please specify the component type of the DPR model to convert: 'ctx_encoder', 'question_encoder' or 'reader'."
- convert(args.type, src_file, dest_dir)
diff --git a/transformers/models/dpr/modeling_dpr.py b/transformers/models/dpr/modeling_dpr.py
deleted file mode 100644
index 0a45ec75207c29a31c32c9856345c6fb5df1e70d..0000000000000000000000000000000000000000
--- a/transformers/models/dpr/modeling_dpr.py
+++ /dev/null
@@ -1,663 +0,0 @@
-# coding=utf-8
-# Copyright 2018 DPR Authors, The Hugging Face Team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DPR model for Open Domain Question Answering."""
-
-
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-from torch import Tensor, nn
-
-from ...modeling_outputs import BaseModelOutputWithPooling
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from ..bert.modeling_bert import BertModel
-from .configuration_dpr import DPRConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "DPRConfig"
-_CHECKPOINT_FOR_DOC = "facebook/dpr-ctx_encoder-single-nq-base"
-
-
-from ..deprecated._archive_maps import ( # noqa: F401, E402
- DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST, # noqa: F401, E402
- DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST, # noqa: F401, E402
- DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST, # noqa: F401, E402
-)
-
-
-##########
-# Outputs
-##########
-
-
-@dataclass
-class DPRContextEncoderOutput(ModelOutput):
- """
- Class for outputs of [`DPRQuestionEncoder`].
-
- Args:
- pooler_output (`torch.FloatTensor` of shape `(batch_size, embeddings_size)`):
- The DPR encoder outputs the *pooler_output* that corresponds to the context representation. Last layer
- hidden-state of the first token of the sequence (classification token) further processed by a Linear layer.
- This output is to be used to embed contexts for nearest neighbors queries with questions embeddings.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- pooler_output: torch.FloatTensor
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-@dataclass
-class DPRQuestionEncoderOutput(ModelOutput):
- """
- Class for outputs of [`DPRQuestionEncoder`].
-
- Args:
- pooler_output (`torch.FloatTensor` of shape `(batch_size, embeddings_size)`):
- The DPR encoder outputs the *pooler_output* that corresponds to the question representation. Last layer
- hidden-state of the first token of the sequence (classification token) further processed by a Linear layer.
- This output is to be used to embed questions for nearest neighbors queries with context embeddings.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- pooler_output: torch.FloatTensor
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-@dataclass
-class DPRReaderOutput(ModelOutput):
- """
- Class for outputs of [`DPRQuestionEncoder`].
-
- Args:
- start_logits (`torch.FloatTensor` of shape `(n_passages, sequence_length)`):
- Logits of the start index of the span for each passage.
- end_logits (`torch.FloatTensor` of shape `(n_passages, sequence_length)`):
- Logits of the end index of the span for each passage.
- relevance_logits (`torch.FloatTensor` of shape `(n_passages, )`):
- Outputs of the QA classifier of the DPRReader that corresponds to the scores of each passage to answer the
- question, compared to all the other passages.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- start_logits: torch.FloatTensor
- end_logits: torch.FloatTensor = None
- relevance_logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-class DPRPreTrainedModel(PreTrainedModel):
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-class DPREncoder(DPRPreTrainedModel):
- base_model_prefix = "bert_model"
-
- def __init__(self, config: DPRConfig):
- super().__init__(config)
- self.bert_model = BertModel(config, add_pooling_layer=False)
- if self.bert_model.config.hidden_size <= 0:
- raise ValueError("Encoder hidden_size can't be zero")
- self.projection_dim = config.projection_dim
- if self.projection_dim > 0:
- self.encode_proj = nn.Linear(self.bert_model.config.hidden_size, config.projection_dim)
- # Initialize weights and apply final processing
- self.post_init()
-
- def forward(
- self,
- input_ids: Tensor,
- attention_mask: Optional[Tensor] = None,
- token_type_ids: Optional[Tensor] = None,
- inputs_embeds: Optional[Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = False,
- ) -> Union[BaseModelOutputWithPooling, Tuple[Tensor, ...]]:
- outputs = self.bert_model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = outputs[0]
- pooled_output = sequence_output[:, 0, :]
-
- if self.projection_dim > 0:
- pooled_output = self.encode_proj(pooled_output)
-
- if not return_dict:
- return (sequence_output, pooled_output) + outputs[2:]
-
- return BaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- @property
- def embeddings_size(self) -> int:
- if self.projection_dim > 0:
- return self.encode_proj.out_features
- return self.bert_model.config.hidden_size
-
-
-class DPRSpanPredictor(DPRPreTrainedModel):
- base_model_prefix = "encoder"
-
- def __init__(self, config: DPRConfig):
- super().__init__(config)
- self.encoder = DPREncoder(config)
- self.qa_outputs = nn.Linear(self.encoder.embeddings_size, 2)
- self.qa_classifier = nn.Linear(self.encoder.embeddings_size, 1)
- # Initialize weights and apply final processing
- self.post_init()
-
- def forward(
- self,
- input_ids: Tensor,
- attention_mask: Tensor,
- inputs_embeds: Optional[Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = False,
- ) -> Union[DPRReaderOutput, Tuple[Tensor, ...]]:
- # notations: N - number of questions in a batch, M - number of passages per questions, L - sequence length
- n_passages, sequence_length = input_ids.size() if input_ids is not None else inputs_embeds.size()[:2]
- # feed encoder
- outputs = self.encoder(
- input_ids,
- attention_mask=attention_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = outputs[0]
-
- # compute logits
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
- relevance_logits = self.qa_classifier(sequence_output[:, 0, :])
-
- # resize
- start_logits = start_logits.view(n_passages, sequence_length)
- end_logits = end_logits.view(n_passages, sequence_length)
- relevance_logits = relevance_logits.view(n_passages)
-
- if not return_dict:
- return (start_logits, end_logits, relevance_logits) + outputs[2:]
-
- return DPRReaderOutput(
- start_logits=start_logits,
- end_logits=end_logits,
- relevance_logits=relevance_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-##################
-# PreTrainedModel
-##################
-
-
-class DPRPretrainedContextEncoder(DPRPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DPRConfig
- load_tf_weights = None
- base_model_prefix = "ctx_encoder"
-
-
-class DPRPretrainedQuestionEncoder(DPRPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DPRConfig
- load_tf_weights = None
- base_model_prefix = "question_encoder"
-
-
-class DPRPretrainedReader(DPRPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DPRConfig
- load_tf_weights = None
- base_model_prefix = "span_predictor"
-
-
-###############
-# Actual Models
-###############
-
-
-DPR_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`DPRConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DPR_ENCODERS_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. To match pretraining, DPR input sequence should be
- formatted with [CLS] and [SEP] tokens as follows:
-
- (a) For sequence pairs (for a pair title+text for example):
-
- ```
- tokens: [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
- token_type_ids: 0 0 0 0 0 0 0 0 1 1 1 1 1 1
- ```
-
- (b) For single sequences (for a question for example):
-
- ```
- tokens: [CLS] the dog is hairy . [SEP]
- token_type_ids: 0 0 0 0 0 0 0
- ```
-
- DPR is a model with absolute position embeddings so it's usually advised to pad the inputs on the right
- rather than the left.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-DPR_READER_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Tuple[torch.LongTensor]` of shapes `(n_passages, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. It has to be a sequence triplet with 1) the question
- and 2) the passages titles and 3) the passages texts To match pretraining, DPR `input_ids` sequence should
- be formatted with [CLS] and [SEP] with the format:
-
- `[CLS] [SEP] [SEP] `
-
- DPR is a model with absolute position embeddings so it's usually advised to pad the inputs on the right
- rather than the left.
-
- Indices can be obtained using [`DPRReaderTokenizer`]. See this class documentation for more details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `(n_passages, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- inputs_embeds (`torch.FloatTensor` of shape `(n_passages, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DPRContextEncoder transformer outputting pooler outputs as context representations.",
- DPR_START_DOCSTRING,
-)
-class DPRContextEncoder(DPRPretrainedContextEncoder):
- def __init__(self, config: DPRConfig):
- super().__init__(config)
- self.config = config
- self.ctx_encoder = DPREncoder(config)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DPR_ENCODERS_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DPRContextEncoderOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[Tensor] = None,
- attention_mask: Optional[Tensor] = None,
- token_type_ids: Optional[Tensor] = None,
- inputs_embeds: Optional[Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[DPRContextEncoderOutput, Tuple[Tensor, ...]]:
- r"""
- Return:
-
- Examples:
-
- ```python
- >>> from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
-
- >>> tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
- >>> model = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
- >>> input_ids = tokenizer("Hello, is my dog cute ?", return_tensors="pt")["input_ids"]
- >>> embeddings = model(input_ids).pooler_output
- ```"""
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if attention_mask is None:
- attention_mask = (
- torch.ones(input_shape, device=device)
- if input_ids is None
- else (input_ids != self.config.pad_token_id)
- )
- if token_type_ids is None:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- outputs = self.ctx_encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- return outputs[1:]
- return DPRContextEncoderOutput(
- pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
-
-@add_start_docstrings(
- "The bare DPRQuestionEncoder transformer outputting pooler outputs as question representations.",
- DPR_START_DOCSTRING,
-)
-class DPRQuestionEncoder(DPRPretrainedQuestionEncoder):
- def __init__(self, config: DPRConfig):
- super().__init__(config)
- self.config = config
- self.question_encoder = DPREncoder(config)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DPR_ENCODERS_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DPRQuestionEncoderOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[Tensor] = None,
- attention_mask: Optional[Tensor] = None,
- token_type_ids: Optional[Tensor] = None,
- inputs_embeds: Optional[Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[DPRQuestionEncoderOutput, Tuple[Tensor, ...]]:
- r"""
- Return:
-
- Examples:
-
- ```python
- >>> from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer
-
- >>> tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
- >>> model = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
- >>> input_ids = tokenizer("Hello, is my dog cute ?", return_tensors="pt")["input_ids"]
- >>> embeddings = model(input_ids).pooler_output
- ```
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if attention_mask is None:
- attention_mask = (
- torch.ones(input_shape, device=device)
- if input_ids is None
- else (input_ids != self.config.pad_token_id)
- )
- if token_type_ids is None:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- outputs = self.question_encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- return outputs[1:]
- return DPRQuestionEncoderOutput(
- pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
-
-@add_start_docstrings(
- "The bare DPRReader transformer outputting span predictions.",
- DPR_START_DOCSTRING,
-)
-class DPRReader(DPRPretrainedReader):
- def __init__(self, config: DPRConfig):
- super().__init__(config)
- self.config = config
- self.span_predictor = DPRSpanPredictor(config)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DPR_READER_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DPRReaderOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[Tensor] = None,
- attention_mask: Optional[Tensor] = None,
- inputs_embeds: Optional[Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[DPRReaderOutput, Tuple[Tensor, ...]]:
- r"""
- Return:
-
- Examples:
-
- ```python
- >>> from transformers import DPRReader, DPRReaderTokenizer
-
- >>> tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base")
- >>> model = DPRReader.from_pretrained("facebook/dpr-reader-single-nq-base")
- >>> encoded_inputs = tokenizer(
- ... questions=["What is love ?"],
- ... titles=["Haddaway"],
- ... texts=["'What Is Love' is a song recorded by the artist Haddaway"],
- ... return_tensors="pt",
- ... )
- >>> outputs = model(**encoded_inputs)
- >>> start_logits = outputs.start_logits
- >>> end_logits = outputs.end_logits
- >>> relevance_logits = outputs.relevance_logits
- ```
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device)
-
- return self.span_predictor(
- input_ids,
- attention_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
diff --git a/transformers/models/dpr/modeling_tf_dpr.py b/transformers/models/dpr/modeling_tf_dpr.py
deleted file mode 100644
index e8cb1464f70da8bd5816fbde1f932c1d6e6fd084..0000000000000000000000000000000000000000
--- a/transformers/models/dpr/modeling_tf_dpr.py
+++ /dev/null
@@ -1,797 +0,0 @@
-# coding=utf-8
-# Copyright 2018 DPR Authors, The Hugging Face Team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-""" TensorFlow DPR model for Open Domain Question Answering."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Tuple, Union
-
-import tensorflow as tf
-
-from ...modeling_tf_outputs import TFBaseModelOutputWithPooling
-from ...modeling_tf_utils import TFModelInputType, TFPreTrainedModel, get_initializer, keras, shape_list, unpack_inputs
-from ...utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from ..bert.modeling_tf_bert import TFBertMainLayer
-from .configuration_dpr import DPRConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "DPRConfig"
-
-
-from ..deprecated._archive_maps import ( # noqa: F401, E402
- TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST, # noqa: F401, E402
- TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST, # noqa: F401, E402
- TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST, # noqa: F401, E402
-)
-
-
-##########
-# Outputs
-##########
-
-
-@dataclass
-class TFDPRContextEncoderOutput(ModelOutput):
- r"""
- Class for outputs of [`TFDPRContextEncoder`].
-
- Args:
- pooler_output (`tf.Tensor` of shape `(batch_size, embeddings_size)`):
- The DPR encoder outputs the *pooler_output* that corresponds to the context representation. Last layer
- hidden-state of the first token of the sequence (classification token) further processed by a Linear layer.
- This output is to be used to embed contexts for nearest neighbors queries with questions embeddings.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- pooler_output: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor, ...] | None = None
- attentions: Tuple[tf.Tensor, ...] | None = None
-
-
-@dataclass
-class TFDPRQuestionEncoderOutput(ModelOutput):
- """
- Class for outputs of [`TFDPRQuestionEncoder`].
-
- Args:
- pooler_output (`tf.Tensor` of shape `(batch_size, embeddings_size)`):
- The DPR encoder outputs the *pooler_output* that corresponds to the question representation. Last layer
- hidden-state of the first token of the sequence (classification token) further processed by a Linear layer.
- This output is to be used to embed questions for nearest neighbors queries with context embeddings.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- pooler_output: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor, ...] | None = None
- attentions: Tuple[tf.Tensor, ...] | None = None
-
-
-@dataclass
-class TFDPRReaderOutput(ModelOutput):
- """
- Class for outputs of [`TFDPRReaderEncoder`].
-
- Args:
- start_logits (`tf.Tensor` of shape `(n_passages, sequence_length)`):
- Logits of the start index of the span for each passage.
- end_logits (`tf.Tensor` of shape `(n_passages, sequence_length)`):
- Logits of the end index of the span for each passage.
- relevance_logits (`tf.Tensor` of shape `(n_passages, )`):
- Outputs of the QA classifier of the DPRReader that corresponds to the scores of each passage to answer the
- question, compared to all the other passages.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- start_logits: tf.Tensor = None
- end_logits: tf.Tensor = None
- relevance_logits: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor, ...] | None = None
- attentions: Tuple[tf.Tensor, ...] | None = None
-
-
-class TFDPREncoderLayer(keras.layers.Layer):
- base_model_prefix = "bert_model"
-
- def __init__(self, config: DPRConfig, **kwargs):
- super().__init__(**kwargs)
-
- # resolve name conflict with TFBertMainLayer instead of TFBertModel
- self.bert_model = TFBertMainLayer(config, add_pooling_layer=False, name="bert_model")
- self.config = config
-
- if self.config.hidden_size <= 0:
- raise ValueError("Encoder hidden_size can't be zero")
- self.projection_dim = config.projection_dim
- if self.projection_dim > 0:
- self.encode_proj = keras.layers.Dense(
- config.projection_dim, kernel_initializer=get_initializer(config.initializer_range), name="encode_proj"
- )
-
- @unpack_inputs
- def call(
- self,
- input_ids: tf.Tensor = None,
- attention_mask: tf.Tensor | None = None,
- token_type_ids: tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: bool = None,
- output_hidden_states: bool = None,
- return_dict: bool = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutputWithPooling, Tuple[tf.Tensor, ...]]:
- outputs = self.bert_model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = outputs[0]
- pooled_output = sequence_output[:, 0, :]
- if self.projection_dim > 0:
- pooled_output = self.encode_proj(pooled_output)
-
- if not return_dict:
- return (sequence_output, pooled_output) + outputs[1:]
-
- return TFBaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- @property
- def embeddings_size(self) -> int:
- if self.projection_dim > 0:
- return self.projection_dim
- return self.bert_model.config.hidden_size
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "bert_model", None) is not None:
- with tf.name_scope(self.bert_model.name):
- self.bert_model.build(None)
- if getattr(self, "encode_proj", None) is not None:
- with tf.name_scope(self.encode_proj.name):
- self.encode_proj.build(None)
-
-
-class TFDPRSpanPredictorLayer(keras.layers.Layer):
- base_model_prefix = "encoder"
-
- def __init__(self, config: DPRConfig, **kwargs):
- super().__init__(**kwargs)
- self.config = config
- self.encoder = TFDPREncoderLayer(config, name="encoder")
-
- self.qa_outputs = keras.layers.Dense(
- 2, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
- )
- self.qa_classifier = keras.layers.Dense(
- 1, kernel_initializer=get_initializer(config.initializer_range), name="qa_classifier"
- )
-
- @unpack_inputs
- def call(
- self,
- input_ids: tf.Tensor = None,
- attention_mask: tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = False,
- training: bool = False,
- ) -> Union[TFDPRReaderOutput, Tuple[tf.Tensor, ...]]:
- # notations: N - number of questions in a batch, M - number of passages per questions, L - sequence length
- n_passages, sequence_length = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)[:2]
- # feed encoder
- outputs = self.encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
-
- # compute logits
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = tf.split(logits, 2, axis=-1)
- start_logits = tf.squeeze(start_logits, axis=-1)
- end_logits = tf.squeeze(end_logits, axis=-1)
- relevance_logits = self.qa_classifier(sequence_output[:, 0, :])
-
- # resize
- start_logits = tf.reshape(start_logits, [n_passages, sequence_length])
- end_logits = tf.reshape(end_logits, [n_passages, sequence_length])
- relevance_logits = tf.reshape(relevance_logits, [n_passages])
-
- if not return_dict:
- return (start_logits, end_logits, relevance_logits) + outputs[2:]
-
- return TFDPRReaderOutput(
- start_logits=start_logits,
- end_logits=end_logits,
- relevance_logits=relevance_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.encoder.embeddings_size])
- if getattr(self, "qa_classifier", None) is not None:
- with tf.name_scope(self.qa_classifier.name):
- self.qa_classifier.build([None, None, self.encoder.embeddings_size])
-
-
-class TFDPRSpanPredictor(TFPreTrainedModel):
- base_model_prefix = "encoder"
-
- def __init__(self, config: DPRConfig, **kwargs):
- super().__init__(config, **kwargs)
- self.encoder = TFDPRSpanPredictorLayer(config)
-
- @unpack_inputs
- def call(
- self,
- input_ids: tf.Tensor = None,
- attention_mask: tf.Tensor | None = None,
- token_type_ids: tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = False,
- training: bool = False,
- ) -> Union[TFDPRReaderOutput, Tuple[tf.Tensor, ...]]:
- outputs = self.encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
-
-class TFDPREncoder(TFPreTrainedModel):
- base_model_prefix = "encoder"
-
- def __init__(self, config: DPRConfig, **kwargs):
- super().__init__(config, **kwargs)
-
- self.encoder = TFDPREncoderLayer(config)
-
- @unpack_inputs
- def call(
- self,
- input_ids: tf.Tensor = None,
- attention_mask: tf.Tensor | None = None,
- token_type_ids: tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = False,
- training: bool = False,
- ) -> Union[TFDPRReaderOutput, Tuple[tf.Tensor, ...]]:
- outputs = self.encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- return outputs
-
-
-##################
-# PreTrainedModel
-##################
-
-
-class TFDPRPretrainedContextEncoder(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DPRConfig
- base_model_prefix = "ctx_encoder"
-
-
-class TFDPRPretrainedQuestionEncoder(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DPRConfig
- base_model_prefix = "question_encoder"
-
-
-class TFDPRPretrainedReader(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DPRConfig
- base_model_prefix = "reader"
-
-
-###############
-# Actual Models
-###############
-
-
-TF_DPR_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a Tensorflow [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model)
- subclass. Use it as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to
- general usage and behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`DPRConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-TF_DPR_ENCODERS_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. To match pretraining, DPR input sequence should be
- formatted with [CLS] and [SEP] tokens as follows:
-
- (a) For sequence pairs (for a pair title+text for example):
-
- ```
- tokens: [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
- token_type_ids: 0 0 0 0 0 0 0 0 1 1 1 1 1 1
- ```
-
- (b) For single sequences (for a question for example):
-
- ```
- tokens: [CLS] the dog is hairy . [SEP]
- token_type_ids: 0 0 0 0 0 0 0
- ```
-
- DPR is a model with absolute position embeddings so it's usually advised to pad the inputs on the right
- rather than the left.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`Numpy array` or `tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`Numpy array` or `tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- inputs_embeds (`Numpy array` or `tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-TF_DPR_READER_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shapes `(n_passages, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. It has to be a sequence triplet with 1) the question
- and 2) the passages titles and 3) the passages texts To match pretraining, DPR `input_ids` sequence should
- be formatted with [CLS] and [SEP] with the format:
-
- `[CLS] [SEP] [SEP] `
-
- DPR is a model with absolute position embeddings so it's usually advised to pad the inputs on the right
- rather than the left.
-
- Indices can be obtained using [`DPRReaderTokenizer`]. See this class documentation for more details.
- attention_mask (`Numpy array` or `tf.Tensor` of shape `(n_passages, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- inputs_embeds (`Numpy array` or `tf.Tensor` of shape `(n_passages, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare DPRContextEncoder transformer outputting pooler outputs as context representations.",
- TF_DPR_START_DOCSTRING,
-)
-class TFDPRContextEncoder(TFDPRPretrainedContextEncoder):
- def __init__(self, config: DPRConfig, *args, **kwargs):
- super().__init__(config, *args, **kwargs)
- self.ctx_encoder = TFDPREncoderLayer(config, name="ctx_encoder")
-
- def get_input_embeddings(self):
- try:
- return self.ctx_encoder.bert_model.get_input_embeddings()
- except AttributeError:
- self.build()
- return self.ctx_encoder.bert_model.get_input_embeddings()
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(TF_DPR_ENCODERS_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFDPRContextEncoderOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: tf.Tensor | None = None,
- token_type_ids: tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: bool | None = None,
- output_hidden_states: bool | None = None,
- return_dict: bool | None = None,
- training: bool = False,
- ) -> TFDPRContextEncoderOutput | Tuple[tf.Tensor, ...]:
- r"""
- Return:
-
- Examples:
-
- ```python
- >>> from transformers import TFDPRContextEncoder, DPRContextEncoderTokenizer
-
- >>> tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
- >>> model = TFDPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", from_pt=True)
- >>> input_ids = tokenizer("Hello, is my dog cute ?", return_tensors="tf")["input_ids"]
- >>> embeddings = model(input_ids).pooler_output
- ```
- """
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = (
- tf.ones(input_shape, dtype=tf.dtypes.int32)
- if input_ids is None
- else (input_ids != self.config.pad_token_id)
- )
- if token_type_ids is None:
- token_type_ids = tf.zeros(input_shape, dtype=tf.dtypes.int32)
-
- outputs = self.ctx_encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- if not return_dict:
- return outputs[1:]
-
- return TFDPRContextEncoderOutput(
- pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "ctx_encoder", None) is not None:
- with tf.name_scope(self.ctx_encoder.name):
- self.ctx_encoder.build(None)
-
-
-@add_start_docstrings(
- "The bare DPRQuestionEncoder transformer outputting pooler outputs as question representations.",
- TF_DPR_START_DOCSTRING,
-)
-class TFDPRQuestionEncoder(TFDPRPretrainedQuestionEncoder):
- def __init__(self, config: DPRConfig, *args, **kwargs):
- super().__init__(config, *args, **kwargs)
- self.question_encoder = TFDPREncoderLayer(config, name="question_encoder")
-
- def get_input_embeddings(self):
- try:
- return self.question_encoder.bert_model.get_input_embeddings()
- except AttributeError:
- self.build()
- return self.question_encoder.bert_model.get_input_embeddings()
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(TF_DPR_ENCODERS_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFDPRQuestionEncoderOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: tf.Tensor | None = None,
- token_type_ids: tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: bool | None = None,
- output_hidden_states: bool | None = None,
- return_dict: bool | None = None,
- training: bool = False,
- ) -> TFDPRQuestionEncoderOutput | Tuple[tf.Tensor, ...]:
- r"""
- Return:
-
- Examples:
-
- ```python
- >>> from transformers import TFDPRQuestionEncoder, DPRQuestionEncoderTokenizer
-
- >>> tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
- >>> model = TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base", from_pt=True)
- >>> input_ids = tokenizer("Hello, is my dog cute ?", return_tensors="tf")["input_ids"]
- >>> embeddings = model(input_ids).pooler_output
- ```
- """
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = (
- tf.ones(input_shape, dtype=tf.dtypes.int32)
- if input_ids is None
- else (input_ids != self.config.pad_token_id)
- )
- if token_type_ids is None:
- token_type_ids = tf.zeros(input_shape, dtype=tf.dtypes.int32)
-
- outputs = self.question_encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- if not return_dict:
- return outputs[1:]
- return TFDPRQuestionEncoderOutput(
- pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "question_encoder", None) is not None:
- with tf.name_scope(self.question_encoder.name):
- self.question_encoder.build(None)
-
-
-@add_start_docstrings(
- "The bare DPRReader transformer outputting span predictions.",
- TF_DPR_START_DOCSTRING,
-)
-class TFDPRReader(TFDPRPretrainedReader):
- def __init__(self, config: DPRConfig, *args, **kwargs):
- super().__init__(config, *args, **kwargs)
- self.span_predictor = TFDPRSpanPredictorLayer(config, name="span_predictor")
-
- def get_input_embeddings(self):
- try:
- return self.span_predictor.encoder.bert_model.get_input_embeddings()
- except AttributeError:
- self.build()
- return self.span_predictor.encoder.bert_model.get_input_embeddings()
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(TF_DPR_READER_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFDPRReaderOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: bool | None = None,
- output_hidden_states: bool | None = None,
- return_dict: bool | None = None,
- training: bool = False,
- ) -> TFDPRReaderOutput | Tuple[tf.Tensor, ...]:
- r"""
- Return:
-
- Examples:
-
- ```python
- >>> from transformers import TFDPRReader, DPRReaderTokenizer
-
- >>> tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base")
- >>> model = TFDPRReader.from_pretrained("facebook/dpr-reader-single-nq-base", from_pt=True)
- >>> encoded_inputs = tokenizer(
- ... questions=["What is love ?"],
- ... titles=["Haddaway"],
- ... texts=["'What Is Love' is a song recorded by the artist Haddaway"],
- ... return_tensors="tf",
- ... )
- >>> outputs = model(encoded_inputs)
- >>> start_logits = outputs.start_logits
- >>> end_logits = outputs.end_logits
- >>> relevance_logits = outputs.relevance_logits
- ```
- """
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = tf.ones(input_shape, dtype=tf.dtypes.int32)
-
- return self.span_predictor(
- input_ids=input_ids,
- attention_mask=attention_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "span_predictor", None) is not None:
- with tf.name_scope(self.span_predictor.name):
- self.span_predictor.build(None)
diff --git a/transformers/models/dpr/tokenization_dpr.py b/transformers/models/dpr/tokenization_dpr.py
deleted file mode 100644
index 1362047ce283e2bacb8611115b80f4ad0bb6de61..0000000000000000000000000000000000000000
--- a/transformers/models/dpr/tokenization_dpr.py
+++ /dev/null
@@ -1,319 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team, The Hugging Face Team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for DPR."""
-
-
-import collections
-from typing import List, Optional, Union
-
-from ...tokenization_utils_base import BatchEncoding
-from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging
-from ..bert.tokenization_bert import BertTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
-
-
-class DPRContextEncoderTokenizer(BertTokenizer):
- r"""
- Construct a DPRContextEncoder tokenizer.
-
- [`DPRContextEncoderTokenizer`] is identical to [`BertTokenizer`] and runs end-to-end tokenization: punctuation
- splitting and wordpiece.
-
- Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
-
-
-class DPRQuestionEncoderTokenizer(BertTokenizer):
- r"""
- Constructs a DPRQuestionEncoder tokenizer.
-
- [`DPRQuestionEncoderTokenizer`] is identical to [`BertTokenizer`] and runs end-to-end tokenization: punctuation
- splitting and wordpiece.
-
- Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
-
-
-DPRSpanPrediction = collections.namedtuple(
- "DPRSpanPrediction", ["span_score", "relevance_score", "doc_id", "start_index", "end_index", "text"]
-)
-
-DPRReaderOutput = collections.namedtuple("DPRReaderOutput", ["start_logits", "end_logits", "relevance_logits"])
-
-
-CUSTOM_DPR_READER_DOCSTRING = r"""
- Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.
- It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),
- using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`
- with the format:
-
- ```
- [CLS] [SEP] [SEP]
- ```
-
- Args:
- questions (`str` or `List[str]`):
- The questions to be encoded. You can specify one question for many passages. In this case, the question
- will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in
- `titles` or `texts`.
- titles (`str` or `List[str]`):
- The passages titles to be encoded. This can be a string or a list of strings if there are several passages.
- texts (`str` or `List[str]`):
- The passages texts to be encoded. This can be a string or a list of strings if there are several passages.
- padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
- Activates and controls padding. Accepts the following values:
-
- - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence
- if provided).
- - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided.
- - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
- lengths).
- truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):
- Activates and controls truncation. Accepts the following values:
-
- - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to
- the maximum acceptable input length for the model if that argument is not provided. This will truncate
- token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch
- of pairs) is provided.
- - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided. This will only truncate the first
- sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
- - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided. This will only truncate the
- second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
- - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths
- greater than the model maximum admissible input size).
- max_length (`int`, *optional*):
- Controls the maximum length to use by one of the truncation/padding parameters.
-
- If left unset or set to `None`, this will use the predefined model maximum length if a maximum length
- is required by one of the truncation/padding parameters. If the model has no specific maximum input
- length (like XLNet) truncation/padding to a maximum length will be deactivated.
- return_tensors (`str` or [`~utils.TensorType`], *optional*):
- If set, will return tensors instead of list of python integers. Acceptable values are:
-
- - `'tf'`: Return TensorFlow `tf.constant` objects.
- - `'pt'`: Return PyTorch `torch.Tensor` objects.
- - `'np'`: Return Numpy `np.ndarray` objects.
- return_attention_mask (`bool`, *optional*):
- Whether or not to return the attention mask. If not set, will return the attention mask according to the
- specific tokenizer's default, defined by the `return_outputs` attribute.
-
- [What are attention masks?](../glossary#attention-mask)
-
- Returns:
- `Dict[str, List[List[int]]]`: A dictionary with the following keys:
-
- - `input_ids`: List of token ids to be fed to a model.
- - `attention_mask`: List of indices specifying which tokens should be attended to by the model.
- """
-
-
-@add_start_docstrings(CUSTOM_DPR_READER_DOCSTRING)
-class CustomDPRReaderTokenizerMixin:
- def __call__(
- self,
- questions,
- titles: Optional[str] = None,
- texts: Optional[str] = None,
- padding: Union[bool, str] = False,
- truncation: Union[bool, str] = False,
- max_length: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_attention_mask: Optional[bool] = None,
- **kwargs,
- ) -> BatchEncoding:
- if titles is None and texts is None:
- return super().__call__(
- questions,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- return_tensors=return_tensors,
- return_attention_mask=return_attention_mask,
- **kwargs,
- )
- elif titles is None or texts is None:
- text_pair = titles if texts is None else texts
- return super().__call__(
- questions,
- text_pair,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- return_tensors=return_tensors,
- return_attention_mask=return_attention_mask,
- **kwargs,
- )
- titles = titles if not isinstance(titles, str) else [titles]
- texts = texts if not isinstance(texts, str) else [texts]
- n_passages = len(titles)
- questions = questions if not isinstance(questions, str) else [questions] * n_passages
- if len(titles) != len(texts):
- raise ValueError(
- f"There should be as many titles than texts but got {len(titles)} titles and {len(texts)} texts."
- )
- encoded_question_and_titles = super().__call__(questions, titles, padding=False, truncation=False)["input_ids"]
- encoded_texts = super().__call__(texts, add_special_tokens=False, padding=False, truncation=False)["input_ids"]
- encoded_inputs = {
- "input_ids": [
- (encoded_question_and_title + encoded_text)[:max_length]
- if max_length is not None and truncation
- else encoded_question_and_title + encoded_text
- for encoded_question_and_title, encoded_text in zip(encoded_question_and_titles, encoded_texts)
- ]
- }
- if return_attention_mask is not False:
- attention_mask = []
- for input_ids in encoded_inputs["input_ids"]:
- attention_mask.append([int(input_id != self.pad_token_id) for input_id in input_ids])
- encoded_inputs["attention_mask"] = attention_mask
- return self.pad(encoded_inputs, padding=padding, max_length=max_length, return_tensors=return_tensors)
-
- def decode_best_spans(
- self,
- reader_input: BatchEncoding,
- reader_output: DPRReaderOutput,
- num_spans: int = 16,
- max_answer_length: int = 64,
- num_spans_per_passage: int = 4,
- ) -> List[DPRSpanPrediction]:
- """
- Get the span predictions for the extractive Q&A model.
-
- Returns: *List* of *DPRReaderOutput* sorted by descending *(relevance_score, span_score)*. Each
- *DPRReaderOutput* is a *Tuple* with:
-
- - **span_score**: `float` that corresponds to the score given by the reader for this span compared to other
- spans in the same passage. It corresponds to the sum of the start and end logits of the span.
- - **relevance_score**: `float` that corresponds to the score of the each passage to answer the question,
- compared to all the other passages. It corresponds to the output of the QA classifier of the DPRReader.
- - **doc_id**: `int` the id of the passage. - **start_index**: `int` the start index of the span
- (inclusive). - **end_index**: `int` the end index of the span (inclusive).
-
- Examples:
-
- ```python
- >>> from transformers import DPRReader, DPRReaderTokenizer
-
- >>> tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base")
- >>> model = DPRReader.from_pretrained("facebook/dpr-reader-single-nq-base")
- >>> encoded_inputs = tokenizer(
- ... questions=["What is love ?"],
- ... titles=["Haddaway"],
- ... texts=["'What Is Love' is a song recorded by the artist Haddaway"],
- ... return_tensors="pt",
- ... )
- >>> outputs = model(**encoded_inputs)
- >>> predicted_spans = tokenizer.decode_best_spans(encoded_inputs, outputs)
- >>> print(predicted_spans[0].text) # best span
- a song
- ```"""
- input_ids = reader_input["input_ids"]
- start_logits, end_logits, relevance_logits = reader_output[:3]
- n_passages = len(relevance_logits)
- sorted_docs = sorted(range(n_passages), reverse=True, key=relevance_logits.__getitem__)
- nbest_spans_predictions: List[DPRReaderOutput] = []
- for doc_id in sorted_docs:
- sequence_ids = list(input_ids[doc_id])
- # assuming question & title information is at the beginning of the sequence
- passage_offset = sequence_ids.index(self.sep_token_id, 2) + 1 # second sep id
- if sequence_ids[-1] == self.pad_token_id:
- sequence_len = sequence_ids.index(self.pad_token_id)
- else:
- sequence_len = len(sequence_ids)
-
- best_spans = self._get_best_spans(
- start_logits=start_logits[doc_id][passage_offset:sequence_len],
- end_logits=end_logits[doc_id][passage_offset:sequence_len],
- max_answer_length=max_answer_length,
- top_spans=num_spans_per_passage,
- )
- for start_index, end_index in best_spans:
- start_index += passage_offset
- end_index += passage_offset
- nbest_spans_predictions.append(
- DPRSpanPrediction(
- span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index],
- relevance_score=relevance_logits[doc_id],
- doc_id=doc_id,
- start_index=start_index,
- end_index=end_index,
- text=self.decode(sequence_ids[start_index : end_index + 1]),
- )
- )
- if len(nbest_spans_predictions) >= num_spans:
- break
- return nbest_spans_predictions[:num_spans]
-
- def _get_best_spans(
- self,
- start_logits: List[int],
- end_logits: List[int],
- max_answer_length: int,
- top_spans: int,
- ) -> List[DPRSpanPrediction]:
- """
- Finds the best answer span for the extractive Q&A model for one passage. It returns the best span by descending
- `span_score` order and keeping max `top_spans` spans. Spans longer that `max_answer_length` are ignored.
- """
- scores = []
- for start_index, start_score in enumerate(start_logits):
- for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]):
- scores.append(((start_index, start_index + answer_length), start_score + end_score))
- scores = sorted(scores, key=lambda x: x[1], reverse=True)
- chosen_span_intervals = []
- for (start_index, end_index), score in scores:
- if start_index > end_index:
- raise ValueError(f"Wrong span indices: [{start_index}:{end_index}]")
- length = end_index - start_index + 1
- if length > max_answer_length:
- raise ValueError(f"Span is too long: {length} > {max_answer_length}")
- if any(
- start_index <= prev_start_index <= prev_end_index <= end_index
- or prev_start_index <= start_index <= end_index <= prev_end_index
- for (prev_start_index, prev_end_index) in chosen_span_intervals
- ):
- continue
- chosen_span_intervals.append((start_index, end_index))
-
- if len(chosen_span_intervals) == top_spans:
- break
- return chosen_span_intervals
-
-
-@add_end_docstrings(CUSTOM_DPR_READER_DOCSTRING)
-class DPRReaderTokenizer(CustomDPRReaderTokenizerMixin, BertTokenizer):
- r"""
- Construct a DPRReader tokenizer.
-
- [`DPRReaderTokenizer`] is almost identical to [`BertTokenizer`] and runs end-to-end tokenization: punctuation
- splitting and wordpiece. The difference is that is has three inputs strings: question, titles and texts that are
- combined to be fed to the [`DPRReader`] model.
-
- Refer to superclass [`BertTokenizer`] for usage examples and documentation concerning parameters.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
diff --git a/transformers/models/dpr/tokenization_dpr_fast.py b/transformers/models/dpr/tokenization_dpr_fast.py
deleted file mode 100644
index 730f200a6876f11bbe30edcfd47b30861374b018..0000000000000000000000000000000000000000
--- a/transformers/models/dpr/tokenization_dpr_fast.py
+++ /dev/null
@@ -1,319 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team, The Hugging Face Team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for DPR."""
-
-
-import collections
-from typing import List, Optional, Union
-
-from ...tokenization_utils_base import BatchEncoding
-from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging
-from ..bert.tokenization_bert_fast import BertTokenizerFast
-from .tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer, DPRReaderTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
-
-
-class DPRContextEncoderTokenizerFast(BertTokenizerFast):
- r"""
- Construct a "fast" DPRContextEncoder tokenizer (backed by HuggingFace's *tokenizers* library).
-
- [`DPRContextEncoderTokenizerFast`] is identical to [`BertTokenizerFast`] and runs end-to-end tokenization:
- punctuation splitting and wordpiece.
-
- Refer to superclass [`BertTokenizerFast`] for usage examples and documentation concerning parameters.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = DPRContextEncoderTokenizer
-
-
-class DPRQuestionEncoderTokenizerFast(BertTokenizerFast):
- r"""
- Constructs a "fast" DPRQuestionEncoder tokenizer (backed by HuggingFace's *tokenizers* library).
-
- [`DPRQuestionEncoderTokenizerFast`] is identical to [`BertTokenizerFast`] and runs end-to-end tokenization:
- punctuation splitting and wordpiece.
-
- Refer to superclass [`BertTokenizerFast`] for usage examples and documentation concerning parameters.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = DPRQuestionEncoderTokenizer
-
-
-DPRSpanPrediction = collections.namedtuple(
- "DPRSpanPrediction", ["span_score", "relevance_score", "doc_id", "start_index", "end_index", "text"]
-)
-
-DPRReaderOutput = collections.namedtuple("DPRReaderOutput", ["start_logits", "end_logits", "relevance_logits"])
-
-
-CUSTOM_DPR_READER_DOCSTRING = r"""
- Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.
- It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),
- using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`
- with the format:
-
- [CLS] [SEP] [SEP]
-
- Args:
- questions (`str` or `List[str]`):
- The questions to be encoded. You can specify one question for many passages. In this case, the question
- will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in
- `titles` or `texts`.
- titles (`str` or `List[str]`):
- The passages titles to be encoded. This can be a string or a list of strings if there are several passages.
- texts (`str` or `List[str]`):
- The passages texts to be encoded. This can be a string or a list of strings if there are several passages.
- padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
- Activates and controls padding. Accepts the following values:
-
- - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence
- if provided).
- - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided.
- - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
- lengths).
- truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):
- Activates and controls truncation. Accepts the following values:
-
- - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to
- the maximum acceptable input length for the model if that argument is not provided. This will truncate
- token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch
- of pairs) is provided.
- - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided. This will only truncate the first
- sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
- - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided. This will only truncate the
- second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
- - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths
- greater than the model maximum admissible input size).
- max_length (`int`, *optional*):
- Controls the maximum length to use by one of the truncation/padding parameters.
-
- If left unset or set to `None`, this will use the predefined model maximum length if a maximum length
- is required by one of the truncation/padding parameters. If the model has no specific maximum input
- length (like XLNet) truncation/padding to a maximum length will be deactivated.
- return_tensors (`str` or [`~utils.TensorType`], *optional*):
- If set, will return tensors instead of list of python integers. Acceptable values are:
-
- - `'tf'`: Return TensorFlow `tf.constant` objects.
- - `'pt'`: Return PyTorch `torch.Tensor` objects.
- - `'np'`: Return Numpy `np.ndarray` objects.
- return_attention_mask (`bool`, *optional*):
- Whether or not to return the attention mask. If not set, will return the attention mask according to the
- specific tokenizer's default, defined by the `return_outputs` attribute.
-
- [What are attention masks?](../glossary#attention-mask)
-
- Return:
- `Dict[str, List[List[int]]]`: A dictionary with the following keys:
-
- - `input_ids`: List of token ids to be fed to a model.
- - `attention_mask`: List of indices specifying which tokens should be attended to by the model.
- """
-
-
-@add_start_docstrings(CUSTOM_DPR_READER_DOCSTRING)
-class CustomDPRReaderTokenizerMixin:
- def __call__(
- self,
- questions,
- titles: Optional[str] = None,
- texts: Optional[str] = None,
- padding: Union[bool, str] = False,
- truncation: Union[bool, str] = False,
- max_length: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- return_attention_mask: Optional[bool] = None,
- **kwargs,
- ) -> BatchEncoding:
- if titles is None and texts is None:
- return super().__call__(
- questions,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- return_tensors=return_tensors,
- return_attention_mask=return_attention_mask,
- **kwargs,
- )
- elif titles is None or texts is None:
- text_pair = titles if texts is None else texts
- return super().__call__(
- questions,
- text_pair,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- return_tensors=return_tensors,
- return_attention_mask=return_attention_mask,
- **kwargs,
- )
- titles = titles if not isinstance(titles, str) else [titles]
- texts = texts if not isinstance(texts, str) else [texts]
- n_passages = len(titles)
- questions = questions if not isinstance(questions, str) else [questions] * n_passages
- assert len(titles) == len(
- texts
- ), f"There should be as many titles than texts but got {len(titles)} titles and {len(texts)} texts."
- encoded_question_and_titles = super().__call__(questions, titles, padding=False, truncation=False)["input_ids"]
- encoded_texts = super().__call__(texts, add_special_tokens=False, padding=False, truncation=False)["input_ids"]
- encoded_inputs = {
- "input_ids": [
- (encoded_question_and_title + encoded_text)[:max_length]
- if max_length is not None and truncation
- else encoded_question_and_title + encoded_text
- for encoded_question_and_title, encoded_text in zip(encoded_question_and_titles, encoded_texts)
- ]
- }
- if return_attention_mask is not False:
- attention_mask = []
- for input_ids in encoded_inputs["input_ids"]:
- attention_mask.append([int(input_id != self.pad_token_id) for input_id in input_ids])
- encoded_inputs["attention_mask"] = attention_mask
- return self.pad(encoded_inputs, padding=padding, max_length=max_length, return_tensors=return_tensors)
-
- def decode_best_spans(
- self,
- reader_input: BatchEncoding,
- reader_output: DPRReaderOutput,
- num_spans: int = 16,
- max_answer_length: int = 64,
- num_spans_per_passage: int = 4,
- ) -> List[DPRSpanPrediction]:
- """
- Get the span predictions for the extractive Q&A model.
-
- Returns: *List* of *DPRReaderOutput* sorted by descending *(relevance_score, span_score)*. Each
- *DPRReaderOutput* is a *Tuple* with:
-
- - **span_score**: `float` that corresponds to the score given by the reader for this span compared to other
- spans in the same passage. It corresponds to the sum of the start and end logits of the span.
- - **relevance_score**: `float` that corresponds to the score of the each passage to answer the question,
- compared to all the other passages. It corresponds to the output of the QA classifier of the DPRReader.
- - **doc_id**: `int` the id of the passage. - ***start_index**: `int` the start index of the span
- (inclusive). - **end_index**: `int` the end index of the span (inclusive).
-
- Examples:
-
- ```python
- >>> from transformers import DPRReader, DPRReaderTokenizer
-
- >>> tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base")
- >>> model = DPRReader.from_pretrained("facebook/dpr-reader-single-nq-base")
- >>> encoded_inputs = tokenizer(
- ... questions=["What is love ?"],
- ... titles=["Haddaway"],
- ... texts=["'What Is Love' is a song recorded by the artist Haddaway"],
- ... return_tensors="pt",
- ... )
- >>> outputs = model(**encoded_inputs)
- >>> predicted_spans = tokenizer.decode_best_spans(encoded_inputs, outputs)
- >>> print(predicted_spans[0].text) # best span
- a song
- ```"""
- input_ids = reader_input["input_ids"]
- start_logits, end_logits, relevance_logits = reader_output[:3]
- n_passages = len(relevance_logits)
- sorted_docs = sorted(range(n_passages), reverse=True, key=relevance_logits.__getitem__)
- nbest_spans_predictions: List[DPRReaderOutput] = []
- for doc_id in sorted_docs:
- sequence_ids = list(input_ids[doc_id])
- # assuming question & title information is at the beginning of the sequence
- passage_offset = sequence_ids.index(self.sep_token_id, 2) + 1 # second sep id
- if sequence_ids[-1] == self.pad_token_id:
- sequence_len = sequence_ids.index(self.pad_token_id)
- else:
- sequence_len = len(sequence_ids)
-
- best_spans = self._get_best_spans(
- start_logits=start_logits[doc_id][passage_offset:sequence_len],
- end_logits=end_logits[doc_id][passage_offset:sequence_len],
- max_answer_length=max_answer_length,
- top_spans=num_spans_per_passage,
- )
- for start_index, end_index in best_spans:
- start_index += passage_offset
- end_index += passage_offset
- nbest_spans_predictions.append(
- DPRSpanPrediction(
- span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index],
- relevance_score=relevance_logits[doc_id],
- doc_id=doc_id,
- start_index=start_index,
- end_index=end_index,
- text=self.decode(sequence_ids[start_index : end_index + 1]),
- )
- )
- if len(nbest_spans_predictions) >= num_spans:
- break
- return nbest_spans_predictions[:num_spans]
-
- def _get_best_spans(
- self,
- start_logits: List[int],
- end_logits: List[int],
- max_answer_length: int,
- top_spans: int,
- ) -> List[DPRSpanPrediction]:
- """
- Finds the best answer span for the extractive Q&A model for one passage. It returns the best span by descending
- `span_score` order and keeping max `top_spans` spans. Spans longer that `max_answer_length` are ignored.
- """
- scores = []
- for start_index, start_score in enumerate(start_logits):
- for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]):
- scores.append(((start_index, start_index + answer_length), start_score + end_score))
- scores = sorted(scores, key=lambda x: x[1], reverse=True)
- chosen_span_intervals = []
- for (start_index, end_index), score in scores:
- assert start_index <= end_index, f"Wrong span indices: [{start_index}:{end_index}]"
- length = end_index - start_index + 1
- assert length <= max_answer_length, f"Span is too long: {length} > {max_answer_length}"
- if any(
- start_index <= prev_start_index <= prev_end_index <= end_index
- or prev_start_index <= start_index <= end_index <= prev_end_index
- for (prev_start_index, prev_end_index) in chosen_span_intervals
- ):
- continue
- chosen_span_intervals.append((start_index, end_index))
-
- if len(chosen_span_intervals) == top_spans:
- break
- return chosen_span_intervals
-
-
-@add_end_docstrings(CUSTOM_DPR_READER_DOCSTRING)
-class DPRReaderTokenizerFast(CustomDPRReaderTokenizerMixin, BertTokenizerFast):
- r"""
- Constructs a "fast" DPRReader tokenizer (backed by HuggingFace's *tokenizers* library).
-
- [`DPRReaderTokenizerFast`] is almost identical to [`BertTokenizerFast`] and runs end-to-end tokenization:
- punctuation splitting and wordpiece. The difference is that is has three inputs strings: question, titles and texts
- that are combined to be fed to the [`DPRReader`] model.
-
- Refer to superclass [`BertTokenizerFast`] for usage examples and documentation concerning parameters.
-
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
- slow_tokenizer_class = DPRReaderTokenizer
diff --git a/transformers/models/dpt/__init__.py b/transformers/models/dpt/__init__.py
deleted file mode 100644
index da53011b87b318bbef0d48557284d290f92a9fe4..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/__init__.py
+++ /dev/null
@@ -1,76 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...file_utils import _LazyModule, is_tokenizers_available, is_torch_available, is_vision_available
-from ...utils import OptionalDependencyNotAvailable
-
-
-_import_structure = {"configuration_dpt": ["DPT_PRETRAINED_CONFIG_ARCHIVE_MAP", "DPTConfig"]}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_dpt"] = ["DPTFeatureExtractor"]
- _import_structure["image_processing_dpt"] = ["DPTImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_dpt"] = [
- "DPT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "DPTForDepthEstimation",
- "DPTForSemanticSegmentation",
- "DPTModel",
- "DPTPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_dpt import DPT_PRETRAINED_CONFIG_ARCHIVE_MAP, DPTConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_dpt import DPTFeatureExtractor
- from .image_processing_dpt import DPTImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_dpt import (
- DPT_PRETRAINED_MODEL_ARCHIVE_LIST,
- DPTForDepthEstimation,
- DPTForSemanticSegmentation,
- DPTModel,
- DPTPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/dpt/__pycache__/__init__.cpython-310.pyc b/transformers/models/dpt/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index a7f0c083fb356a1e7e347355d953faff126db269..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/configuration_dpt.cpython-310.pyc b/transformers/models/dpt/__pycache__/configuration_dpt.cpython-310.pyc
deleted file mode 100644
index 9fef0f94ff33fc376e22613e44f31e00504415ab..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/configuration_dpt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/convert_dinov2_depth_to_hf.cpython-310.pyc b/transformers/models/dpt/__pycache__/convert_dinov2_depth_to_hf.cpython-310.pyc
deleted file mode 100644
index 49f0d7ef2b89cd94ce66b864bbb302123f77b0f7..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/convert_dinov2_depth_to_hf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/convert_dpt_beit_to_hf.cpython-310.pyc b/transformers/models/dpt/__pycache__/convert_dpt_beit_to_hf.cpython-310.pyc
deleted file mode 100644
index c8a031ed8e11bf99a7f67d6f6733692c31116ec3..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/convert_dpt_beit_to_hf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/convert_dpt_hybrid_to_pytorch.cpython-310.pyc b/transformers/models/dpt/__pycache__/convert_dpt_hybrid_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 1b8fe9810002a5125eccbf2ccbc0cdf43026a123..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/convert_dpt_hybrid_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/convert_dpt_swinv2_to_hf.cpython-310.pyc b/transformers/models/dpt/__pycache__/convert_dpt_swinv2_to_hf.cpython-310.pyc
deleted file mode 100644
index 143fde2bb7beae750549ac7fbe88cc53d799090e..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/convert_dpt_swinv2_to_hf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/convert_dpt_to_pytorch.cpython-310.pyc b/transformers/models/dpt/__pycache__/convert_dpt_to_pytorch.cpython-310.pyc
deleted file mode 100644
index f681bb66515949b949d24619abe3a8da173cdd2f..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/convert_dpt_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/feature_extraction_dpt.cpython-310.pyc b/transformers/models/dpt/__pycache__/feature_extraction_dpt.cpython-310.pyc
deleted file mode 100644
index fb3c673808ea9b0a6e4a5e1b16e71dce082fa4e3..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/feature_extraction_dpt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/image_processing_dpt.cpython-310.pyc b/transformers/models/dpt/__pycache__/image_processing_dpt.cpython-310.pyc
deleted file mode 100644
index 69b6588006deceaf7f3fa71f8e8384ef2b06465a..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/image_processing_dpt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/__pycache__/modeling_dpt.cpython-310.pyc b/transformers/models/dpt/__pycache__/modeling_dpt.cpython-310.pyc
deleted file mode 100644
index dd9300cba5e2bbccbe6c7c9ab62bc88371240664..0000000000000000000000000000000000000000
Binary files a/transformers/models/dpt/__pycache__/modeling_dpt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/dpt/configuration_dpt.py b/transformers/models/dpt/configuration_dpt.py
deleted file mode 100644
index 9bdc8d1ef0affbf6ea7cc952b174cf57f1237bb7..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/configuration_dpt.py
+++ /dev/null
@@ -1,284 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" DPT model configuration"""
-
-import copy
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-from ..auto.configuration_auto import CONFIG_MAPPING
-from ..bit import BitConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import DPT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DPTConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`DPTModel`]. It is used to instantiate an DPT
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the DPT
- [Intel/dpt-large](https://huggingface.co/Intel/dpt-large) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- image_size (`int`, *optional*, defaults to 384):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 16):
- The size (resolution) of each patch.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- is_hybrid (`bool`, *optional*, defaults to `False`):
- Whether to use a hybrid backbone. Useful in the context of loading DPT-Hybrid models.
- qkv_bias (`bool`, *optional*, defaults to `True`):
- Whether to add a bias to the queries, keys and values.
- backbone_out_indices (`List[int]`, *optional*, defaults to `[2, 5, 8, 11]`):
- Indices of the intermediate hidden states to use from backbone.
- readout_type (`str`, *optional*, defaults to `"project"`):
- The readout type to use when processing the readout token (CLS token) of the intermediate hidden states of
- the ViT backbone. Can be one of [`"ignore"`, `"add"`, `"project"`].
-
- - "ignore" simply ignores the CLS token.
- - "add" passes the information from the CLS token to all other tokens by adding the representations.
- - "project" passes information to the other tokens by concatenating the readout to all other tokens before
- projecting the
- representation to the original feature dimension D using a linear layer followed by a GELU non-linearity.
- reassemble_factors (`List[int]`, *optional*, defaults to `[4, 2, 1, 0.5]`):
- The up/downsampling factors of the reassemble layers.
- neck_hidden_sizes (`List[str]`, *optional*, defaults to `[96, 192, 384, 768]`):
- The hidden sizes to project to for the feature maps of the backbone.
- fusion_hidden_size (`int`, *optional*, defaults to 256):
- The number of channels before fusion.
- head_in_index (`int`, *optional*, defaults to -1):
- The index of the features to use in the heads.
- use_batch_norm_in_fusion_residual (`bool`, *optional*, defaults to `False`):
- Whether to use batch normalization in the pre-activate residual units of the fusion blocks.
- use_bias_in_fusion_residual (`bool`, *optional*, defaults to `True`):
- Whether to use bias in the pre-activate residual units of the fusion blocks.
- add_projection (`bool`, *optional*, defaults to `False`):
- Whether to add a projection layer before the depth estimation head.
- use_auxiliary_head (`bool`, *optional*, defaults to `True`):
- Whether to use an auxiliary head during training.
- auxiliary_loss_weight (`float`, *optional*, defaults to 0.4):
- Weight of the cross-entropy loss of the auxiliary head.
- semantic_loss_ignore_index (`int`, *optional*, defaults to 255):
- The index that is ignored by the loss function of the semantic segmentation model.
- semantic_classifier_dropout (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the semantic classification head.
- backbone_featmap_shape (`List[int]`, *optional*, defaults to `[1, 1024, 24, 24]`):
- Used only for the `hybrid` embedding type. The shape of the feature maps of the backbone.
- neck_ignore_stages (`List[int]`, *optional*, defaults to `[0, 1]`):
- Used only for the `hybrid` embedding type. The stages of the readout layers to ignore.
- backbone_config (`Union[Dict[str, Any], PretrainedConfig]`, *optional*):
- The configuration of the backbone model. Only used in case `is_hybrid` is `True` or in case you want to
- leverage the [`AutoBackbone`] API.
- backbone (`str`, *optional*):
- Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
- will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
- is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
- use_pretrained_backbone (`bool`, *optional*, defaults to `False`):
- Whether to use pretrained weights for the backbone.
- use_timm_backbone (`bool`, *optional*, defaults to `False`):
- Whether to load `backbone` from the timm library. If `False`, the backbone is loaded from the transformers
- library.
- backbone_kwargs (`dict`, *optional*):
- Keyword arguments to be passed to AutoBackbone when loading from a checkpoint
- e.g. `{'out_indices': (0, 1, 2, 3)}`. Cannot be specified if `backbone_config` is set.
-
- Example:
-
- ```python
- >>> from transformers import DPTModel, DPTConfig
-
- >>> # Initializing a DPT dpt-large style configuration
- >>> configuration = DPTConfig()
-
- >>> # Initializing a model from the dpt-large style configuration
- >>> model = DPTModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "dpt"
-
- def __init__(
- self,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.0,
- attention_probs_dropout_prob=0.0,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- image_size=384,
- patch_size=16,
- num_channels=3,
- is_hybrid=False,
- qkv_bias=True,
- backbone_out_indices=[2, 5, 8, 11],
- readout_type="project",
- reassemble_factors=[4, 2, 1, 0.5],
- neck_hidden_sizes=[96, 192, 384, 768],
- fusion_hidden_size=256,
- head_in_index=-1,
- use_batch_norm_in_fusion_residual=False,
- use_bias_in_fusion_residual=None,
- add_projection=False,
- use_auxiliary_head=True,
- auxiliary_loss_weight=0.4,
- semantic_loss_ignore_index=255,
- semantic_classifier_dropout=0.1,
- backbone_featmap_shape=[1, 1024, 24, 24],
- neck_ignore_stages=[0, 1],
- backbone_config=None,
- backbone=None,
- use_pretrained_backbone=False,
- use_timm_backbone=False,
- backbone_kwargs=None,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.is_hybrid = is_hybrid
-
- if use_pretrained_backbone:
- raise ValueError("Pretrained backbones are not supported yet.")
-
- use_autobackbone = False
- if self.is_hybrid:
- if backbone_config is None and backbone is None:
- logger.info("Initializing the config with a `BiT` backbone.")
- backbone_config = {
- "global_padding": "same",
- "layer_type": "bottleneck",
- "depths": [3, 4, 9],
- "out_features": ["stage1", "stage2", "stage3"],
- "embedding_dynamic_padding": True,
- }
- backbone_config = BitConfig(**backbone_config)
- elif isinstance(backbone_config, dict):
- logger.info("Initializing the config with a `BiT` backbone.")
- backbone_config = BitConfig(**backbone_config)
- elif isinstance(backbone_config, PretrainedConfig):
- backbone_config = backbone_config
- else:
- raise ValueError(
- f"backbone_config must be a dictionary or a `PretrainedConfig`, got {backbone_config.__class__}."
- )
- self.backbone_config = backbone_config
- self.backbone_featmap_shape = backbone_featmap_shape
- self.neck_ignore_stages = neck_ignore_stages
-
- if readout_type != "project":
- raise ValueError("Readout type must be 'project' when using `DPT-hybrid` mode.")
-
- elif backbone_config is not None:
- use_autobackbone = True
-
- if isinstance(backbone_config, dict):
- backbone_model_type = backbone_config.get("model_type")
- config_class = CONFIG_MAPPING[backbone_model_type]
- backbone_config = config_class.from_dict(backbone_config)
-
- self.backbone_config = backbone_config
- self.backbone_featmap_shape = None
- self.neck_ignore_stages = []
- else:
- self.backbone_config = backbone_config
- self.backbone_featmap_shape = None
- self.neck_ignore_stages = []
-
- if use_autobackbone and backbone_config is not None and backbone is not None:
- raise ValueError("You can't specify both `backbone` and `backbone_config`.")
-
- if backbone_kwargs is not None and backbone_kwargs and backbone_config is not None:
- raise ValueError("You can't specify both `backbone_kwargs` and `backbone_config`.")
-
- self.backbone = backbone
- self.use_pretrained_backbone = use_pretrained_backbone
- self.use_timm_backbone = use_timm_backbone
- self.backbone_kwargs = backbone_kwargs
- self.num_hidden_layers = None if use_autobackbone else num_hidden_layers
- self.num_attention_heads = None if use_autobackbone else num_attention_heads
- self.intermediate_size = None if use_autobackbone else intermediate_size
- self.hidden_dropout_prob = None if use_autobackbone else hidden_dropout_prob
- self.attention_probs_dropout_prob = None if use_autobackbone else attention_probs_dropout_prob
- self.layer_norm_eps = None if use_autobackbone else layer_norm_eps
- self.image_size = None if use_autobackbone else image_size
- self.patch_size = None if use_autobackbone else patch_size
- self.num_channels = None if use_autobackbone else num_channels
- self.qkv_bias = None if use_autobackbone else qkv_bias
- self.backbone_out_indices = None if use_autobackbone else backbone_out_indices
-
- if readout_type not in ["ignore", "add", "project"]:
- raise ValueError("Readout_type must be one of ['ignore', 'add', 'project']")
- self.hidden_act = hidden_act
- self.initializer_range = initializer_range
- self.readout_type = readout_type
- self.reassemble_factors = reassemble_factors
- self.neck_hidden_sizes = neck_hidden_sizes
- self.fusion_hidden_size = fusion_hidden_size
- self.head_in_index = head_in_index
- self.use_batch_norm_in_fusion_residual = use_batch_norm_in_fusion_residual
- self.use_bias_in_fusion_residual = use_bias_in_fusion_residual
- self.add_projection = add_projection
-
- # auxiliary head attributes (semantic segmentation)
- self.use_auxiliary_head = use_auxiliary_head
- self.auxiliary_loss_weight = auxiliary_loss_weight
- self.semantic_loss_ignore_index = semantic_loss_ignore_index
- self.semantic_classifier_dropout = semantic_classifier_dropout
-
- def to_dict(self):
- """
- Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. Returns:
- `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
- """
- output = copy.deepcopy(self.__dict__)
-
- if output["backbone_config"] is not None:
- output["backbone_config"] = self.backbone_config.to_dict()
-
- output["model_type"] = self.__class__.model_type
- return output
diff --git a/transformers/models/dpt/convert_dinov2_depth_to_hf.py b/transformers/models/dpt/convert_dinov2_depth_to_hf.py
deleted file mode 100644
index 2bd147096c86aebf1f2cedb3b412deaed35da15a..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/convert_dinov2_depth_to_hf.py
+++ /dev/null
@@ -1,384 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DINOv2 + DPT checkpoints from the original repository. URL:
-https://github.com/facebookresearch/dinov2/tree/main"""
-
-
-import argparse
-import itertools
-import math
-from pathlib import Path
-
-import requests
-import torch
-from PIL import Image
-from torchvision import transforms
-
-from transformers import Dinov2Config, DPTConfig, DPTForDepthEstimation, DPTImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_dpt_config(model_name):
- if "small" in model_name:
- # equivalent to stage 3, stage 6, stage 9, stage 12
- backbone_config = Dinov2Config.from_pretrained(
- "facebook/dinov2-small", out_indices=[3, 6, 9, 12], apply_layernorm=False, reshape_hidden_states=False
- )
- neck_hidden_sizes = [48, 96, 192, 384]
- elif "base" in model_name:
- backbone_config = Dinov2Config.from_pretrained(
- "facebook/dinov2-base", out_indices=[3, 6, 9, 12], apply_layernorm=False, reshape_hidden_states=False
- )
- neck_hidden_sizes = [96, 192, 384, 768]
- elif "large" in model_name:
- backbone_config = Dinov2Config.from_pretrained(
- "facebook/dinov2-large", out_indices=[5, 12, 18, 24], apply_layernorm=False, reshape_hidden_states=False
- )
- neck_hidden_sizes = [128, 256, 512, 1024]
- elif "giant" in model_name:
- backbone_config = Dinov2Config.from_pretrained(
- "facebook/dinov2-giant", out_indices=[10, 20, 30, 40], apply_layernorm=False, reshape_hidden_states=False
- )
- neck_hidden_sizes = [192, 384, 768, 1536]
- else:
- raise NotImplementedError("To do")
-
- config = DPTConfig(
- backbone_config=backbone_config,
- neck_hidden_sizes=neck_hidden_sizes,
- use_bias_in_fusion_residual=False,
- add_projection=True,
- )
-
- return config
-
-
-# here we list all DPT keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys_dpt(config):
- rename_keys = []
-
- # fmt: off
- # activation postprocessing (projections, readout projections + resize blocks)
- for i in range(4):
- rename_keys.append((f"decode_head.reassemble_blocks.projects.{i}.conv.weight", f"neck.reassemble_stage.layers.{i}.projection.weight"))
- rename_keys.append((f"decode_head.reassemble_blocks.projects.{i}.conv.bias", f"neck.reassemble_stage.layers.{i}.projection.bias"))
-
- rename_keys.append((f"decode_head.reassemble_blocks.readout_projects.{i}.0.weight", f"neck.reassemble_stage.readout_projects.{i}.0.weight"))
- rename_keys.append((f"decode_head.reassemble_blocks.readout_projects.{i}.0.bias", f"neck.reassemble_stage.readout_projects.{i}.0.bias"))
-
- if i != 2:
- rename_keys.append((f"decode_head.reassemble_blocks.resize_layers.{i}.weight", f"neck.reassemble_stage.layers.{i}.resize.weight"))
- rename_keys.append((f"decode_head.reassemble_blocks.resize_layers.{i}.bias", f"neck.reassemble_stage.layers.{i}.resize.bias"))
-
- # fusion layers
- for i in range(4):
- rename_keys.append((f"decode_head.fusion_blocks.{i}.project.conv.weight", f"neck.fusion_stage.layers.{i}.projection.weight"))
- rename_keys.append((f"decode_head.fusion_blocks.{i}.project.conv.bias", f"neck.fusion_stage.layers.{i}.projection.bias"))
- if i != 0:
- rename_keys.append((f"decode_head.fusion_blocks.{i}.res_conv_unit1.conv1.conv.weight", f"neck.fusion_stage.layers.{i}.residual_layer1.convolution1.weight"))
- rename_keys.append((f"decode_head.fusion_blocks.{i}.res_conv_unit1.conv2.conv.weight", f"neck.fusion_stage.layers.{i}.residual_layer1.convolution2.weight"))
- rename_keys.append((f"decode_head.fusion_blocks.{i}.res_conv_unit2.conv1.conv.weight", f"neck.fusion_stage.layers.{i}.residual_layer2.convolution1.weight"))
- rename_keys.append((f"decode_head.fusion_blocks.{i}.res_conv_unit2.conv2.conv.weight", f"neck.fusion_stage.layers.{i}.residual_layer2.convolution2.weight"))
-
- # neck convolutions
- for i in range(4):
- rename_keys.append((f"decode_head.convs.{i}.conv.weight", f"neck.convs.{i}.weight"))
-
- # head
- rename_keys.append(("decode_head.project.conv.weight", "head.projection.weight"))
- rename_keys.append(("decode_head.project.conv.bias", "head.projection.bias"))
-
- for i in range(0, 5, 2):
- rename_keys.append((f"decode_head.conv_depth.head.{i}.weight", f"head.head.{i}.weight"))
- rename_keys.append((f"decode_head.conv_depth.head.{i}.bias", f"head.head.{i}.bias"))
- # fmt: on
-
- return rename_keys
-
-
-# here we list all backbone keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys_backbone(config):
- rename_keys = []
-
- # fmt: off
- # patch embedding layer
- rename_keys.append(("cls_token", "backbone.embeddings.cls_token"))
- rename_keys.append(("mask_token", "backbone.embeddings.mask_token"))
- rename_keys.append(("pos_embed", "backbone.embeddings.position_embeddings"))
- rename_keys.append(("patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight"))
- rename_keys.append(("patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias"))
-
- # Transfomer encoder
- for i in range(config.backbone_config.num_hidden_layers):
- # layernorms
- rename_keys.append((f"blocks.{i}.norm1.weight", f"backbone.encoder.layer.{i}.norm1.weight"))
- rename_keys.append((f"blocks.{i}.norm1.bias", f"backbone.encoder.layer.{i}.norm1.bias"))
- rename_keys.append((f"blocks.{i}.norm2.weight", f"backbone.encoder.layer.{i}.norm2.weight"))
- rename_keys.append((f"blocks.{i}.norm2.bias", f"backbone.encoder.layer.{i}.norm2.bias"))
- # MLP
- if config.backbone_config.use_swiglu_ffn:
- rename_keys.append((f"blocks.{i}.mlp.w12.weight", f"backbone.encoder.layer.{i}.mlp.w12.weight"))
- rename_keys.append((f"blocks.{i}.mlp.w12.bias", f"backbone.encoder.layer.{i}.mlp.w12.bias"))
- rename_keys.append((f"blocks.{i}.mlp.w3.weight", f"backbone.encoder.layer.{i}.mlp.w3.weight"))
- rename_keys.append((f"blocks.{i}.mlp.w3.bias", f"backbone.encoder.layer.{i}.mlp.w3.bias"))
- else:
- rename_keys.append((f"blocks.{i}.mlp.fc1.weight", f"backbone.encoder.layer.{i}.mlp.fc1.weight"))
- rename_keys.append((f"blocks.{i}.mlp.fc1.bias", f"backbone.encoder.layer.{i}.mlp.fc1.bias"))
- rename_keys.append((f"blocks.{i}.mlp.fc2.weight", f"backbone.encoder.layer.{i}.mlp.fc2.weight"))
- rename_keys.append((f"blocks.{i}.mlp.fc2.bias", f"backbone.encoder.layer.{i}.mlp.fc2.bias"))
- # layerscale
- rename_keys.append((f"blocks.{i}.ls1.gamma", f"backbone.encoder.layer.{i}.layer_scale1.lambda1"))
- rename_keys.append((f"blocks.{i}.ls2.gamma", f"backbone.encoder.layer.{i}.layer_scale2.lambda1"))
- # attention projection layer
- rename_keys.append((f"blocks.{i}.attn.proj.weight", f"backbone.encoder.layer.{i}.attention.output.dense.weight"))
- rename_keys.append((f"blocks.{i}.attn.proj.bias", f"backbone.encoder.layer.{i}.attention.output.dense.bias"))
- # fmt: on
-
- rename_keys.append(("norm.weight", "backbone.layernorm.weight"))
- rename_keys.append(("norm.bias", "backbone.layernorm.bias"))
-
- return rename_keys
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config):
- for i in range(config.backbone_config.num_hidden_layers):
- # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"blocks.{i}.attn.qkv.weight")
- in_proj_bias = state_dict.pop(f"blocks.{i}.attn.qkv.bias")
- hidden_size = config.backbone_config.hidden_size
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[:hidden_size, :]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[:hidden_size]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- hidden_size : hidden_size * 2, :
- ]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
- hidden_size : hidden_size * 2
- ]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-hidden_size:, :]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-hidden_size:]
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "https://dl.fbaipublicfiles.com/dinov2/images/example.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-name_to_url = {
- "dpt-dinov2-small-nyu": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_nyu_dpt_head.pth",
- "dpt-dinov2-small-kitti": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_kitti_dpt_head.pth",
- "dpt-dinov2-base-nyu": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_nyu_dpt_head.pth",
- "dpt-dinov2-base-kitti": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_kitti_dpt_head.pth",
- "dpt-dinov2-large-nyu": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_nyu_dpt_head.pth",
- "dpt-dinov2-large-kitti": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_kitti_dpt_head.pth",
- "dpt-dinov2-giant-nyu": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_nyu_dpt_head.pth",
- "dpt-dinov2-giant-kitti": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitg14/dinov2_vitg14_kitti_dpt_head.pth",
-}
-
-
-def get_original_pixel_values(image):
- class CenterPadding(object):
- def __init__(self, multiple):
- super().__init__()
- self.multiple = multiple
-
- def _get_pad(self, size):
- new_size = math.ceil(size / self.multiple) * self.multiple
- pad_size = new_size - size
- pad_size_left = pad_size // 2
- pad_size_right = pad_size - pad_size_left
- return pad_size_left, pad_size_right
-
- def __call__(self, img):
- pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in img.shape[-2:][::-1]))
- output = torch.nn.functional.pad(img, pads)
- return output
-
- def __repr__(self):
- return self.__class__.__name__ + "()"
-
- def make_depth_transform() -> transforms.Compose:
- return transforms.Compose(
- [
- transforms.ToTensor(),
- lambda x: 255.0 * x[:3], # Discard alpha component and scale by 255
- transforms.Normalize(
- mean=(123.675, 116.28, 103.53),
- std=(58.395, 57.12, 57.375),
- ),
- CenterPadding(multiple=14),
- ]
- )
-
- transform = make_depth_transform()
- original_pixel_values = transform(image).unsqueeze(0)
-
- return original_pixel_values
-
-
-@torch.no_grad()
-def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub, verify_logits):
- """
- Copy/paste/tweak model's weights to our DPT structure.
- """
-
- # define DPT configuration based on URL
- checkpoint_url = name_to_url[model_name]
- config = get_dpt_config(model_name)
-
- # load original DPT state_dict from URL
- print("URL:", checkpoint_url)
- dpt_state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["state_dict"]
- # rename keys
- rename_keys = create_rename_keys_dpt(config)
- for src, dest in rename_keys:
- rename_key(dpt_state_dict, src, dest)
-
- # load original backbone state_dict from URL
- if "small" in model_name:
- original_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
- elif "base" in model_name:
- original_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitb14")
- elif "large" in model_name:
- original_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14")
- elif "giant" in model_name:
- original_model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitg14")
- else:
- raise NotImplementedError("To do")
- original_model.eval()
- backbone_state_dict = original_model.state_dict()
-
- # rename keys
- rename_keys = create_rename_keys_backbone(config)
- for src, dest in rename_keys:
- rename_key(backbone_state_dict, src, dest)
-
- # read in qkv matrices
- read_in_q_k_v(backbone_state_dict, config)
-
- for key, val in backbone_state_dict.copy().items():
- val = backbone_state_dict.pop(key)
- if "w12" in key:
- key = key.replace("w12", "weights_in")
- if "w3" in key:
- key = key.replace("w3", "weights_out")
- backbone_state_dict[key] = val
-
- # merge state_dicts
- state_dict = {**backbone_state_dict, **dpt_state_dict}
-
- # load HuggingFace model
- model = DPTForDepthEstimation(config)
- missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
- print("Missing keys:", missing_keys)
- print("Unexpected keys:", unexpected_keys)
- assert missing_keys == [
- "neck.fusion_stage.layers.0.residual_layer1.convolution1.weight",
- "neck.fusion_stage.layers.0.residual_layer1.convolution2.weight",
- ]
- model.eval()
-
- # Verify image processor
- processor = DPTImageProcessor(
- do_resize=False,
- do_rescale=False,
- do_pad=True,
- size_divisor=14,
- do_normalize=True,
- image_mean=(123.675, 116.28, 103.53),
- image_std=(58.395, 57.12, 57.375),
- )
-
- image = prepare_img()
- pixel_values = processor(image, return_tensors="pt").pixel_values.float()
- original_pixel_values = get_original_pixel_values(image)
-
- assert torch.allclose(pixel_values, original_pixel_values)
-
- # Verify forward pass
- with torch.no_grad():
- outputs = model(pixel_values)
-
- predicted_depth = outputs.predicted_depth
-
- print("Shape of predicted depth:", predicted_depth.shape)
- print("First values of predicted depth:", predicted_depth[0, :3, :3])
-
- # assert logits
- if verify_logits:
- if model_name == "dpt-dinov2-small-nyu":
- expected_shape = torch.Size([1, 576, 736])
- expected_slice = torch.tensor(
- [[3.3576, 3.4741, 3.4345], [3.4324, 3.5012, 3.2775], [3.2560, 3.3563, 3.2354]]
- )
-
- assert predicted_depth.shape == torch.Size(expected_shape)
- assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-5)
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model and processor to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- print("Pushing model and processor to hub...")
- model.push_to_hub(repo_id=f"facebook/{model_name}")
- processor.push_to_hub(repo_id=f"facebook/{model_name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="dpt-dinov2-small-nyu",
- type=str,
- choices=name_to_url.keys(),
- help="Name of the model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- help="Whether to push the model to the hub after conversion.",
- )
- parser.add_argument(
- "--verify_logits",
- action="store_true",
- required=False,
- help="Path to the output PyTorch model directory.",
- )
-
- args = parser.parse_args()
- convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.verify_logits)
diff --git a/transformers/models/dpt/convert_dpt_beit_to_hf.py b/transformers/models/dpt/convert_dpt_beit_to_hf.py
deleted file mode 100644
index eb335a0ea2aeeb82dfa43287b5697c755ce1d5ff..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/convert_dpt_beit_to_hf.py
+++ /dev/null
@@ -1,306 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DPT 3.1 checkpoints from the MiDaS repository. URL: https://github.com/isl-org/MiDaS"""
-
-
-import argparse
-from pathlib import Path
-
-import requests
-import torch
-from PIL import Image
-
-from transformers import BeitConfig, DPTConfig, DPTForDepthEstimation, DPTImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_dpt_config(model_name):
- hidden_size = 768
- num_hidden_layers = 12
- num_attention_heads = 12
- intermediate_size = 3072
- out_features = ["stage3", "stage6", "stage9", "stage12"] # beit-base-384 uses [2, 5, 8, 11]
-
- if "large" in model_name:
- hidden_size = 1024
- num_hidden_layers = 24
- num_attention_heads = 16
- intermediate_size = 4096
- out_features = ["stage6", "stage12", "stage18", "stage24"] # beit-large-512 uses [5, 11, 17, 23]
-
- if "512" in model_name:
- image_size = 512
- elif "384" in model_name:
- image_size = 384
- else:
- raise ValueError("Model not supported")
-
- backbone_config = BeitConfig(
- image_size=image_size,
- num_hidden_layers=num_hidden_layers,
- hidden_size=hidden_size,
- intermediate_size=intermediate_size,
- num_attention_heads=num_attention_heads,
- use_relative_position_bias=True,
- reshape_hidden_states=False,
- out_features=out_features,
- )
-
- neck_hidden_sizes = [256, 512, 1024, 1024] if "large" in model_name else [96, 192, 384, 768]
- config = DPTConfig(backbone_config=backbone_config, neck_hidden_sizes=neck_hidden_sizes)
-
- return config, image_size
-
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys(config):
- rename_keys = []
-
- # fmt: off
- # stem
- rename_keys.append(("pretrained.model.cls_token", "backbone.embeddings.cls_token"))
- rename_keys.append(("pretrained.model.patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight"))
- rename_keys.append(("pretrained.model.patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias"))
-
- # Transfomer encoder
- for i in range(config.backbone_config.num_hidden_layers):
- rename_keys.append((f"pretrained.model.blocks.{i}.gamma_1", f"backbone.encoder.layer.{i}.lambda_1"))
- rename_keys.append((f"pretrained.model.blocks.{i}.gamma_2", f"backbone.encoder.layer.{i}.lambda_2"))
- rename_keys.append((f"pretrained.model.blocks.{i}.norm1.weight", f"backbone.encoder.layer.{i}.layernorm_before.weight"))
- rename_keys.append((f"pretrained.model.blocks.{i}.norm1.bias", f"backbone.encoder.layer.{i}.layernorm_before.bias"))
- rename_keys.append((f"pretrained.model.blocks.{i}.norm2.weight", f"backbone.encoder.layer.{i}.layernorm_after.weight"))
- rename_keys.append((f"pretrained.model.blocks.{i}.norm2.bias", f"backbone.encoder.layer.{i}.layernorm_after.bias"))
- rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc1.weight", f"backbone.encoder.layer.{i}.intermediate.dense.weight"))
- rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc1.bias", f"backbone.encoder.layer.{i}.intermediate.dense.bias"))
- rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc2.weight", f"backbone.encoder.layer.{i}.output.dense.weight"))
- rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc2.bias", f"backbone.encoder.layer.{i}.output.dense.bias"))
- rename_keys.append((f"pretrained.model.blocks.{i}.attn.proj.weight", f"backbone.encoder.layer.{i}.attention.output.dense.weight"))
- rename_keys.append((f"pretrained.model.blocks.{i}.attn.proj.bias", f"backbone.encoder.layer.{i}.attention.output.dense.bias"))
- rename_keys.append((f"pretrained.model.blocks.{i}.attn.relative_position_bias_table", f"backbone.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_bias_table"))
- rename_keys.append((f"pretrained.model.blocks.{i}.attn.relative_position_index", f"backbone.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_index"))
-
- # activation postprocessing (readout projections + resize blocks)
- for i in range(4):
- rename_keys.append((f"pretrained.act_postprocess{i+1}.0.project.0.weight", f"neck.reassemble_stage.readout_projects.{i}.0.weight"))
- rename_keys.append((f"pretrained.act_postprocess{i+1}.0.project.0.bias", f"neck.reassemble_stage.readout_projects.{i}.0.bias"))
-
- rename_keys.append((f"pretrained.act_postprocess{i+1}.3.weight", f"neck.reassemble_stage.layers.{i}.projection.weight"))
- rename_keys.append((f"pretrained.act_postprocess{i+1}.3.bias", f"neck.reassemble_stage.layers.{i}.projection.bias"))
-
- if i != 2:
- rename_keys.append((f"pretrained.act_postprocess{i+1}.4.weight", f"neck.reassemble_stage.layers.{i}.resize.weight"))
- rename_keys.append((f"pretrained.act_postprocess{i+1}.4.bias", f"neck.reassemble_stage.layers.{i}.resize.bias"))
-
- # refinenet (tricky here)
- mapping = {1:3, 2:2, 3:1, 4:0}
-
- for i in range(1, 5):
- j = mapping[i]
- rename_keys.append((f"scratch.refinenet{i}.out_conv.weight", f"neck.fusion_stage.layers.{j}.projection.weight"))
- rename_keys.append((f"scratch.refinenet{i}.out_conv.bias", f"neck.fusion_stage.layers.{j}.projection.bias"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.weight"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.bias"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.weight"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.bias"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.weight"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.bias"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.weight"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.bias"))
-
- # scratch convolutions
- for i in range(4):
- rename_keys.append((f"scratch.layer{i+1}_rn.weight", f"neck.convs.{i}.weight"))
-
- # head
- for i in range(0, 5, 2):
- rename_keys.append((f"scratch.output_conv.{i}.weight", f"head.head.{i}.weight"))
- rename_keys.append((f"scratch.output_conv.{i}.bias", f"head.head.{i}.bias"))
-
- return rename_keys
-
-
-def remove_ignore_keys_(state_dict):
- ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"]
- for k in ignore_keys:
- state_dict.pop(k, None)
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config):
- hidden_size = config.backbone_config.hidden_size
- for i in range(config.backbone_config.num_hidden_layers):
- # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"pretrained.model.blocks.{i}.attn.qkv.weight")
- q_bias = state_dict.pop(f"pretrained.model.blocks.{i}.attn.q_bias")
- v_bias = state_dict.pop(f"pretrained.model.blocks.{i}.attn.v_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[:hidden_size, :]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.bias"] = q_bias
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- hidden_size : hidden_size * 2, :
- ]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-hidden_size:, :]
- state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.bias"] = v_bias
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-@torch.no_grad()
-def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub):
- """
- Copy/paste/tweak model's weights to our DPT structure.
- """
-
- name_to_url = {
- "dpt-beit-large-512": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_512.pt",
- "dpt-beit-large-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_384.pt",
- "dpt-beit-base-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_base_384.pt",
- }
-
- # define DPT configuration based on URL
- checkpoint_url = name_to_url[model_name]
- config, image_size = get_dpt_config(model_name)
- # load original state_dict from URL
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
- # remove certain keys
- remove_ignore_keys_(state_dict)
- # rename keys
- rename_keys = create_rename_keys(config)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- # read in qkv matrices
- read_in_q_k_v(state_dict, config)
-
- # load HuggingFace model
- model = DPTForDepthEstimation(config)
- missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
- print("Missing keys:", missing_keys)
- print("Unexpected keys:", unexpected_keys)
- assert missing_keys == []
- # assert unexpected_keys == ["pretrained.model.fc_norm.weight", "pretrained.model.fc_norm.bias"]
- model.eval()
-
- # Check outputs on an image
- # We set `keep_aspect_ratio=False` as our current BEiT does not support arbitrary window sizes
- processor = DPTImageProcessor(
- size={"height": image_size, "width": image_size}, keep_aspect_ratio=False, ensure_multiple_of=32
- )
-
- image = prepare_img()
- pixel_values = processor(image, return_tensors="pt").pixel_values
-
- print("First values of pixel values:", pixel_values[0, 0, :3, :3])
- print("Mean of pixel values:", pixel_values.mean().item())
- print("Shape of pixel values:", pixel_values.shape)
-
- import requests
- from PIL import Image
- from torchvision import transforms
-
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- image = Image.open(requests.get(url, stream=True).raw)
-
- transforms = transforms.Compose(
- [
- transforms.Resize((image_size, image_size)),
- transforms.ToTensor(),
- ]
- )
- pixel_values = transforms(image).unsqueeze(0)
-
- # forward pass
- with torch.no_grad():
- outputs = model(pixel_values)
-
- predicted_depth = outputs.predicted_depth
-
- print("Shape of predicted depth:", predicted_depth.shape)
- print("First values of predicted depth:", predicted_depth[0, :3, :3])
-
- # assert logits
- # TODO there's still a small difference with the original logits
- if model_name == "dpt-beit-large-512":
- # OK, checked
- expected_shape = torch.Size([1, 512, 512])
- expected_slice = torch.tensor(
- [[2804.6260, 2792.5708, 2812.9263], [2772.0288, 2780.1118, 2796.2529], [2748.1094, 2766.6558, 2766.9834]]
- )
- elif model_name == "dpt-beit-large-384":
- # OK, checked
- expected_shape = torch.Size([1, 384, 384])
- expected_slice = torch.tensor(
- [[1783.2273, 1780.5729, 1792.6453], [1759.9817, 1765.5359, 1778.5002], [1739.1633, 1754.7903, 1757.1990]],
- )
- elif model_name == "dpt-beit-base-384":
- # OK, checked
- expected_shape = torch.Size([1, 384, 384])
- expected_slice = torch.tensor(
- [[2898.4482, 2891.3750, 2904.8079], [2858.6685, 2877.2615, 2894.4507], [2842.1235, 2854.1023, 2861.6328]],
- )
-
- assert predicted_depth.shape == torch.Size(expected_shape)
- assert torch.allclose(predicted_depth[0, :3, :3], expected_slice)
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model and processor to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- print("Pushing model and processor to hub...")
- model.push_to_hub(repo_id=f"nielsr/{model_name}")
- processor.push_to_hub(repo_id=f"nielsr/{model_name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="dpt-beit-large-512",
- type=str,
- choices=["dpt-beit-large-512", "dpt-beit-large-384", "dpt-beit-base-384"],
- help="Name of the model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- help="Whether to push the model to the hub after conversion.",
- )
-
- args = parser.parse_args()
- convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/dpt/convert_dpt_hybrid_to_pytorch.py b/transformers/models/dpt/convert_dpt_hybrid_to_pytorch.py
deleted file mode 100644
index 0fa69adfaf39d54a8417c21328a30a6f5993eac4..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/convert_dpt_hybrid_to_pytorch.py
+++ /dev/null
@@ -1,316 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DPT checkpoints from the original repository. URL: https://github.com/isl-org/DPT"""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import cached_download, hf_hub_url
-from PIL import Image
-
-from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_dpt_config(checkpoint_url):
- config = DPTConfig(embedding_type="hybrid")
-
- if "large" in checkpoint_url:
- config.hidden_size = 1024
- config.intermediate_size = 4096
- config.num_hidden_layers = 24
- config.num_attention_heads = 16
- config.backbone_out_indices = [5, 11, 17, 23]
- config.neck_hidden_sizes = [256, 512, 1024, 1024]
- expected_shape = (1, 384, 384)
-
- if "nyu" or "midas" in checkpoint_url:
- config.hidden_size = 768
- config.reassemble_factors = [1, 1, 1, 0.5]
- config.neck_hidden_sizes = [256, 512, 768, 768]
- config.num_labels = 150
- config.patch_size = 16
- expected_shape = (1, 384, 384)
- config.use_batch_norm_in_fusion_residual = False
- config.readout_type = "project"
-
- if "ade" in checkpoint_url:
- config.use_batch_norm_in_fusion_residual = True
- config.hidden_size = 768
- config.reassemble_stage = [1, 1, 1, 0.5]
- config.num_labels = 150
- config.patch_size = 16
- repo_id = "huggingface/label-files"
- filename = "ade20k-id2label.json"
- id2label = json.load(open(cached_download(hf_hub_url(repo_id, filename, repo_type="dataset")), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
- expected_shape = [1, 150, 480, 480]
-
- return config, expected_shape
-
-
-def remove_ignore_keys_(state_dict):
- ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"]
- for k in ignore_keys:
- state_dict.pop(k, None)
-
-
-def rename_key(name):
- if (
- "pretrained.model" in name
- and "cls_token" not in name
- and "pos_embed" not in name
- and "patch_embed" not in name
- ):
- name = name.replace("pretrained.model", "dpt.encoder")
- if "pretrained.model" in name:
- name = name.replace("pretrained.model", "dpt.embeddings")
- if "patch_embed" in name:
- name = name.replace("patch_embed", "")
- if "pos_embed" in name:
- name = name.replace("pos_embed", "position_embeddings")
- if "attn.proj" in name:
- name = name.replace("attn.proj", "attention.output.dense")
- if "proj" in name and "project" not in name:
- name = name.replace("proj", "projection")
- if "blocks" in name:
- name = name.replace("blocks", "layer")
- if "mlp.fc1" in name:
- name = name.replace("mlp.fc1", "intermediate.dense")
- if "mlp.fc2" in name:
- name = name.replace("mlp.fc2", "output.dense")
- if "norm1" in name and "backbone" not in name:
- name = name.replace("norm1", "layernorm_before")
- if "norm2" in name and "backbone" not in name:
- name = name.replace("norm2", "layernorm_after")
- if "scratch.output_conv" in name:
- name = name.replace("scratch.output_conv", "head")
- if "scratch" in name:
- name = name.replace("scratch", "neck")
- if "layer1_rn" in name:
- name = name.replace("layer1_rn", "convs.0")
- if "layer2_rn" in name:
- name = name.replace("layer2_rn", "convs.1")
- if "layer3_rn" in name:
- name = name.replace("layer3_rn", "convs.2")
- if "layer4_rn" in name:
- name = name.replace("layer4_rn", "convs.3")
- if "refinenet" in name:
- layer_idx = int(name[len("neck.refinenet") : len("neck.refinenet") + 1])
- # tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3
- name = name.replace(f"refinenet{layer_idx}", f"fusion_stage.layers.{abs(layer_idx-4)}")
- if "out_conv" in name:
- name = name.replace("out_conv", "projection")
- if "resConfUnit1" in name:
- name = name.replace("resConfUnit1", "residual_layer1")
- if "resConfUnit2" in name:
- name = name.replace("resConfUnit2", "residual_layer2")
- if "conv1" in name:
- name = name.replace("conv1", "convolution1")
- if "conv2" in name:
- name = name.replace("conv2", "convolution2")
- # readout blocks
- if "pretrained.act_postprocess1.0.project.0" in name:
- name = name.replace("pretrained.act_postprocess1.0.project.0", "neck.reassemble_stage.readout_projects.0.0")
- if "pretrained.act_postprocess2.0.project.0" in name:
- name = name.replace("pretrained.act_postprocess2.0.project.0", "neck.reassemble_stage.readout_projects.1.0")
- if "pretrained.act_postprocess3.0.project.0" in name:
- name = name.replace("pretrained.act_postprocess3.0.project.0", "neck.reassemble_stage.readout_projects.2.0")
- if "pretrained.act_postprocess4.0.project.0" in name:
- name = name.replace("pretrained.act_postprocess4.0.project.0", "neck.reassemble_stage.readout_projects.3.0")
-
- # resize blocks
- if "pretrained.act_postprocess1.3" in name:
- name = name.replace("pretrained.act_postprocess1.3", "neck.reassemble_stage.layers.0.projection")
- if "pretrained.act_postprocess1.4" in name:
- name = name.replace("pretrained.act_postprocess1.4", "neck.reassemble_stage.layers.0.resize")
- if "pretrained.act_postprocess2.3" in name:
- name = name.replace("pretrained.act_postprocess2.3", "neck.reassemble_stage.layers.1.projection")
- if "pretrained.act_postprocess2.4" in name:
- name = name.replace("pretrained.act_postprocess2.4", "neck.reassemble_stage.layers.1.resize")
- if "pretrained.act_postprocess3.3" in name:
- name = name.replace("pretrained.act_postprocess3.3", "neck.reassemble_stage.layers.2.projection")
- if "pretrained.act_postprocess4.3" in name:
- name = name.replace("pretrained.act_postprocess4.3", "neck.reassemble_stage.layers.3.projection")
- if "pretrained.act_postprocess4.4" in name:
- name = name.replace("pretrained.act_postprocess4.4", "neck.reassemble_stage.layers.3.resize")
- if "pretrained" in name:
- name = name.replace("pretrained", "dpt")
- if "bn" in name:
- name = name.replace("bn", "batch_norm")
- if "head" in name:
- name = name.replace("head", "head.head")
- if "encoder.norm" in name:
- name = name.replace("encoder.norm", "layernorm")
- if "auxlayer" in name:
- name = name.replace("auxlayer", "auxiliary_head.head")
- if "backbone" in name:
- name = name.replace("backbone", "backbone.bit.encoder")
-
- if ".." in name:
- name = name.replace("..", ".")
-
- if "stem.conv" in name:
- name = name.replace("stem.conv", "bit.embedder.convolution")
- if "blocks" in name:
- name = name.replace("blocks", "layers")
- if "convolution" in name and "backbone" in name:
- name = name.replace("convolution", "conv")
- if "layer" in name and "backbone" in name:
- name = name.replace("layer", "layers")
- if "backbone.bit.encoder.bit" in name:
- name = name.replace("backbone.bit.encoder.bit", "backbone.bit")
- if "embedder.conv" in name:
- name = name.replace("embedder.conv", "embedder.convolution")
- if "backbone.bit.encoder.stem.norm" in name:
- name = name.replace("backbone.bit.encoder.stem.norm", "backbone.bit.embedder.norm")
- return name
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config):
- for i in range(config.num_hidden_layers):
- # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.weight")
- in_proj_bias = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[: config.hidden_size, :]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- config.hidden_size : config.hidden_size * 2, :
- ]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
- config.hidden_size : config.hidden_size * 2
- ]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
- -config.hidden_size :, :
- ]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-@torch.no_grad()
-def convert_dpt_checkpoint(checkpoint_url, pytorch_dump_folder_path, push_to_hub, model_name, show_prediction):
- """
- Copy/paste/tweak model's weights to our DPT structure.
- """
-
- # define DPT configuration based on URL
- config, expected_shape = get_dpt_config(checkpoint_url)
- # load original state_dict from URL
- # state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
- state_dict = torch.load(checkpoint_url, map_location="cpu")
- # remove certain keys
- remove_ignore_keys_(state_dict)
- # rename keys
- for key in state_dict.copy().keys():
- val = state_dict.pop(key)
- state_dict[rename_key(key)] = val
- # read in qkv matrices
- read_in_q_k_v(state_dict, config)
-
- # load HuggingFace model
- model = DPTForSemanticSegmentation(config) if "ade" in checkpoint_url else DPTForDepthEstimation(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- # Check outputs on an image
- size = 480 if "ade" in checkpoint_url else 384
- image_processor = DPTImageProcessor(size=size)
-
- image = prepare_img()
- encoding = image_processor(image, return_tensors="pt")
-
- # forward pass
- outputs = model(**encoding).logits if "ade" in checkpoint_url else model(**encoding).predicted_depth
-
- if show_prediction:
- prediction = (
- torch.nn.functional.interpolate(
- outputs.unsqueeze(1),
- size=(image.size[1], image.size[0]),
- mode="bicubic",
- align_corners=False,
- )
- .squeeze()
- .cpu()
- .numpy()
- )
-
- Image.fromarray((prediction / prediction.max()) * 255).show()
-
- if pytorch_dump_folder_path is not None:
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- print(f"Saving image processor to {pytorch_dump_folder_path}")
- image_processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- model.push_to_hub("ybelkada/dpt-hybrid-midas")
- image_processor.push_to_hub("ybelkada/dpt-hybrid-midas")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--checkpoint_url",
- default="https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt",
- type=str,
- help="URL of the original DPT checkpoint you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- required=False,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- )
- parser.add_argument(
- "--model_name",
- default="dpt-large",
- type=str,
- help="Name of the model, in case you're pushing to the hub.",
- )
- parser.add_argument(
- "--show_prediction",
- action="store_true",
- )
-
- args = parser.parse_args()
- convert_dpt_checkpoint(
- args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name, args.show_prediction
- )
diff --git a/transformers/models/dpt/convert_dpt_swinv2_to_hf.py b/transformers/models/dpt/convert_dpt_swinv2_to_hf.py
deleted file mode 100644
index fd6522ab6be331dc605eb8b87b3c1784a744cac2..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/convert_dpt_swinv2_to_hf.py
+++ /dev/null
@@ -1,322 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DPT 3.1 checkpoints from the MiDaS repository. URL: https://github.com/isl-org/MiDaS"""
-
-
-import argparse
-from pathlib import Path
-
-import requests
-import torch
-from PIL import Image
-
-from transformers import DPTConfig, DPTForDepthEstimation, DPTImageProcessor, Swinv2Config
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_dpt_config(model_name):
- if "tiny" in model_name:
- embed_dim = 96
- depths = (2, 2, 6, 2)
- num_heads = (3, 6, 12, 24)
- window_size = 16
- # note: for Swinv2-tiny authors used the window_size = 16 variant
- # as seen here: https://github.com/isl-org/MiDaS/blob/bdc4ed64c095e026dc0a2f17cabb14d58263decb/midas/backbones/swin2.py#L26
- pretrained_window_sizes = (0, 0, 0, 0)
- elif "base" in model_name:
- embed_dim = 128
- depths = (2, 2, 18, 2)
- num_heads = (4, 8, 16, 32)
- window_size = 24
- pretrained_window_sizes = (12, 12, 12, 6)
- elif "large" in model_name:
- embed_dim = 192
- depths = (2, 2, 18, 2)
- num_heads = (6, 12, 24, 48)
- window_size = 24
- pretrained_window_sizes = (12, 12, 12, 6)
-
- if "384" in model_name:
- image_size = 384
- elif "256" in model_name:
- image_size = 256
- else:
- raise ValueError("Model not supported, to do")
-
- backbone_config = Swinv2Config(
- image_size=image_size,
- embed_dim=embed_dim,
- depths=depths,
- window_size=window_size,
- pretrained_window_sizes=pretrained_window_sizes,
- num_heads=num_heads,
- out_features=["stage1", "stage2", "stage3", "stage4"],
- )
-
- if model_name == "dpt-swinv2-tiny-256":
- neck_hidden_sizes = [96, 192, 384, 768]
- elif model_name == "dpt-swinv2-base-384":
- neck_hidden_sizes = [128, 256, 512, 1024]
- elif model_name == "dpt-swinv2-large-384":
- neck_hidden_sizes = [192, 384, 768, 1536]
-
- config = DPTConfig(backbone_config=backbone_config, neck_hidden_sizes=neck_hidden_sizes)
-
- return config, image_size
-
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys(config):
- rename_keys = []
-
- # fmt: off
- # stem
- rename_keys.append(("pretrained.model.patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight"))
- rename_keys.append(("pretrained.model.patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias"))
- rename_keys.append(("pretrained.model.patch_embed.norm.weight", "backbone.embeddings.norm.weight"))
- rename_keys.append(("pretrained.model.patch_embed.norm.bias", "backbone.embeddings.norm.bias"))
-
- # transformer encoder
- for i in range(len(config.backbone_config.depths)):
- for j in range(config.backbone_config.depths[i]):
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.logit_scale", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.logit_scale"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.cpb_mlp.0.weight", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.continuous_position_bias_mlp.0.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.cpb_mlp.0.bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.continuous_position_bias_mlp.0.bias"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.cpb_mlp.2.weight", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.continuous_position_bias_mlp.2.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.q_bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.query.bias"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.v_bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.value.bias"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.proj.weight", f"backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.attn.proj.bias", f"backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.norm1.weight", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.norm1.bias", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.mlp.fc1.weight", f"backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.mlp.fc1.bias", f"backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.mlp.fc2.weight", f"backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.mlp.fc2.bias", f"backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.norm2.weight", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.blocks.{j}.norm2.bias", f"backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias"))
-
- # downsample parameters
- if i in [0,1,2]:
- rename_keys.append((f"pretrained.model.layers.{i}.downsample.reduction.weight", f"backbone.encoder.layers.{i}.downsample.reduction.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.downsample.norm.weight", f"backbone.encoder.layers.{i}.downsample.norm.weight"))
- rename_keys.append((f"pretrained.model.layers.{i}.downsample.norm.bias", f"backbone.encoder.layers.{i}.downsample.norm.bias"))
-
- # note: non-Transformer backbones like Swinv2, LeViT et al don't require activation postprocessing (readout projections + resize blocks)
-
- # refinenet (tricky here)
- mapping = {1:3, 2:2, 3:1, 4:0}
-
- for i in range(1, 5):
- j = mapping[i]
- rename_keys.append((f"scratch.refinenet{i}.out_conv.weight", f"neck.fusion_stage.layers.{j}.projection.weight"))
- rename_keys.append((f"scratch.refinenet{i}.out_conv.bias", f"neck.fusion_stage.layers.{j}.projection.bias"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.weight"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.bias"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.weight"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.bias"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.weight"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.bias"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.weight"))
- rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.bias"))
-
- # scratch convolutions
- for i in range(4):
- rename_keys.append((f"scratch.layer{i+1}_rn.weight", f"neck.convs.{i}.weight"))
-
- # head
- for i in range(0, 5, 2):
- rename_keys.append((f"scratch.output_conv.{i}.weight", f"head.head.{i}.weight"))
- rename_keys.append((f"scratch.output_conv.{i}.bias", f"head.head.{i}.bias"))
-
- return rename_keys
-
-
-def remove_ignore_keys_(state_dict):
- ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"]
- for k in ignore_keys:
- state_dict.pop(k, None)
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config, model):
- for i in range(len(config.backbone_config.depths)):
- for j in range(config.backbone_config.depths[i]):
- dim = model.backbone.encoder.layers[i].blocks[j].attention.self.all_head_size
- # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"pretrained.model.layers.{i}.blocks.{j}.attn.qkv.weight")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.query.weight"] = in_proj_weight[:dim, :]
- state_dict[f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.key.weight"] = in_proj_weight[
- dim : dim * 2, :
- ]
- state_dict[f"backbone.encoder.layers.{i}.blocks.{j}.attention.self.value.weight"] = in_proj_weight[
- -dim:, :
- ]
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-@torch.no_grad()
-def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, verify_logits, push_to_hub):
- """
- Copy/paste/tweak model's weights to our DPT structure.
- """
-
- name_to_url = {
- "dpt-swinv2-tiny-256": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_tiny_256.pt",
- "dpt-swinv2-base-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_base_384.pt",
- "dpt-swinv2-large-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_large_384.pt",
- }
-
- # define DPT configuration based on URL
- checkpoint_url = name_to_url[model_name]
- config, image_size = get_dpt_config(model_name)
- # load original state_dict from URL
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
-
- # load HuggingFace model
- model = DPTForDepthEstimation(config)
-
- # remove certain keys
- remove_ignore_keys_(state_dict)
- # rename keys
- rename_keys = create_rename_keys(config)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- # read in qkv matrices
- read_in_q_k_v(state_dict, config, model)
-
- missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
- print("Missing keys:", missing_keys)
- print("Unexpected keys:", unexpected_keys)
- model.eval()
-
- # Check outputs on an image
- processor = DPTImageProcessor(size={"height": image_size, "width": image_size})
-
- image = prepare_img()
- processor(image, return_tensors="pt")
-
- if verify_logits:
- from torchvision import transforms
-
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- image = Image.open(requests.get(url, stream=True).raw)
-
- transforms = transforms.Compose(
- [
- transforms.Resize((image_size, image_size)),
- transforms.ToTensor(),
- ]
- )
- pixel_values = transforms(image).unsqueeze(0)
-
- # forward pass
- with torch.no_grad():
- outputs = model(pixel_values)
-
- predicted_depth = outputs.predicted_depth
-
- print("Shape of predicted depth:", predicted_depth.shape)
- print("First values of predicted depth:", predicted_depth[0, :3, :3])
-
- # assert logits
- if model_name == "dpt-swinv2-base-384":
- # OK, checked
- expected_shape = torch.Size([1, 384, 384])
- expected_slice = torch.tensor(
- [
- [1998.5575, 1997.3887, 2009.2981],
- [1952.8607, 1979.6488, 2001.0854],
- [1953.7697, 1961.7711, 1968.8904],
- ],
- )
- elif model_name == "dpt-swinv2-tiny-256":
- # OK, checked
- expected_shape = torch.Size([1, 256, 256])
- expected_slice = torch.tensor(
- [[978.9163, 976.5215, 978.5349], [974.1859, 971.7249, 975.8046], [971.3419, 970.3118, 971.6830]],
- )
- elif model_name == "dpt-swinv2-large-384":
- # OK, checked
- expected_shape = torch.Size([1, 384, 384])
- expected_slice = torch.tensor(
- [
- [1203.7206, 1200.1495, 1197.8234],
- [1196.2484, 1183.5033, 1186.4640],
- [1178.8131, 1182.3260, 1174.3975],
- ],
- )
-
- assert predicted_depth.shape == torch.Size(expected_shape)
- assert torch.allclose(predicted_depth[0, :3, :3], expected_slice)
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model and processor to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- print("Pushing model and processor to hub...")
- model.push_to_hub(repo_id=f"Intel/{model_name}")
- processor.push_to_hub(repo_id=f"Intel/{model_name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="dpt-swinv2-base-384",
- type=str,
- choices=["dpt-swinv2-tiny-256", "dpt-swinv2-base-384", "dpt-swinv2-large-384"],
- help="Name of the model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--verify_logits",
- action="store_true",
- help="Whether to verify logits after conversion.",
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- help="Whether to push the model to the hub after conversion.",
- )
-
- args = parser.parse_args()
- convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.verify_logits, args.push_to_hub)
diff --git a/transformers/models/dpt/convert_dpt_to_pytorch.py b/transformers/models/dpt/convert_dpt_to_pytorch.py
deleted file mode 100644
index 42637cb21587245f1bd613d25b4094d8a5aebe31..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/convert_dpt_to_pytorch.py
+++ /dev/null
@@ -1,286 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert DPT checkpoints from the original repository. URL: https://github.com/isl-org/DPT"""
-
-
-import argparse
-import json
-from pathlib import Path
-
-import requests
-import torch
-from huggingface_hub import cached_download, hf_hub_url
-from PIL import Image
-
-from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_dpt_config(checkpoint_url):
- config = DPTConfig()
-
- if "large" in checkpoint_url:
- config.hidden_size = 1024
- config.intermediate_size = 4096
- config.num_hidden_layers = 24
- config.num_attention_heads = 16
- config.backbone_out_indices = [5, 11, 17, 23]
- config.neck_hidden_sizes = [256, 512, 1024, 1024]
- expected_shape = (1, 384, 384)
-
- if "ade" in checkpoint_url:
- config.use_batch_norm_in_fusion_residual = True
-
- config.num_labels = 150
- repo_id = "huggingface/label-files"
- filename = "ade20k-id2label.json"
- id2label = json.load(open(cached_download(hf_hub_url(repo_id, filename, repo_type="dataset")), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
- expected_shape = [1, 150, 480, 480]
-
- return config, expected_shape
-
-
-def remove_ignore_keys_(state_dict):
- ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"]
- for k in ignore_keys:
- state_dict.pop(k, None)
-
-
-def rename_key(name):
- if (
- "pretrained.model" in name
- and "cls_token" not in name
- and "pos_embed" not in name
- and "patch_embed" not in name
- ):
- name = name.replace("pretrained.model", "dpt.encoder")
- if "pretrained.model" in name:
- name = name.replace("pretrained.model", "dpt.embeddings")
- if "patch_embed" in name:
- name = name.replace("patch_embed", "patch_embeddings")
- if "pos_embed" in name:
- name = name.replace("pos_embed", "position_embeddings")
- if "attn.proj" in name:
- name = name.replace("attn.proj", "attention.output.dense")
- if "proj" in name and "project" not in name:
- name = name.replace("proj", "projection")
- if "blocks" in name:
- name = name.replace("blocks", "layer")
- if "mlp.fc1" in name:
- name = name.replace("mlp.fc1", "intermediate.dense")
- if "mlp.fc2" in name:
- name = name.replace("mlp.fc2", "output.dense")
- if "norm1" in name:
- name = name.replace("norm1", "layernorm_before")
- if "norm2" in name:
- name = name.replace("norm2", "layernorm_after")
- if "scratch.output_conv" in name:
- name = name.replace("scratch.output_conv", "head")
- if "scratch" in name:
- name = name.replace("scratch", "neck")
- if "layer1_rn" in name:
- name = name.replace("layer1_rn", "convs.0")
- if "layer2_rn" in name:
- name = name.replace("layer2_rn", "convs.1")
- if "layer3_rn" in name:
- name = name.replace("layer3_rn", "convs.2")
- if "layer4_rn" in name:
- name = name.replace("layer4_rn", "convs.3")
- if "refinenet" in name:
- layer_idx = int(name[len("neck.refinenet") : len("neck.refinenet") + 1])
- # tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3
- name = name.replace(f"refinenet{layer_idx}", f"fusion_stage.layers.{abs(layer_idx-4)}")
- if "out_conv" in name:
- name = name.replace("out_conv", "projection")
- if "resConfUnit1" in name:
- name = name.replace("resConfUnit1", "residual_layer1")
- if "resConfUnit2" in name:
- name = name.replace("resConfUnit2", "residual_layer2")
- if "conv1" in name:
- name = name.replace("conv1", "convolution1")
- if "conv2" in name:
- name = name.replace("conv2", "convolution2")
- # readout blocks
- if "pretrained.act_postprocess1.0.project.0" in name:
- name = name.replace("pretrained.act_postprocess1.0.project.0", "neck.reassemble_stage.readout_projects.0.0")
- if "pretrained.act_postprocess2.0.project.0" in name:
- name = name.replace("pretrained.act_postprocess2.0.project.0", "neck.reassemble_stage.readout_projects.1.0")
- if "pretrained.act_postprocess3.0.project.0" in name:
- name = name.replace("pretrained.act_postprocess3.0.project.0", "neck.reassemble_stage.readout_projects.2.0")
- if "pretrained.act_postprocess4.0.project.0" in name:
- name = name.replace("pretrained.act_postprocess4.0.project.0", "neck.reassemble_stage.readout_projects.3.0")
- # resize blocks
- if "pretrained.act_postprocess1.3" in name:
- name = name.replace("pretrained.act_postprocess1.3", "neck.reassemble_stage.layers.0.projection")
- if "pretrained.act_postprocess1.4" in name:
- name = name.replace("pretrained.act_postprocess1.4", "neck.reassemble_stage.layers.0.resize")
- if "pretrained.act_postprocess2.3" in name:
- name = name.replace("pretrained.act_postprocess2.3", "neck.reassemble_stage.layers.1.projection")
- if "pretrained.act_postprocess2.4" in name:
- name = name.replace("pretrained.act_postprocess2.4", "neck.reassemble_stage.layers.1.resize")
- if "pretrained.act_postprocess3.3" in name:
- name = name.replace("pretrained.act_postprocess3.3", "neck.reassemble_stage.layers.2.projection")
- if "pretrained.act_postprocess4.3" in name:
- name = name.replace("pretrained.act_postprocess4.3", "neck.reassemble_stage.layers.3.projection")
- if "pretrained.act_postprocess4.4" in name:
- name = name.replace("pretrained.act_postprocess4.4", "neck.reassemble_stage.layers.3.resize")
- if "pretrained" in name:
- name = name.replace("pretrained", "dpt")
- if "bn" in name:
- name = name.replace("bn", "batch_norm")
- if "head" in name:
- name = name.replace("head", "head.head")
- if "encoder.norm" in name:
- name = name.replace("encoder.norm", "layernorm")
- if "auxlayer" in name:
- name = name.replace("auxlayer", "auxiliary_head.head")
-
- return name
-
-
-# we split up the matrix of each encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config):
- for i in range(config.num_hidden_layers):
- # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.weight")
- in_proj_bias = state_dict.pop(f"dpt.encoder.layer.{i}.attn.qkv.bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[: config.hidden_size, :]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.query.bias"] = in_proj_bias[: config.hidden_size]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
- config.hidden_size : config.hidden_size * 2, :
- ]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.key.bias"] = in_proj_bias[
- config.hidden_size : config.hidden_size * 2
- ]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
- -config.hidden_size :, :
- ]
- state_dict[f"dpt.encoder.layer.{i}.attention.attention.value.bias"] = in_proj_bias[-config.hidden_size :]
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-@torch.no_grad()
-def convert_dpt_checkpoint(checkpoint_url, pytorch_dump_folder_path, push_to_hub, model_name):
- """
- Copy/paste/tweak model's weights to our DPT structure.
- """
-
- # define DPT configuration based on URL
- config, expected_shape = get_dpt_config(checkpoint_url)
- # load original state_dict from URL
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
- # remove certain keys
- remove_ignore_keys_(state_dict)
- # rename keys
- for key in state_dict.copy().keys():
- val = state_dict.pop(key)
- state_dict[rename_key(key)] = val
- # read in qkv matrices
- read_in_q_k_v(state_dict, config)
-
- # load HuggingFace model
- model = DPTForSemanticSegmentation(config) if "ade" in checkpoint_url else DPTForDepthEstimation(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- # Check outputs on an image
- size = 480 if "ade" in checkpoint_url else 384
- image_processor = DPTImageProcessor(size=size)
-
- image = prepare_img()
- encoding = image_processor(image, return_tensors="pt")
-
- # forward pass
- outputs = model(**encoding).logits if "ade" in checkpoint_url else model(**encoding).predicted_depth
-
- # Assert logits
- expected_slice = torch.tensor([[6.3199, 6.3629, 6.4148], [6.3850, 6.3615, 6.4166], [6.3519, 6.3176, 6.3575]])
- if "ade" in checkpoint_url:
- expected_slice = torch.tensor([[4.0480, 4.2420, 4.4360], [4.3124, 4.5693, 4.8261], [4.5768, 4.8965, 5.2163]])
- assert outputs.shape == torch.Size(expected_shape)
- assert (
- torch.allclose(outputs[0, 0, :3, :3], expected_slice, atol=1e-4)
- if "ade" in checkpoint_url
- else torch.allclose(outputs[0, :3, :3], expected_slice)
- )
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- print(f"Saving image processor to {pytorch_dump_folder_path}")
- image_processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- print("Pushing model to hub...")
- model.push_to_hub(
- repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
- organization="nielsr",
- commit_message="Add model",
- use_temp_dir=True,
- )
- image_processor.push_to_hub(
- repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
- organization="nielsr",
- commit_message="Add image processor",
- use_temp_dir=True,
- )
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--checkpoint_url",
- default="https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt",
- type=str,
- help="URL of the original DPT checkpoint you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- required=False,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- )
- parser.add_argument(
- "--model_name",
- default="dpt-large",
- type=str,
- required=False,
- help="Name of the model, in case you're pushing to the hub.",
- )
-
- args = parser.parse_args()
- convert_dpt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name)
diff --git a/transformers/models/dpt/feature_extraction_dpt.py b/transformers/models/dpt/feature_extraction_dpt.py
deleted file mode 100644
index d375d8229f5ee9b3278af363c40043815ff0cf29..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/feature_extraction_dpt.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for DPT."""
-
-import warnings
-
-from ...utils import logging
-from .image_processing_dpt import DPTImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-class DPTFeatureExtractor(DPTImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class DPTFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
- " use DPTImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/dpt/image_processing_dpt.py b/transformers/models/dpt/image_processing_dpt.py
deleted file mode 100644
index 96f43a796e3886b8bf78599bc425a15f63db25ba..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/image_processing_dpt.py
+++ /dev/null
@@ -1,484 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for DPT."""
-
-import math
-from typing import Dict, Iterable, List, Optional, Tuple, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
-from ...image_transforms import pad, resize, to_channel_dimension_format
-from ...image_utils import (
- IMAGENET_STANDARD_MEAN,
- IMAGENET_STANDARD_STD,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- get_image_size,
- infer_channel_dimension_format,
- is_scaled_image,
- is_torch_available,
- is_torch_tensor,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import TensorType, is_vision_available, logging
-
-
-if is_torch_available():
- import torch
-
-if is_vision_available():
- import PIL
-
-
-logger = logging.get_logger(__name__)
-
-
-def get_resize_output_image_size(
- input_image: np.ndarray,
- output_size: Union[int, Iterable[int]],
- keep_aspect_ratio: bool,
- multiple: int,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
-) -> Tuple[int, int]:
- def constraint_to_multiple_of(val, multiple, min_val=0, max_val=None):
- x = round(val / multiple) * multiple
-
- if max_val is not None and x > max_val:
- x = math.floor(val / multiple) * multiple
-
- if x < min_val:
- x = math.ceil(val / multiple) * multiple
-
- return x
-
- output_size = (output_size, output_size) if isinstance(output_size, int) else output_size
-
- input_height, input_width = get_image_size(input_image, input_data_format)
- output_height, output_width = output_size
-
- # determine new height and width
- scale_height = output_height / input_height
- scale_width = output_width / input_width
-
- if keep_aspect_ratio:
- # scale as little as possible
- if abs(1 - scale_width) < abs(1 - scale_height):
- # fit width
- scale_height = scale_width
- else:
- # fit height
- scale_width = scale_height
-
- new_height = constraint_to_multiple_of(scale_height * input_height, multiple=multiple)
- new_width = constraint_to_multiple_of(scale_width * input_width, multiple=multiple)
-
- return (new_height, new_width)
-
-
-class DPTImageProcessor(BaseImageProcessor):
- r"""
- Constructs a DPT image processor.
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the image's (height, width) dimensions. Can be overidden by `do_resize` in `preprocess`.
- size (`Dict[str, int]` *optional*, defaults to `{"height": 384, "width": 384}`):
- Size of the image after resizing. Can be overidden by `size` in `preprocess`.
- resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
- Defines the resampling filter to use if resizing the image. Can be overidden by `resample` in `preprocess`.
- keep_aspect_ratio (`bool`, *optional*, defaults to `False`):
- If `True`, the image is resized to the largest possible size such that the aspect ratio is preserved. Can
- be overidden by `keep_aspect_ratio` in `preprocess`.
- ensure_multiple_of (`int`, *optional*, defaults to 1):
- If `do_resize` is `True`, the image is resized to a size that is a multiple of this value. Can be overidden
- by `ensure_multiple_of` in `preprocess`.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image by the specified scale `rescale_factor`. Can be overidden by `do_rescale` in
- `preprocess`.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overidden by `rescale_factor` in `preprocess`.
- do_normalize (`bool`, *optional*, defaults to `True`):
- Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
- method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
- Mean to use if normalizing the image. This is a float or list of floats the length of the number of
- channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
- Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
- number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
- do_pad (`bool`, *optional*, defaults to `False`):
- Whether to apply center padding. This was introduced in the DINOv2 paper, which uses the model in
- combination with DPT.
- size_divisor (`int`, *optional*):
- If `do_pad` is `True`, pads the image dimensions to be divisible by this value. This was introduced in the
- DINOv2 paper, which uses the model in combination with DPT.
- """
-
- model_input_names = ["pixel_values"]
-
- def __init__(
- self,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- keep_aspect_ratio: bool = False,
- ensure_multiple_of: int = 1,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- do_normalize: bool = True,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- do_pad: bool = False,
- size_divisor: int = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- size = size if size is not None else {"height": 384, "width": 384}
- size = get_size_dict(size)
- self.do_resize = do_resize
- self.size = size
- self.keep_aspect_ratio = keep_aspect_ratio
- self.ensure_multiple_of = ensure_multiple_of
- self.resample = resample
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_normalize = do_normalize
- self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
- self.do_pad = do_pad
- self.size_divisor = size_divisor
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size",
- "keep_aspect_ratio",
- "ensure_multiple_of",
- "resample",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "image_mean",
- "image_std",
- "do_pad",
- "size_divisor",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- keep_aspect_ratio: bool = False,
- ensure_multiple_of: int = 1,
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize an image to target size `(size["height"], size["width"])`. If `keep_aspect_ratio` is `True`, the image
- is resized to the largest possible size such that the aspect ratio is preserved. If `ensure_multiple_of` is
- set, the image is resized to a size that is a multiple of this value.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Target size of the output image.
- keep_aspect_ratio (`bool`, *optional*, defaults to `False`):
- If `True`, the image is resized to the largest possible size such that the aspect ratio is preserved.
- ensure_multiple_of (`int`, *optional*, defaults to 1):
- The image is resized to a size that is a multiple of this value.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
- Defines the resampling filter to use if resizing the image. Otherwise, the image is resized to size
- specified in `size`.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
- Resampling filter to use when resiizing the image.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the image. If not provided, it will be the same as the input image.
- input_data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- size = get_size_dict(size)
- if "height" not in size or "width" not in size:
- raise ValueError(f"The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}")
-
- output_size = get_resize_output_image_size(
- image,
- output_size=(size["height"], size["width"]),
- keep_aspect_ratio=keep_aspect_ratio,
- multiple=ensure_multiple_of,
- input_data_format=input_data_format,
- )
- return resize(
- image,
- size=output_size,
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
-
- def pad_image(
- self,
- image: np.array,
- size_divisor: int,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ):
- """
- Center pad an image to be a multiple of `multiple`.
-
- Args:
- image (`np.ndarray`):
- Image to pad.
- size_divisor (`int`):
- The width and height of the image will be padded to a multiple of this number.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - Unset: Use the channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
-
- def _get_pad(size, size_divisor):
- new_size = math.ceil(size / size_divisor) * size_divisor
- pad_size = new_size - size
- pad_size_left = pad_size // 2
- pad_size_right = pad_size - pad_size_left
- return pad_size_left, pad_size_right
-
- if input_data_format is None:
- input_data_format = infer_channel_dimension_format(image)
-
- height, width = get_image_size(image, input_data_format)
-
- pad_size_left, pad_size_right = _get_pad(height, size_divisor)
- pad_size_top, pad_size_bottom = _get_pad(width, size_divisor)
-
- return pad(image, ((pad_size_left, pad_size_right), (pad_size_top, pad_size_bottom)), data_format=data_format)
-
- def preprocess(
- self,
- images: ImageInput,
- do_resize: bool = None,
- size: int = None,
- keep_aspect_ratio: bool = None,
- ensure_multiple_of: int = None,
- resample: PILImageResampling = None,
- do_rescale: bool = None,
- rescale_factor: float = None,
- do_normalize: bool = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- do_pad: bool = None,
- size_divisor: int = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: ChannelDimension = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> PIL.Image.Image:
- """
- Preprocess an image or batch of images.
-
- Args:
- images (`ImageInput`):
- Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
- passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
- Size of the image after reszing. If `keep_aspect_ratio` is `True`, the image is resized to the largest
- possible size such that the aspect ratio is preserved. If `ensure_multiple_of` is set, the image is
- resized to a size that is a multiple of this value.
- keep_aspect_ratio (`bool`, *optional*, defaults to `self.keep_aspect_ratio`):
- Whether to keep the aspect ratio of the image. If False, the image will be resized to (size, size). If
- True, the image will be resized to keep the aspect ratio and the size will be the maximum possible.
- ensure_multiple_of (`int`, *optional*, defaults to `self.ensure_multiple_of`):
- Ensure that the image size is a multiple of this value.
- resample (`int`, *optional*, defaults to `self.resample`):
- Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`, Only
- has an effect if `do_resize` is set to `True`.
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether to rescale the image values between [0 - 1].
- rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
- Rescale factor to rescale the image by if `do_rescale` is set to `True`.
- do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
- Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
- Image mean.
- image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
- Image standard deviation.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- do_resize = do_resize if do_resize is not None else self.do_resize
- size = size if size is not None else self.size
- size = get_size_dict(size)
- keep_aspect_ratio = keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
- ensure_multiple_of = ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
- resample = resample if resample is not None else self.resample
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- do_normalize = do_normalize if do_normalize is not None else self.do_normalize
- image_mean = image_mean if image_mean is not None else self.image_mean
- image_std = image_std if image_std is not None else self.image_std
- do_pad = do_pad if do_pad is not None else self.do_pad
- size_divisor = size_divisor if size_divisor is not None else self.size_divisor
-
- images = make_list_of_images(images)
-
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_pad=do_pad,
- size_divisibility=size_divisor,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
- # All transformations expect numpy arrays.
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- if do_resize:
- images = [
- self.resize(
- image=image,
- size=size,
- resample=resample,
- keep_aspect_ratio=keep_aspect_ratio,
- ensure_multiple_of=ensure_multiple_of,
- input_data_format=input_data_format,
- )
- for image in images
- ]
-
- if do_rescale:
- images = [
- self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_normalize:
- images = [
- self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_pad:
- images = [
- self.pad_image(image=image, size_divisor=size_divisor, input_data_format=input_data_format)
- for image in images
- ]
-
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
- ]
-
- data = {"pixel_values": images}
- return BatchFeature(data=data, tensor_type=return_tensors)
-
- # Copied from transformers.models.beit.image_processing_beit.BeitImageProcessor.post_process_semantic_segmentation with Beit->DPT
- def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
- """
- Converts the output of [`DPTForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.
-
- Args:
- outputs ([`DPTForSemanticSegmentation`]):
- Raw outputs of the model.
- target_sizes (`List[Tuple]` of length `batch_size`, *optional*):
- List of tuples corresponding to the requested final size (height, width) of each prediction. If unset,
- predictions will not be resized.
-
- Returns:
- semantic_segmentation: `List[torch.Tensor]` of length `batch_size`, where each item is a semantic
- segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is
- specified). Each entry of each `torch.Tensor` correspond to a semantic class id.
- """
- # TODO: add support for other frameworks
- logits = outputs.logits
-
- # Resize logits and compute semantic segmentation maps
- if target_sizes is not None:
- if len(logits) != len(target_sizes):
- raise ValueError(
- "Make sure that you pass in as many target sizes as the batch dimension of the logits"
- )
-
- if is_torch_tensor(target_sizes):
- target_sizes = target_sizes.numpy()
-
- semantic_segmentation = []
-
- for idx in range(len(logits)):
- resized_logits = torch.nn.functional.interpolate(
- logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
- )
- semantic_map = resized_logits[0].argmax(dim=0)
- semantic_segmentation.append(semantic_map)
- else:
- semantic_segmentation = logits.argmax(dim=1)
- semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]
-
- return semantic_segmentation
diff --git a/transformers/models/dpt/modeling_dpt.py b/transformers/models/dpt/modeling_dpt.py
deleted file mode 100644
index aad3330279f051c48d1dae64fd39649ffaa7ee62..0000000000000000000000000000000000000000
--- a/transformers/models/dpt/modeling_dpt.py
+++ /dev/null
@@ -1,1378 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Intel Labs, OpenMMLab and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch DPT (Dense Prediction Transformers) model.
-
-This implementation is heavily inspired by OpenMMLab's implementation, found here:
-https://github.com/open-mmlab/mmsegmentation/blob/master/mmseg/models/decode_heads/dpt_head.py.
-
-"""
-
-
-import collections.abc
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Set, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from ...activations import ACT2FN
-from ...file_utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- replace_return_docstrings,
-)
-from ...modeling_outputs import BaseModelOutput, DepthEstimatorOutput, SemanticSegmenterOutput
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import ModelOutput, logging
-from ...utils.backbone_utils import load_backbone
-from .configuration_dpt import DPTConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "DPTConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "Intel/dpt-large"
-_EXPECTED_OUTPUT_SHAPE = [1, 577, 1024]
-
-
-from ..deprecated._archive_maps import DPT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class BaseModelOutputWithIntermediateActivations(ModelOutput):
- """
- Base class for model's outputs that also contains intermediate activations that can be used at later stages. Useful
- in the context of Vision models.:
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- intermediate_activations (`tuple(torch.FloatTensor)`, *optional*):
- Intermediate activations that can be used to compute hidden states of the model at various layers.
- """
-
- last_hidden_states: torch.FloatTensor = None
- intermediate_activations: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-@dataclass
-class BaseModelOutputWithPoolingAndIntermediateActivations(ModelOutput):
- """
- Base class for model's outputs that also contains a pooling of the last hidden states as well as intermediate
- activations that can be used by the model at later stages.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):
- Last layer hidden-state of the first token of the sequence (classification token) after further processing
- through the layers used for the auxiliary pretraining task. E.g. for BERT-family of models, this returns
- the classification token after processing through a linear layer and a tanh activation function. The linear
- layer weights are trained from the next sentence prediction (classification) objective during pretraining.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
- one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- intermediate_activations (`tuple(torch.FloatTensor)`, *optional*):
- Intermediate activations that can be used to compute hidden states of the model at various layers.
- """
-
- last_hidden_state: torch.FloatTensor = None
- pooler_output: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
- intermediate_activations: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-class DPTViTHybridEmbeddings(nn.Module):
- """
- This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
- `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
- Transformer.
- """
-
- def __init__(self, config, feature_size=None):
- super().__init__()
- image_size, patch_size = config.image_size, config.patch_size
- num_channels, hidden_size = config.num_channels, config.hidden_size
-
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
-
- self.backbone = load_backbone(config)
- feature_dim = self.backbone.channels[-1]
- if len(self.backbone.channels) != 3:
- raise ValueError(f"Expected backbone to have 3 output features, got {len(self.backbone.channels)}")
- self.residual_feature_map_index = [0, 1] # Always take the output of the first and second backbone stage
-
- if feature_size is None:
- feat_map_shape = config.backbone_featmap_shape
- feature_size = feat_map_shape[-2:]
- feature_dim = feat_map_shape[1]
- else:
- feature_size = (
- feature_size if isinstance(feature_size, collections.abc.Iterable) else (feature_size, feature_size)
- )
- feature_dim = self.backbone.channels[-1]
-
- self.image_size = image_size
- self.patch_size = patch_size[0]
- self.num_channels = num_channels
-
- self.projection = nn.Conv2d(feature_dim, hidden_size, kernel_size=1)
-
- self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
- self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))
-
- def _resize_pos_embed(self, posemb, grid_size_height, grid_size_width, start_index=1):
- posemb_tok = posemb[:, :start_index]
- posemb_grid = posemb[0, start_index:]
-
- old_grid_size = int(math.sqrt(len(posemb_grid)))
-
- posemb_grid = posemb_grid.reshape(1, old_grid_size, old_grid_size, -1).permute(0, 3, 1, 2)
- posemb_grid = nn.functional.interpolate(posemb_grid, size=(grid_size_height, grid_size_width), mode="bilinear")
- posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, grid_size_height * grid_size_width, -1)
-
- posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
-
- return posemb
-
- def forward(
- self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False, return_dict: bool = False
- ) -> torch.Tensor:
- batch_size, num_channels, height, width = pixel_values.shape
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- if not interpolate_pos_encoding:
- if height != self.image_size[0] or width != self.image_size[1]:
- raise ValueError(
- f"Input image size ({height}*{width}) doesn't match model"
- f" ({self.image_size[0]}*{self.image_size[1]})."
- )
-
- position_embeddings = self._resize_pos_embed(
- self.position_embeddings, height // self.patch_size, width // self.patch_size
- )
-
- backbone_output = self.backbone(pixel_values)
-
- features = backbone_output.feature_maps[-1]
-
- # Retrieve also the intermediate activations to use them at later stages
- output_hidden_states = [backbone_output.feature_maps[index] for index in self.residual_feature_map_index]
-
- embeddings = self.projection(features).flatten(2).transpose(1, 2)
-
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- embeddings = torch.cat((cls_tokens, embeddings), dim=1)
-
- # add positional encoding to each token
- embeddings = embeddings + position_embeddings
-
- if not return_dict:
- return (embeddings, output_hidden_states)
-
- # Return hidden states and intermediate activations
- return BaseModelOutputWithIntermediateActivations(
- last_hidden_states=embeddings,
- intermediate_activations=output_hidden_states,
- )
-
-
-class DPTViTEmbeddings(nn.Module):
- """
- Construct the CLS token, position and patch embeddings.
-
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
- self.patch_embeddings = DPTViTPatchEmbeddings(config)
- num_patches = self.patch_embeddings.num_patches
- self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.config = config
-
- def _resize_pos_embed(self, posemb, grid_size_height, grid_size_width, start_index=1):
- posemb_tok = posemb[:, :start_index]
- posemb_grid = posemb[0, start_index:]
-
- old_grid_size = int(math.sqrt(len(posemb_grid)))
-
- posemb_grid = posemb_grid.reshape(1, old_grid_size, old_grid_size, -1).permute(0, 3, 1, 2)
- posemb_grid = nn.functional.interpolate(posemb_grid, size=(grid_size_height, grid_size_width), mode="bilinear")
- posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, grid_size_height * grid_size_width, -1)
-
- posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
-
- return posemb
-
- def forward(self, pixel_values, return_dict=False):
- batch_size, num_channels, height, width = pixel_values.shape
-
- # possibly interpolate position encodings to handle varying image sizes
- patch_size = self.config.patch_size
- position_embeddings = self._resize_pos_embed(
- self.position_embeddings, height // patch_size, width // patch_size
- )
-
- embeddings = self.patch_embeddings(pixel_values)
-
- batch_size, seq_len, _ = embeddings.size()
-
- # add the [CLS] token to the embedded patch tokens
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- embeddings = torch.cat((cls_tokens, embeddings), dim=1)
-
- # add positional encoding to each token
- embeddings = embeddings + position_embeddings
-
- embeddings = self.dropout(embeddings)
-
- if not return_dict:
- return (embeddings,)
-
- return BaseModelOutputWithIntermediateActivations(last_hidden_states=embeddings)
-
-
-class DPTViTPatchEmbeddings(nn.Module):
- """
- Image to Patch Embedding.
-
- """
-
- def __init__(self, config):
- super().__init__()
- image_size, patch_size = config.image_size, config.patch_size
- num_channels, hidden_size = config.num_channels, config.hidden_size
-
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.num_patches = num_patches
-
- self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
-
- def forward(self, pixel_values):
- batch_size, num_channels, height, width = pixel_values.shape
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
- return embeddings
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->DPT
-class DPTViTSelfAttention(nn.Module):
- def __init__(self, config: DPTConfig) -> None:
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
- f"heads {config.num_attention_heads}."
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
- self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
- self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self, hidden_states, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- mixed_query_layer = self.query(hidden_states)
-
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->DPT
-class DPTViTSelfOutput(nn.Module):
- """
- The residual connection is defined in DPTLayer instead of here (as is the case with other models), due to the
- layernorm applied before each block.
- """
-
- def __init__(self, config: DPTConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-class DPTViTAttention(nn.Module):
- def __init__(self, config: DPTConfig) -> None:
- super().__init__()
- self.attention = DPTViTSelfAttention(config)
- self.output = DPTViTSelfOutput(config)
- self.pruned_heads = set()
-
- # Copied from transformers.models.vit.modeling_vit.ViTAttention.prune_heads
- def prune_heads(self, heads: Set[int]) -> None:
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.attention.query = prune_linear_layer(self.attention.query, index)
- self.attention.key = prune_linear_layer(self.attention.key, index)
- self.attention.value = prune_linear_layer(self.attention.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
- self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- # Copied from transformers.models.vit.modeling_vit.ViTAttention.forward
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- self_outputs = self.attention(hidden_states, head_mask, output_attentions)
-
- attention_output = self.output(self_outputs[0], hidden_states)
-
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTIntermediate with ViT->DPT
-class DPTViTIntermediate(nn.Module):
- def __init__(self, config: DPTConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTOutput with ViT->DPT
-class DPTViTOutput(nn.Module):
- def __init__(self, config: DPTConfig) -> None:
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- hidden_states = hidden_states + input_tensor
-
- return hidden_states
-
-
-# copied from transformers.models.vit.modeling_vit.ViTLayer with ViTConfig->DPTConfig, ViTAttention->DPTViTAttention, ViTIntermediate->DPTViTIntermediate, ViTOutput->DPTViTOutput
-class DPTViTLayer(nn.Module):
- """This corresponds to the Block class in the timm implementation."""
-
- def __init__(self, config: DPTConfig) -> None:
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = DPTViTAttention(config)
- self.intermediate = DPTViTIntermediate(config)
- self.output = DPTViTOutput(config)
- self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- self_attention_outputs = self.attention(
- self.layernorm_before(hidden_states), # in ViT, layernorm is applied before self-attention
- head_mask,
- output_attentions=output_attentions,
- )
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- # first residual connection
- hidden_states = attention_output + hidden_states
-
- # in ViT, layernorm is also applied after self-attention
- layer_output = self.layernorm_after(hidden_states)
- layer_output = self.intermediate(layer_output)
-
- # second residual connection is done here
- layer_output = self.output(layer_output, hidden_states)
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-# copied from transformers.models.vit.modeling_vit.ViTEncoder with ViTConfig -> DPTConfig, ViTLayer->DPTViTLayer
-class DPTViTEncoder(nn.Module):
- def __init__(self, config: DPTConfig) -> None:
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([DPTViTLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[tuple, BaseModelOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- layer_head_mask,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions)
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-class DPTReassembleStage(nn.Module):
- """
- This class reassembles the hidden states of the backbone into image-like feature representations at various
- resolutions.
-
- This happens in 3 stages:
- 1. Map the N + 1 tokens to a set of N tokens, by taking into account the readout ([CLS]) token according to
- `config.readout_type`.
- 2. Project the channel dimension of the hidden states according to `config.neck_hidden_sizes`.
- 3. Resizing the spatial dimensions (height, width).
-
- Args:
- config (`[DPTConfig]`):
- Model configuration class defining the model architecture.
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.config = config
- self.layers = nn.ModuleList()
- if config.is_hybrid:
- self._init_reassemble_dpt_hybrid(config)
- else:
- self._init_reassemble_dpt(config)
-
- self.neck_ignore_stages = config.neck_ignore_stages
-
- def _init_reassemble_dpt_hybrid(self, config):
- r""" "
- For DPT-Hybrid the first 2 reassemble layers are set to `nn.Identity()`, please check the official
- implementation: https://github.com/isl-org/DPT/blob/f43ef9e08d70a752195028a51be5e1aff227b913/dpt/vit.py#L438
- for more details.
- """
- for i, factor in zip(range(len(config.neck_hidden_sizes)), config.reassemble_factors):
- if i <= 1:
- self.layers.append(nn.Identity())
- elif i > 1:
- self.layers.append(DPTReassembleLayer(config, channels=config.neck_hidden_sizes[i], factor=factor))
-
- if config.readout_type != "project":
- raise ValueError(f"Readout type {config.readout_type} is not supported for DPT-Hybrid.")
-
- # When using DPT-Hybrid the readout type is set to "project". The sanity check is done on the config file
- self.readout_projects = nn.ModuleList()
- hidden_size = _get_backbone_hidden_size(config)
- for i in range(len(config.neck_hidden_sizes)):
- if i <= 1:
- self.readout_projects.append(nn.Sequential(nn.Identity()))
- elif i > 1:
- self.readout_projects.append(
- nn.Sequential(nn.Linear(2 * hidden_size, hidden_size), ACT2FN[config.hidden_act])
- )
-
- def _init_reassemble_dpt(self, config):
- for i, factor in zip(range(len(config.neck_hidden_sizes)), config.reassemble_factors):
- self.layers.append(DPTReassembleLayer(config, channels=config.neck_hidden_sizes[i], factor=factor))
-
- if config.readout_type == "project":
- self.readout_projects = nn.ModuleList()
- hidden_size = _get_backbone_hidden_size(config)
- for _ in range(len(config.neck_hidden_sizes)):
- self.readout_projects.append(
- nn.Sequential(nn.Linear(2 * hidden_size, hidden_size), ACT2FN[config.hidden_act])
- )
-
- def forward(self, hidden_states: List[torch.Tensor], patch_height=None, patch_width=None) -> List[torch.Tensor]:
- """
- Args:
- hidden_states (`List[torch.FloatTensor]`, each of shape `(batch_size, sequence_length + 1, hidden_size)`):
- List of hidden states from the backbone.
- """
- out = []
-
- for i, hidden_state in enumerate(hidden_states):
- if i not in self.neck_ignore_stages:
- # reshape to (batch_size, num_channels, height, width)
- cls_token, hidden_state = hidden_state[:, 0], hidden_state[:, 1:]
- batch_size, sequence_length, num_channels = hidden_state.shape
- if patch_height is not None and patch_width is not None:
- hidden_state = hidden_state.reshape(batch_size, patch_height, patch_width, num_channels)
- else:
- size = int(math.sqrt(sequence_length))
- hidden_state = hidden_state.reshape(batch_size, size, size, num_channels)
- hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
-
- feature_shape = hidden_state.shape
- if self.config.readout_type == "project":
- # reshape to (batch_size, height*width, num_channels)
- hidden_state = hidden_state.flatten(2).permute((0, 2, 1))
- readout = cls_token.unsqueeze(1).expand_as(hidden_state)
- # concatenate the readout token to the hidden states and project
- hidden_state = self.readout_projects[i](torch.cat((hidden_state, readout), -1))
- # reshape back to (batch_size, num_channels, height, width)
- hidden_state = hidden_state.permute(0, 2, 1).reshape(feature_shape)
- elif self.config.readout_type == "add":
- hidden_state = hidden_state.flatten(2) + cls_token.unsqueeze(-1)
- hidden_state = hidden_state.reshape(feature_shape)
- hidden_state = self.layers[i](hidden_state)
- out.append(hidden_state)
-
- return out
-
-
-def _get_backbone_hidden_size(config):
- if config.backbone_config is not None and config.is_hybrid is False:
- return config.backbone_config.hidden_size
- else:
- return config.hidden_size
-
-
-class DPTReassembleLayer(nn.Module):
- def __init__(self, config, channels, factor):
- super().__init__()
- # projection
- hidden_size = _get_backbone_hidden_size(config)
- self.projection = nn.Conv2d(in_channels=hidden_size, out_channels=channels, kernel_size=1)
-
- # up/down sampling depending on factor
- if factor > 1:
- self.resize = nn.ConvTranspose2d(channels, channels, kernel_size=factor, stride=factor, padding=0)
- elif factor == 1:
- self.resize = nn.Identity()
- elif factor < 1:
- # so should downsample
- self.resize = nn.Conv2d(channels, channels, kernel_size=3, stride=int(1 / factor), padding=1)
-
- def forward(self, hidden_state):
- hidden_state = self.projection(hidden_state)
- hidden_state = self.resize(hidden_state)
- return hidden_state
-
-
-class DPTFeatureFusionStage(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.layers = nn.ModuleList()
- for _ in range(len(config.neck_hidden_sizes)):
- self.layers.append(DPTFeatureFusionLayer(config))
-
- def forward(self, hidden_states):
- # reversing the hidden_states, we start from the last
- hidden_states = hidden_states[::-1]
-
- fused_hidden_states = []
- # first layer only uses the last hidden_state
- fused_hidden_state = self.layers[0](hidden_states[0])
- fused_hidden_states.append(fused_hidden_state)
- # looping from the last layer to the second
- for hidden_state, layer in zip(hidden_states[1:], self.layers[1:]):
- fused_hidden_state = layer(fused_hidden_state, hidden_state)
- fused_hidden_states.append(fused_hidden_state)
-
- return fused_hidden_states
-
-
-class DPTPreActResidualLayer(nn.Module):
- """
- ResidualConvUnit, pre-activate residual unit.
-
- Args:
- config (`[DPTConfig]`):
- Model configuration class defining the model architecture.
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.use_batch_norm = config.use_batch_norm_in_fusion_residual
- use_bias_in_fusion_residual = (
- config.use_bias_in_fusion_residual
- if config.use_bias_in_fusion_residual is not None
- else not self.use_batch_norm
- )
-
- self.activation1 = nn.ReLU()
- self.convolution1 = nn.Conv2d(
- config.fusion_hidden_size,
- config.fusion_hidden_size,
- kernel_size=3,
- stride=1,
- padding=1,
- bias=use_bias_in_fusion_residual,
- )
-
- self.activation2 = nn.ReLU()
- self.convolution2 = nn.Conv2d(
- config.fusion_hidden_size,
- config.fusion_hidden_size,
- kernel_size=3,
- stride=1,
- padding=1,
- bias=use_bias_in_fusion_residual,
- )
-
- if self.use_batch_norm:
- self.batch_norm1 = nn.BatchNorm2d(config.fusion_hidden_size)
- self.batch_norm2 = nn.BatchNorm2d(config.fusion_hidden_size)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- residual = hidden_state
- hidden_state = self.activation1(hidden_state)
-
- hidden_state = self.convolution1(hidden_state)
-
- if self.use_batch_norm:
- hidden_state = self.batch_norm1(hidden_state)
-
- hidden_state = self.activation2(hidden_state)
- hidden_state = self.convolution2(hidden_state)
-
- if self.use_batch_norm:
- hidden_state = self.batch_norm2(hidden_state)
-
- return hidden_state + residual
-
-
-class DPTFeatureFusionLayer(nn.Module):
- """Feature fusion layer, merges feature maps from different stages.
-
- Args:
- config (`[DPTConfig]`):
- Model configuration class defining the model architecture.
- align_corners (`bool`, *optional*, defaults to `True`):
- The align_corner setting for bilinear upsample.
- """
-
- def __init__(self, config, align_corners=True):
- super().__init__()
-
- self.align_corners = align_corners
-
- self.projection = nn.Conv2d(config.fusion_hidden_size, config.fusion_hidden_size, kernel_size=1, bias=True)
-
- self.residual_layer1 = DPTPreActResidualLayer(config)
- self.residual_layer2 = DPTPreActResidualLayer(config)
-
- def forward(self, hidden_state, residual=None):
- if residual is not None:
- if hidden_state.shape != residual.shape:
- residual = nn.functional.interpolate(
- residual, size=(hidden_state.shape[2], hidden_state.shape[3]), mode="bilinear", align_corners=False
- )
- hidden_state = hidden_state + self.residual_layer1(residual)
-
- hidden_state = self.residual_layer2(hidden_state)
- hidden_state = nn.functional.interpolate(
- hidden_state, scale_factor=2, mode="bilinear", align_corners=self.align_corners
- )
- hidden_state = self.projection(hidden_state)
-
- return hidden_state
-
-
-class DPTPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = DPTConfig
- base_model_prefix = "dpt"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-DPT_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`ViTConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-DPT_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`DPTImageProcessor.__call__`]
- for details.
-
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare DPT Model transformer outputting raw hidden-states without any specific head on top.",
- DPT_START_DOCSTRING,
-)
-class DPTModel(DPTPreTrainedModel):
- def __init__(self, config, add_pooling_layer=True):
- super().__init__(config)
- self.config = config
-
- # vit encoder
- if config.is_hybrid:
- self.embeddings = DPTViTHybridEmbeddings(config)
- else:
- self.embeddings = DPTViTEmbeddings(config)
- self.encoder = DPTViTEncoder(config)
-
- self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.pooler = DPTViTPooler(config) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- if self.config.is_hybrid:
- return self.embeddings
- else:
- return self.embeddings.patch_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(DPT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPoolingAndIntermediateActivations,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- head_mask: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPoolingAndIntermediateActivations]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- embedding_output = self.embeddings(pixel_values, return_dict=return_dict)
-
- embedding_last_hidden_states = embedding_output[0] if not return_dict else embedding_output.last_hidden_states
-
- encoder_outputs = self.encoder(
- embedding_last_hidden_states,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,)
- return head_outputs + encoder_outputs[1:] + embedding_output[1:]
-
- return BaseModelOutputWithPoolingAndIntermediateActivations(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- intermediate_activations=embedding_output.intermediate_activations,
- )
-
-
-# Copied from transformers.models.vit.modeling_vit.ViTPooler with ViT->DPT
-class DPTViTPooler(nn.Module):
- def __init__(self, config: DPTConfig):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = nn.Tanh()
-
- def forward(self, hidden_states):
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(first_token_tensor)
- pooled_output = self.activation(pooled_output)
- return pooled_output
-
-
-class DPTNeck(nn.Module):
- """
- DPTNeck. A neck is a module that is normally used between the backbone and the head. It takes a list of tensors as
- input and produces another list of tensors as output. For DPT, it includes 2 stages:
-
- * DPTReassembleStage
- * DPTFeatureFusionStage.
-
- Args:
- config (dict): config dict.
- """
-
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- # postprocessing: only required in case of a non-hierarchical backbone (e.g. ViT, BEiT)
- if config.backbone_config is not None and config.backbone_config.model_type in ["swinv2"]:
- self.reassemble_stage = None
- else:
- self.reassemble_stage = DPTReassembleStage(config)
-
- self.convs = nn.ModuleList()
- for channel in config.neck_hidden_sizes:
- self.convs.append(nn.Conv2d(channel, config.fusion_hidden_size, kernel_size=3, padding=1, bias=False))
-
- # fusion
- self.fusion_stage = DPTFeatureFusionStage(config)
-
- def forward(self, hidden_states: List[torch.Tensor], patch_height=None, patch_width=None) -> List[torch.Tensor]:
- """
- Args:
- hidden_states (`List[torch.FloatTensor]`, each of shape `(batch_size, sequence_length, hidden_size)` or `(batch_size, hidden_size, height, width)`):
- List of hidden states from the backbone.
- """
- if not isinstance(hidden_states, (tuple, list)):
- raise ValueError("hidden_states should be a tuple or list of tensors")
-
- if len(hidden_states) != len(self.config.neck_hidden_sizes):
- raise ValueError("The number of hidden states should be equal to the number of neck hidden sizes.")
-
- # postprocess hidden states
- if self.reassemble_stage is not None:
- hidden_states = self.reassemble_stage(hidden_states, patch_height, patch_width)
-
- features = [self.convs[i](feature) for i, feature in enumerate(hidden_states)]
-
- # fusion blocks
- output = self.fusion_stage(features)
-
- return output
-
-
-class DPTDepthEstimationHead(nn.Module):
- """
- Output head head consisting of 3 convolutional layers. It progressively halves the feature dimension and upsamples
- the predictions to the input resolution after the first convolutional layer (details can be found in the paper's
- supplementary material).
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.config = config
-
- self.projection = None
- if config.add_projection:
- self.projection = nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
-
- features = config.fusion_hidden_size
- self.head = nn.Sequential(
- nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
- nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
- nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
- nn.ReLU(),
- nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
- nn.ReLU(),
- )
-
- def forward(self, hidden_states: List[torch.Tensor]) -> torch.Tensor:
- # use last features
- hidden_states = hidden_states[self.config.head_in_index]
-
- if self.projection is not None:
- hidden_states = self.projection(hidden_states)
- hidden_states = nn.ReLU()(hidden_states)
-
- predicted_depth = self.head(hidden_states)
-
- predicted_depth = predicted_depth.squeeze(dim=1)
-
- return predicted_depth
-
-
-@add_start_docstrings(
- """
- DPT Model with a depth estimation head on top (consisting of 3 convolutional layers) e.g. for KITTI, NYUv2.
- """,
- DPT_START_DOCSTRING,
-)
-class DPTForDepthEstimation(DPTPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.backbone = None
- if config.backbone_config is not None and config.is_hybrid is False:
- self.backbone = load_backbone(config)
- else:
- self.dpt = DPTModel(config, add_pooling_layer=False)
-
- # Neck
- self.neck = DPTNeck(config)
-
- # Depth estimation head
- self.head = DPTDepthEstimationHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DPT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=DepthEstimatorOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- head_mask: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], DepthEstimatorOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
- Ground truth depth estimation maps for computing the loss.
-
- Returns:
-
- Examples:
- ```python
- >>> from transformers import AutoImageProcessor, DPTForDepthEstimation
- >>> import torch
- >>> import numpy as np
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("Intel/dpt-large")
- >>> model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
-
- >>> # prepare image for the model
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> with torch.no_grad():
- ... outputs = model(**inputs)
- ... predicted_depth = outputs.predicted_depth
-
- >>> # interpolate to original size
- >>> prediction = torch.nn.functional.interpolate(
- ... predicted_depth.unsqueeze(1),
- ... size=image.size[::-1],
- ... mode="bicubic",
- ... align_corners=False,
- ... )
-
- >>> # visualize the prediction
- >>> output = prediction.squeeze().cpu().numpy()
- >>> formatted = (output * 255 / np.max(output)).astype("uint8")
- >>> depth = Image.fromarray(formatted)
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-
- if self.backbone is not None:
- outputs = self.backbone.forward_with_filtered_kwargs(
- pixel_values, output_hidden_states=output_hidden_states, output_attentions=output_attentions
- )
- hidden_states = outputs.feature_maps
- else:
- outputs = self.dpt(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=True, # we need the intermediate hidden states
- return_dict=return_dict,
- )
- hidden_states = outputs.hidden_states if return_dict else outputs[1]
- # only keep certain features based on config.backbone_out_indices
- # note that the hidden_states also include the initial embeddings
- if not self.config.is_hybrid:
- hidden_states = [
- feature for idx, feature in enumerate(hidden_states[1:]) if idx in self.config.backbone_out_indices
- ]
- else:
- backbone_hidden_states = outputs.intermediate_activations if return_dict else list(outputs[-1])
- backbone_hidden_states.extend(
- feature
- for idx, feature in enumerate(hidden_states[1:])
- if idx in self.config.backbone_out_indices[2:]
- )
-
- hidden_states = backbone_hidden_states
-
- patch_height, patch_width = None, None
- if self.config.backbone_config is not None and self.config.is_hybrid is False:
- _, _, height, width = pixel_values.shape
- patch_size = self.config.backbone_config.patch_size
- patch_height = height // patch_size
- patch_width = width // patch_size
-
- hidden_states = self.neck(hidden_states, patch_height, patch_width)
-
- predicted_depth = self.head(hidden_states)
-
- loss = None
- if labels is not None:
- raise NotImplementedError("Training is not implemented yet")
-
- if not return_dict:
- if output_hidden_states:
- output = (predicted_depth,) + outputs[1:]
- else:
- output = (predicted_depth,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return DepthEstimatorOutput(
- loss=loss,
- predicted_depth=predicted_depth,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=outputs.attentions,
- )
-
-
-class DPTSemanticSegmentationHead(nn.Module):
- def __init__(self, config):
- super().__init__()
-
- self.config = config
-
- features = config.fusion_hidden_size
- self.head = nn.Sequential(
- nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False),
- nn.BatchNorm2d(features),
- nn.ReLU(),
- nn.Dropout(config.semantic_classifier_dropout),
- nn.Conv2d(features, config.num_labels, kernel_size=1),
- nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
- )
-
- def forward(self, hidden_states: List[torch.Tensor]) -> torch.Tensor:
- # use last features
- hidden_states = hidden_states[self.config.head_in_index]
-
- logits = self.head(hidden_states)
-
- return logits
-
-
-class DPTAuxiliaryHead(nn.Module):
- def __init__(self, config):
- super().__init__()
-
- features = config.fusion_hidden_size
- self.head = nn.Sequential(
- nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False),
- nn.BatchNorm2d(features),
- nn.ReLU(),
- nn.Dropout(0.1, False),
- nn.Conv2d(features, config.num_labels, kernel_size=1),
- )
-
- def forward(self, hidden_states):
- logits = self.head(hidden_states)
-
- return logits
-
-
-@add_start_docstrings(
- """
- DPT Model with a semantic segmentation head on top e.g. for ADE20k, CityScapes.
- """,
- DPT_START_DOCSTRING,
-)
-class DPTForSemanticSegmentation(DPTPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.dpt = DPTModel(config, add_pooling_layer=False)
-
- # Neck
- self.neck = DPTNeck(config)
-
- # Segmentation head(s)
- self.head = DPTSemanticSegmentationHead(config)
- self.auxiliary_head = DPTAuxiliaryHead(config) if config.use_auxiliary_head else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(DPT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=SemanticSegmenterOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], SemanticSegmenterOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
- Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy).
-
- Returns:
-
- Examples:
- ```python
- >>> from transformers import AutoImageProcessor, DPTForSemanticSegmentation
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("Intel/dpt-large-ade")
- >>> model = DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade")
-
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> outputs = model(**inputs)
- >>> logits = outputs.logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- outputs = self.dpt(
- pixel_values,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=True, # we need the intermediate hidden states
- return_dict=return_dict,
- )
-
- hidden_states = outputs.hidden_states if return_dict else outputs[1]
-
- # only keep certain features based on config.backbone_out_indices
- # note that the hidden_states also include the initial embeddings
- if not self.config.is_hybrid:
- hidden_states = [
- feature for idx, feature in enumerate(hidden_states[1:]) if idx in self.config.backbone_out_indices
- ]
- else:
- backbone_hidden_states = outputs.intermediate_activations if return_dict else list(outputs[-1])
- backbone_hidden_states.extend(
- feature for idx, feature in enumerate(hidden_states[1:]) if idx in self.config.backbone_out_indices[2:]
- )
-
- hidden_states = backbone_hidden_states
-
- hidden_states = self.neck(hidden_states=hidden_states)
-
- logits = self.head(hidden_states)
-
- auxiliary_logits = None
- if self.auxiliary_head is not None:
- auxiliary_logits = self.auxiliary_head(hidden_states[-1])
-
- loss = None
- if labels is not None:
- if self.config.num_labels == 1:
- raise ValueError("The number of labels should be greater than one")
- else:
- # upsample logits to the images' original size
- upsampled_logits = nn.functional.interpolate(
- logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
- )
- if auxiliary_logits is not None:
- upsampled_auxiliary_logits = nn.functional.interpolate(
- auxiliary_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
- )
- # compute weighted loss
- loss_fct = CrossEntropyLoss(ignore_index=self.config.semantic_loss_ignore_index)
- main_loss = loss_fct(upsampled_logits, labels)
- auxiliary_loss = loss_fct(upsampled_auxiliary_logits, labels)
- loss = main_loss + self.config.auxiliary_loss_weight * auxiliary_loss
-
- if not return_dict:
- if output_hidden_states:
- output = (logits,) + outputs[1:]
- else:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SemanticSegmenterOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/efficientformer/__init__.py b/transformers/models/efficientformer/__init__.py
deleted file mode 100644
index 25d60d1ee765efb08eaa6242530bf9e8a93fafa9..0000000000000000000000000000000000000000
--- a/transformers/models/efficientformer/__init__.py
+++ /dev/null
@@ -1,109 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_tf_available,
- is_torch_available,
- is_vision_available,
-)
-
-
-_import_structure = {
- "configuration_efficientformer": [
- "EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "EfficientFormerConfig",
- ]
-}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["image_processing_efficientformer"] = ["EfficientFormerImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_efficientformer"] = [
- "EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "EfficientFormerForImageClassification",
- "EfficientFormerForImageClassificationWithTeacher",
- "EfficientFormerModel",
- "EfficientFormerPreTrainedModel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_efficientformer"] = [
- "TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFEfficientFormerForImageClassification",
- "TFEfficientFormerForImageClassificationWithTeacher",
- "TFEfficientFormerModel",
- "TFEfficientFormerPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_efficientformer import EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, EfficientFormerConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .image_processing_efficientformer import EfficientFormerImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_efficientformer import (
- EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
- EfficientFormerForImageClassification,
- EfficientFormerForImageClassificationWithTeacher,
- EfficientFormerModel,
- EfficientFormerPreTrainedModel,
- )
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_efficientformer import (
- TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFEfficientFormerForImageClassification,
- TFEfficientFormerForImageClassificationWithTeacher,
- TFEfficientFormerModel,
- TFEfficientFormerPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/efficientformer/__pycache__/__init__.cpython-310.pyc b/transformers/models/efficientformer/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 3270b1688e04762f0482a3ebf50da4aa625b697d..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientformer/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientformer/__pycache__/configuration_efficientformer.cpython-310.pyc b/transformers/models/efficientformer/__pycache__/configuration_efficientformer.cpython-310.pyc
deleted file mode 100644
index 108966c0ee22b1e71048c4e7a23fcf0d02d0b475..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientformer/__pycache__/configuration_efficientformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientformer/__pycache__/convert_efficientformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/efficientformer/__pycache__/convert_efficientformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 1e65dbc55eb9f7291e0a28e654bc2de8ec12de14..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientformer/__pycache__/convert_efficientformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientformer/__pycache__/image_processing_efficientformer.cpython-310.pyc b/transformers/models/efficientformer/__pycache__/image_processing_efficientformer.cpython-310.pyc
deleted file mode 100644
index 9e27fa24c05c66efabb13683b9541b2c265432bc..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientformer/__pycache__/image_processing_efficientformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientformer/__pycache__/modeling_efficientformer.cpython-310.pyc b/transformers/models/efficientformer/__pycache__/modeling_efficientformer.cpython-310.pyc
deleted file mode 100644
index c8b5689453991946a548f7042d1723ba1cb7a96c..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientformer/__pycache__/modeling_efficientformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientformer/__pycache__/modeling_tf_efficientformer.cpython-310.pyc b/transformers/models/efficientformer/__pycache__/modeling_tf_efficientformer.cpython-310.pyc
deleted file mode 100644
index 9a8cbe700ae38bb018db9a6cad8052caddcf73e9..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientformer/__pycache__/modeling_tf_efficientformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientformer/configuration_efficientformer.py b/transformers/models/efficientformer/configuration_efficientformer.py
deleted file mode 100644
index 1641c90711f5d4c9e0f8619ab1ff8a5c450f9959..0000000000000000000000000000000000000000
--- a/transformers/models/efficientformer/configuration_efficientformer.py
+++ /dev/null
@@ -1,170 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" EfficientFormer model configuration"""
-
-from typing import List
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class EfficientFormerConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of an [`EfficientFormerModel`]. It is used to
- instantiate an EfficientFormer model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the EfficientFormer
- [snap-research/efficientformer-l1](https://huggingface.co/snap-research/efficientformer-l1) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- depths (`List(int)`, *optional*, defaults to `[3, 2, 6, 4]`)
- Depth of each stage.
- hidden_sizes (`List(int)`, *optional*, defaults to `[48, 96, 224, 448]`)
- Dimensionality of each stage.
- downsamples (`List(bool)`, *optional*, defaults to `[True, True, True, True]`)
- Whether or not to downsample inputs between two stages.
- dim (`int`, *optional*, defaults to 448):
- Number of channels in Meta3D layers
- key_dim (`int`, *optional*, defaults to 32):
- The size of the key in meta3D block.
- attention_ratio (`int`, *optional*, defaults to 4):
- Ratio of the dimension of the query and value to the dimension of the key in MSHA block
- resolution (`int`, *optional*, defaults to 7)
- Size of each patch
- num_hidden_layers (`int`, *optional*, defaults to 5):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 8):
- Number of attention heads for each attention layer in the 3D MetaBlock.
- mlp_expansion_ratio (`int`, *optional*, defaults to 4):
- Ratio of size of the hidden dimensionality of an MLP to the dimensionality of its input.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings and encoder.
- patch_size (`int`, *optional*, defaults to 16):
- The size (resolution) of each patch.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- pool_size (`int`, *optional*, defaults to 3):
- Kernel size of pooling layers.
- downsample_patch_size (`int`, *optional*, defaults to 3):
- The size of patches in downsampling layers.
- downsample_stride (`int`, *optional*, defaults to 2):
- The stride of convolution kernels in downsampling layers.
- downsample_pad (`int`, *optional*, defaults to 1):
- Padding in downsampling layers.
- drop_path_rate (`int`, *optional*, defaults to 0):
- Rate at which to increase dropout probability in DropPath.
- num_meta3d_blocks (`int`, *optional*, defaults to 1):
- The number of 3D MetaBlocks in the last stage.
- distillation (`bool`, *optional*, defaults to `True`):
- Whether to add a distillation head.
- use_layer_scale (`bool`, *optional*, defaults to `True`):
- Whether to scale outputs from token mixers.
- layer_scale_init_value (`float`, *optional*, defaults to 1e-5):
- Factor by which outputs from token mixers are scaled.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- image_size (`int`, *optional*, defaults to `224`):
- The size (resolution) of each image.
-
- Example:
-
- ```python
- >>> from transformers import EfficientFormerConfig, EfficientFormerModel
-
- >>> # Initializing a EfficientFormer efficientformer-l1 style configuration
- >>> configuration = EfficientFormerConfig()
-
- >>> # Initializing a EfficientFormerModel (with random weights) from the efficientformer-l3 style configuration
- >>> model = EfficientFormerModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "efficientformer"
-
- def __init__(
- self,
- depths: List[int] = [3, 2, 6, 4],
- hidden_sizes: List[int] = [48, 96, 224, 448],
- downsamples: List[bool] = [True, True, True, True],
- dim: int = 448,
- key_dim: int = 32,
- attention_ratio: int = 4,
- resolution: int = 7,
- num_hidden_layers: int = 5,
- num_attention_heads: int = 8,
- mlp_expansion_ratio: int = 4,
- hidden_dropout_prob: float = 0.0,
- patch_size: int = 16,
- num_channels: int = 3,
- pool_size: int = 3,
- downsample_patch_size: int = 3,
- downsample_stride: int = 2,
- downsample_pad: int = 1,
- drop_path_rate: float = 0.0,
- num_meta3d_blocks: int = 1,
- distillation: bool = True,
- use_layer_scale: bool = True,
- layer_scale_init_value: float = 1e-5,
- hidden_act: str = "gelu",
- initializer_range: float = 0.02,
- layer_norm_eps: float = 1e-12,
- image_size: int = 224,
- batch_norm_eps: float = 1e-05,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
-
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.hidden_sizes = hidden_sizes
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.depths = depths
- self.mlp_expansion_ratio = mlp_expansion_ratio
- self.downsamples = downsamples
- self.dim = dim
- self.key_dim = key_dim
- self.attention_ratio = attention_ratio
- self.resolution = resolution
- self.pool_size = pool_size
- self.downsample_patch_size = downsample_patch_size
- self.downsample_stride = downsample_stride
- self.downsample_pad = downsample_pad
- self.drop_path_rate = drop_path_rate
- self.num_meta3d_blocks = num_meta3d_blocks
- self.distillation = distillation
- self.use_layer_scale = use_layer_scale
- self.layer_scale_init_value = layer_scale_init_value
- self.image_size = image_size
- self.batch_norm_eps = batch_norm_eps
diff --git a/transformers/models/efficientformer/convert_efficientformer_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/efficientformer/convert_efficientformer_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index 7431cd6136a593e7bd65f33d847e6b9346abfe46..0000000000000000000000000000000000000000
--- a/transformers/models/efficientformer/convert_efficientformer_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,252 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Convert EfficientFormer checkpoints from the original repository.
-
-URL: https://github.com/snap-research/EfficientFormer
-"""
-
-import argparse
-import re
-from pathlib import Path
-
-import requests
-import torch
-from PIL import Image
-from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
-
-from transformers import (
- EfficientFormerConfig,
- EfficientFormerForImageClassificationWithTeacher,
- EfficientFormerImageProcessor,
-)
-from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling
-
-
-def rename_key(old_name, num_meta4D_last_stage):
- new_name = old_name
-
- if "patch_embed" in old_name:
- _, layer, param = old_name.split(".")
-
- if layer == "0":
- new_name = old_name.replace("0", "convolution1")
- elif layer == "1":
- new_name = old_name.replace("1", "batchnorm_before")
- elif layer == "3":
- new_name = old_name.replace("3", "convolution2")
- else:
- new_name = old_name.replace("4", "batchnorm_after")
-
- if "network" in old_name and re.search(r"\d\.\d", old_name):
- two_digit_num = r"\b\d{2}\b"
- if bool(re.search(two_digit_num, old_name)):
- match = re.search(r"\d\.\d\d.", old_name).group()
- else:
- match = re.search(r"\d\.\d.", old_name).group()
- if int(match[0]) < 6:
- trimmed_name = old_name.replace(match, "")
- trimmed_name = trimmed_name.replace("network", match[0] + ".meta4D_layers.blocks." + match[2:-1])
- new_name = "intermediate_stages." + trimmed_name
- else:
- trimmed_name = old_name.replace(match, "")
- if int(match[2]) < num_meta4D_last_stage:
- trimmed_name = trimmed_name.replace("network", "meta4D_layers.blocks." + match[2])
- else:
- layer_index = str(int(match[2]) - num_meta4D_last_stage)
- trimmed_name = trimmed_name.replace("network", "meta3D_layers.blocks." + layer_index)
- if "norm1" in old_name:
- trimmed_name = trimmed_name.replace("norm1", "layernorm1")
- elif "norm2" in old_name:
- trimmed_name = trimmed_name.replace("norm2", "layernorm2")
- elif "fc1" in old_name:
- trimmed_name = trimmed_name.replace("fc1", "linear_in")
- elif "fc2" in old_name:
- trimmed_name = trimmed_name.replace("fc2", "linear_out")
-
- new_name = "last_stage." + trimmed_name
-
- elif "network" in old_name and re.search(r".\d.", old_name):
- new_name = old_name.replace("network", "intermediate_stages")
-
- if "fc" in new_name:
- new_name = new_name.replace("fc", "convolution")
- elif ("norm1" in new_name) and ("layernorm1" not in new_name):
- new_name = new_name.replace("norm1", "batchnorm_before")
- elif ("norm2" in new_name) and ("layernorm2" not in new_name):
- new_name = new_name.replace("norm2", "batchnorm_after")
- if "proj" in new_name:
- new_name = new_name.replace("proj", "projection")
- if "dist_head" in new_name:
- new_name = new_name.replace("dist_head", "distillation_classifier")
- elif "head" in new_name:
- new_name = new_name.replace("head", "classifier")
- elif "patch_embed" in new_name:
- new_name = "efficientformer." + new_name
- elif new_name == "norm.weight" or new_name == "norm.bias":
- new_name = new_name.replace("norm", "layernorm")
- new_name = "efficientformer." + new_name
- else:
- new_name = "efficientformer.encoder." + new_name
-
- return new_name
-
-
-def convert_torch_checkpoint(checkpoint, num_meta4D_last_stage):
- for key in checkpoint.copy().keys():
- val = checkpoint.pop(key)
- checkpoint[rename_key(key, num_meta4D_last_stage)] = val
-
- return checkpoint
-
-
-# We will verify our results on a COCO image
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- image = Image.open(requests.get(url, stream=True).raw)
-
- return image
-
-
-def convert_efficientformer_checkpoint(
- checkpoint_path: Path, efficientformer_config_file: Path, pytorch_dump_path: Path, push_to_hub: bool
-):
- orig_state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
- config = EfficientFormerConfig.from_json_file(efficientformer_config_file)
- model = EfficientFormerForImageClassificationWithTeacher(config)
- model_name = "_".join(checkpoint_path.split("/")[-1].split(".")[0].split("_")[:-1])
-
- num_meta4D_last_stage = config.depths[-1] - config.num_meta3d_blocks + 1
- new_state_dict = convert_torch_checkpoint(orig_state_dict, num_meta4D_last_stage)
-
- model.load_state_dict(new_state_dict)
- model.eval()
-
- pillow_resamplings = {
- "bilinear": PILImageResampling.BILINEAR,
- "bicubic": PILImageResampling.BICUBIC,
- "nearest": PILImageResampling.NEAREST,
- }
-
- # prepare image
- image = prepare_img()
- image_size = 256
- crop_size = 224
- processor = EfficientFormerImageProcessor(
- size={"shortest_edge": image_size},
- crop_size={"height": crop_size, "width": crop_size},
- resample=pillow_resamplings["bicubic"],
- )
- pixel_values = processor(images=image, return_tensors="pt").pixel_values
-
- # original processing pipeline
- image_transforms = Compose(
- [
- Resize(image_size, interpolation=pillow_resamplings["bicubic"]),
- CenterCrop(crop_size),
- ToTensor(),
- Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD),
- ]
- )
- original_pixel_values = image_transforms(image).unsqueeze(0)
-
- assert torch.allclose(original_pixel_values, pixel_values)
-
- outputs = model(pixel_values)
- logits = outputs.logits
-
- expected_shape = (1, 1000)
-
- if "l1" in model_name:
- expected_logits = torch.Tensor(
- [-0.1312, 0.4353, -1.0499, -0.5124, 0.4183, -0.6793, -1.3777, -0.0893, -0.7358, -2.4328]
- )
- assert torch.allclose(logits[0, :10], expected_logits, atol=1e-3)
- assert logits.shape == expected_shape
- elif "l3" in model_name:
- expected_logits = torch.Tensor(
- [-1.3150, -1.5456, -1.2556, -0.8496, -0.7127, -0.7897, -0.9728, -0.3052, 0.3751, -0.3127]
- )
- assert torch.allclose(logits[0, :10], expected_logits, atol=1e-3)
- assert logits.shape == expected_shape
- elif "l7" in model_name:
- expected_logits = torch.Tensor(
- [-1.0283, -1.4131, -0.5644, -1.3115, -0.5785, -1.2049, -0.7528, 0.1992, -0.3822, -0.0878]
- )
- assert logits.shape == expected_shape
- else:
- raise ValueError(
- f"Unknown model checkpoint: {checkpoint_path}. Supported version of efficientformer are l1, l3 and l7"
- )
-
- # Save Checkpoints
- Path(pytorch_dump_path).mkdir(exist_ok=True)
- model.save_pretrained(pytorch_dump_path)
- print(f"Checkpoint successfuly converted. Model saved at {pytorch_dump_path}")
- processor.save_pretrained(pytorch_dump_path)
- print(f"Processor successfuly saved at {pytorch_dump_path}")
-
- if push_to_hub:
- print("Pushing model to the hub...")
-
- model.push_to_hub(
- repo_id=f"Bearnardd/{pytorch_dump_path}",
- commit_message="Add model",
- use_temp_dir=True,
- )
- processor.push_to_hub(
- repo_id=f"Bearnardd/{pytorch_dump_path}",
- commit_message="Add image processor",
- use_temp_dir=True,
- )
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--pytorch_model_path",
- default=None,
- type=str,
- required=True,
- help="Path to EfficientFormer pytorch checkpoint.",
- )
- parser.add_argument(
- "--config_file",
- default=None,
- type=str,
- required=True,
- help="The json file for EfficientFormer model config.",
- )
- parser.add_argument(
- "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
-
- parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")
- parser.add_argument(
- "--no-push_to_hub",
- dest="push_to_hub",
- action="store_false",
- help="Do not push model and image processor to the hub",
- )
- parser.set_defaults(push_to_hub=True)
-
- args = parser.parse_args()
- convert_efficientformer_checkpoint(
- checkpoint_path=args.pytorch_model_path,
- efficientformer_config_file=args.config_file,
- pytorch_dump_path=args.pytorch_dump_path,
- push_to_hub=args.push_to_hub,
- )
diff --git a/transformers/models/efficientformer/image_processing_efficientformer.py b/transformers/models/efficientformer/image_processing_efficientformer.py
deleted file mode 100644
index 38756f7c958f5d1441c6f2c1d4ec5987c664c7ae..0000000000000000000000000000000000000000
--- a/transformers/models/efficientformer/image_processing_efficientformer.py
+++ /dev/null
@@ -1,321 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for EfficientFormer."""
-
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
-from ...image_transforms import (
- get_resize_output_image_size,
- resize,
- to_channel_dimension_format,
-)
-from ...image_utils import (
- IMAGENET_DEFAULT_MEAN,
- IMAGENET_DEFAULT_STD,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- infer_channel_dimension_format,
- is_batched,
- is_scaled_image,
- to_numpy_array,
- valid_images,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import TensorType, logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class EfficientFormerImageProcessor(BaseImageProcessor):
- r"""
- Constructs a EfficientFormer image processor.
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the image's (height, width) dimensions to the specified `(size["height"],
- size["width"])`. Can be overridden by the `do_resize` parameter in the `preprocess` method.
- size (`dict`, *optional*, defaults to `{"height": 224, "width": 224}`):
- Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess`
- method.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in the
- `preprocess` method.
- do_center_crop (`bool`, *optional*, defaults to `True`):
- Whether to center crop the image to the specified `crop_size`. Can be overridden by `do_center_crop` in the
- `preprocess` method.
- crop_size (`Dict[str, int]` *optional*, defaults to 224):
- Size of the output image after applying `center_crop`. Can be overridden by `crop_size` in the `preprocess`
- method.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`
- parameter in the `preprocess` method.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the
- `preprocess` method.
- do_normalize:
- Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
- method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
- Mean to use if normalizing the image. This is a float or list of floats the length of the number of
- channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
- Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
- number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
- """
-
- model_input_names = ["pixel_values"]
-
- def __init__(
- self,
- do_resize: bool = True,
- size: Optional[Dict[str, int]] = None,
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- do_center_crop: bool = True,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- crop_size: Dict[str, int] = None,
- do_normalize: bool = True,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- size = size if size is not None else {"height": 224, "width": 224}
- size = get_size_dict(size)
- crop_size = crop_size if crop_size is not None else {"height": 224, "width": 224}
- crop_size = get_size_dict(crop_size, default_to_square=True, param_name="crop_size")
-
- self.do_resize = do_resize
- self.do_rescale = do_rescale
- self.do_normalize = do_normalize
- self.do_center_crop = do_center_crop
- self.crop_size = crop_size
- self.size = size
- self.resample = resample
- self.rescale_factor = rescale_factor
- self.image_mean = image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size",
- "resample",
- "do_center_crop",
- "crop_size",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "image_mean",
- "image_std",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize an image to `(size["height"], size["width"])`.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
- resample:
- `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BILINEAR`.
- data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
-
- Returns:
- `np.ndarray`: The resized image.
- """
- size = get_size_dict(size)
-
- if "shortest_edge" in size:
- size = get_resize_output_image_size(
- image, size=size["shortest_edge"], default_to_square=False, input_data_format=input_data_format
- )
- # size = get_resize_output_image_size(image, size["shortest_edge"], size["longest_edge"])
- elif "height" in size and "width" in size:
- size = (size["height"], size["width"])
- else:
- raise ValueError(f"Size must contain 'height' and 'width' keys or 'shortest_edge' key. Got {size.keys()}")
- return resize(
- image, size=size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs
- )
-
- def preprocess(
- self,
- images: ImageInput,
- do_resize: Optional[bool] = None,
- size: Dict[str, int] = None,
- resample: PILImageResampling = None,
- do_center_crop: bool = None,
- crop_size: int = None,
- do_rescale: Optional[bool] = None,
- rescale_factor: Optional[float] = None,
- do_normalize: Optional[bool] = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> BatchFeature:
- """
- Preprocess an image or batch of images.
-
- Args:
- images (`ImageInput`):
- Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
- passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
- Dictionary in the format `{"height": h, "width": w}` specifying the size of the output image after
- resizing.
- resample (`PILImageResampling` filter, *optional*, defaults to `self.resample`):
- `PILImageResampling` filter to use if resizing the image e.g. `PILImageResampling.BILINEAR`. Only has
- an effect if `do_resize` is set to `True`.
- do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`):
- Whether to center crop the image.
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether to rescale the image values between [0 - 1].
- rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
- Rescale factor to rescale the image by if `do_rescale` is set to `True`.
- crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
- Size of the center crop. Only has an effect if `do_center_crop` is set to `True`.
- do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
- Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
- Image mean to use if `do_normalize` is set to `True`.
- image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
- Image standard deviation to use if `do_normalize` is set to `True`.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - Unset: Use the channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- do_resize = do_resize if do_resize is not None else self.do_resize
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- do_normalize = do_normalize if do_normalize is not None else self.do_normalize
- do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
- crop_size = crop_size if crop_size is not None else self.crop_size
- crop_size = get_size_dict(crop_size, param_name="crop_size", default_to_square=True)
- resample = resample if resample is not None else self.resample
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- image_mean = image_mean if image_mean is not None else self.image_mean
- image_std = image_std if image_std is not None else self.image_std
-
- size = size if size is not None else self.size
- size_dict = get_size_dict(size)
-
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- if not is_batched(images):
- images = [images]
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_center_crop=do_center_crop,
- crop_size=crop_size,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
- # All transformations expect numpy arrays.
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- if do_resize:
- images = [
- self.resize(image=image, size=size_dict, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_center_crop:
- images = [
- self.center_crop(image=image, size=crop_size, input_data_format=input_data_format) for image in images
- ]
-
- if do_rescale:
- images = [
- self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_normalize:
- images = [
- self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
- for image in images
- ]
-
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
- ]
-
- data = {"pixel_values": images}
- return BatchFeature(data=data, tensor_type=return_tensors)
diff --git a/transformers/models/efficientformer/modeling_efficientformer.py b/transformers/models/efficientformer/modeling_efficientformer.py
deleted file mode 100644
index 70075cff55d7d9838e5757e21302af9aa650cad7..0000000000000000000000000000000000000000
--- a/transformers/models/efficientformer/modeling_efficientformer.py
+++ /dev/null
@@ -1,803 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Snapchat Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch EfficientFormer model."""
-
-import itertools
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_efficientformer import EfficientFormerConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "EfficientFormerConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "snap-research/efficientformer-l1-300"
-_EXPECTED_OUTPUT_SHAPE = [1, 49, 448]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "snap-research/efficientformer-l1-300"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "Egyptian cat"
-
-
-from ..deprecated._archive_maps import EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class EfficientFormerPatchEmbeddings(nn.Module):
- """
- This class performs downsampling between two stages. For the input tensor with the shape [batch_size, num_channels,
- height, width] it produces output tensor with the shape [batch_size, num_channels, height/stride, width/stride]
- """
-
- def __init__(self, config: EfficientFormerConfig, num_channels: int, embed_dim: int, apply_norm: bool = True):
- super().__init__()
- self.num_channels = num_channels
-
- self.projection = nn.Conv2d(
- num_channels,
- embed_dim,
- kernel_size=config.downsample_patch_size,
- stride=config.downsample_stride,
- padding=config.downsample_pad,
- )
- self.norm = nn.BatchNorm2d(embed_dim, eps=config.batch_norm_eps) if apply_norm else nn.Identity()
-
- def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
- batch_size, num_channels, height, width = pixel_values.shape
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
-
- embeddings = self.projection(pixel_values)
- embeddings = self.norm(embeddings)
-
- return embeddings
-
-
-class EfficientFormerSelfAttention(nn.Module):
- def __init__(self, dim: int, key_dim: int, num_heads: int, attention_ratio: int, resolution: int):
- super().__init__()
-
- self.num_heads = num_heads
- self.key_dim = key_dim
- self.attention_ratio = attention_ratio
- self.scale = key_dim**-0.5
- self.total_key_dim = key_dim * num_heads
- self.expanded_key_dim = int(attention_ratio * key_dim)
- self.total_expanded_key_dim = int(self.expanded_key_dim * num_heads)
- hidden_size = self.total_expanded_key_dim + self.total_key_dim * 2
- self.qkv = nn.Linear(dim, hidden_size)
- self.projection = nn.Linear(self.total_expanded_key_dim, dim)
- points = list(itertools.product(range(resolution), range(resolution)))
- num_points = len(points)
- attention_offsets = {}
- idxs = []
- for point_1 in points:
- for point_2 in points:
- offset = (abs(point_1[0] - point_2[0]), abs(point_1[1] - point_2[1]))
- if offset not in attention_offsets:
- attention_offsets[offset] = len(attention_offsets)
- idxs.append(attention_offsets[offset])
- self.attention_biases = torch.nn.Parameter(torch.zeros(num_heads, len(attention_offsets)))
- self.register_buffer("attention_bias_idxs", torch.LongTensor(idxs).view(num_points, num_points))
-
- @torch.no_grad()
- def train(self, mode=True):
- super().train(mode)
- if mode and hasattr(self, "ab"):
- del self.ab
- else:
- self.ab = self.attention_biases[:, self.attention_bias_idxs]
-
- def forward(self, hidden_states: torch.Tensor, output_attentions: bool = False) -> Tuple[torch.Tensor]:
- batch_size, sequence_length, num_channels = hidden_states.shape
- qkv = self.qkv(hidden_states)
- query_layer, key_layer, value_layer = qkv.reshape(batch_size, sequence_length, self.num_heads, -1).split(
- [self.key_dim, self.key_dim, self.expanded_key_dim], dim=3
- )
- query_layer = query_layer.permute(0, 2, 1, 3)
- key_layer = key_layer.permute(0, 2, 1, 3)
- value_layer = value_layer.permute(0, 2, 1, 3)
-
- # set `model.to(torch_device)` won't change `self.ab.device`, if there is no follow-up `train` or `eval` call.
- # Let's do it manually here, so users won't have to do this everytime.
- if not self.training:
- self.ab = self.ab.to(self.attention_biases.device)
- attention_probs = (torch.matmul(query_layer, key_layer.transpose(-2, -1))) * self.scale + (
- self.attention_biases[:, self.attention_bias_idxs] if self.training else self.ab
- )
-
- attention_probs = attention_probs.softmax(dim=-1)
-
- context_layer = torch.matmul(attention_probs, value_layer).transpose(1, 2)
- context_layer = context_layer.reshape(batch_size, sequence_length, self.total_expanded_key_dim)
- context_layer = self.projection(context_layer)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-class EfficientFormerConvStem(nn.Module):
- def __init__(self, config: EfficientFormerConfig, out_channels: int):
- super().__init__()
-
- self.convolution1 = nn.Conv2d(config.num_channels, out_channels // 2, kernel_size=3, stride=2, padding=1)
- self.batchnorm_before = nn.BatchNorm2d(out_channels // 2, eps=config.batch_norm_eps)
-
- self.convolution2 = nn.Conv2d(out_channels // 2, out_channels, kernel_size=3, stride=2, padding=1)
- self.batchnorm_after = nn.BatchNorm2d(out_channels, eps=config.batch_norm_eps)
-
- self.activation = nn.ReLU()
-
- def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
- features = self.batchnorm_before(self.convolution1(pixel_values))
- features = self.activation(features)
- features = self.batchnorm_after(self.convolution2(features))
- features = self.activation(features)
-
- return features
-
-
-class EfficientFormerPooling(nn.Module):
- def __init__(self, pool_size: int):
- super().__init__()
- self.pool = nn.AvgPool2d(pool_size, stride=1, padding=pool_size // 2, count_include_pad=False)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- output = self.pool(hidden_states) - hidden_states
- return output
-
-
-class EfficientFormerDenseMlp(nn.Module):
- def __init__(
- self,
- config: EfficientFormerConfig,
- in_features: int,
- hidden_features: Optional[int] = None,
- out_features: Optional[int] = None,
- ):
- super().__init__()
- out_features = out_features or in_features
- hidden_features = hidden_features or in_features
-
- self.linear_in = nn.Linear(in_features, hidden_features)
- self.activation = ACT2FN[config.hidden_act]
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.linear_out = nn.Linear(hidden_features, out_features)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.linear_in(hidden_states)
- hidden_states = self.activation(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.linear_out(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-class EfficientFormerConvMlp(nn.Module):
- def __init__(
- self,
- config: EfficientFormerConfig,
- in_features: int,
- hidden_features: Optional[int] = None,
- out_features: Optional[int] = None,
- drop: float = 0.0,
- ):
- super().__init__()
- out_features = out_features or in_features
- hidden_features = hidden_features or in_features
-
- self.convolution1 = nn.Conv2d(in_features, hidden_features, 1)
- self.activation = ACT2FN[config.hidden_act]
- self.convolution2 = nn.Conv2d(hidden_features, out_features, 1)
- self.dropout = nn.Dropout(drop)
-
- self.batchnorm_before = nn.BatchNorm2d(hidden_features, eps=config.batch_norm_eps)
- self.batchnorm_after = nn.BatchNorm2d(out_features, eps=config.batch_norm_eps)
-
- def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
- hidden_state = self.convolution1(hidden_state)
- hidden_state = self.batchnorm_before(hidden_state)
-
- hidden_state = self.activation(hidden_state)
- hidden_state = self.dropout(hidden_state)
- hidden_state = self.convolution2(hidden_state)
-
- hidden_state = self.batchnorm_after(hidden_state)
- hidden_state = self.dropout(hidden_state)
-
- return hidden_state
-
-
-# Copied from transformers.models.convnext.modeling_convnext.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->EfficientFormer
-class EfficientFormerDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-class EfficientFormerFlat(nn.Module):
- def __init__(self):
- super().__init__()
-
- def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor]:
- hidden_states = hidden_states.flatten(2).transpose(1, 2)
- return hidden_states
-
-
-class EfficientFormerMeta3D(nn.Module):
- def __init__(self, config: EfficientFormerConfig, dim: int, drop_path: float = 0.0):
- super().__init__()
-
- self.token_mixer = EfficientFormerSelfAttention(
- dim=config.dim,
- key_dim=config.key_dim,
- num_heads=config.num_attention_heads,
- attention_ratio=config.attention_ratio,
- resolution=config.resolution,
- )
-
- self.layernorm1 = nn.LayerNorm(dim, eps=config.layer_norm_eps)
- self.layernorm2 = nn.LayerNorm(dim, eps=config.layer_norm_eps)
-
- mlp_hidden_dim = int(dim * config.mlp_expansion_ratio)
- self.mlp = EfficientFormerDenseMlp(config, in_features=dim, hidden_features=mlp_hidden_dim)
-
- self.drop_path = EfficientFormerDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
- self.use_layer_scale = config.use_layer_scale
- if config.use_layer_scale:
- self.layer_scale_1 = nn.Parameter(config.layer_scale_init_value * torch.ones((dim)), requires_grad=True)
- self.layer_scale_2 = nn.Parameter(config.layer_scale_init_value * torch.ones((dim)), requires_grad=True)
-
- def forward(self, hidden_states: torch.Tensor, output_attentions: bool = False) -> Tuple[torch.Tensor]:
- self_attention_outputs = self.token_mixer(self.layernorm1(hidden_states), output_attentions)
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- if self.use_layer_scale:
- layer_output = hidden_states + self.drop_path(
- self.layer_scale_1.unsqueeze(0).unsqueeze(0) * attention_output
- )
- layer_output = layer_output + self.drop_path(
- self.layer_scale_2.unsqueeze(0).unsqueeze(0) * self.mlp(self.layernorm2(layer_output))
- )
- else:
- layer_output = hidden_states + self.drop_path(attention_output)
- layer_output = layer_output + self.drop_path(self.mlp(self.layernorm2(layer_output)))
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-class EfficientFormerMeta3DLayers(nn.Module):
- def __init__(self, config: EfficientFormerConfig):
- super().__init__()
- drop_paths = [
- config.drop_path_rate * (block_idx + sum(config.depths[:-1]))
- for block_idx in range(config.num_meta3d_blocks)
- ]
- self.blocks = nn.ModuleList(
- [EfficientFormerMeta3D(config, config.hidden_sizes[-1], drop_path=drop_path) for drop_path in drop_paths]
- )
-
- def forward(self, hidden_states: torch.Tensor, output_attentions: bool = False) -> Tuple[torch.Tensor]:
- all_attention_outputs = () if output_attentions else None
-
- for layer_module in self.blocks:
- if isinstance(hidden_states, tuple):
- hidden_states = hidden_states[0]
-
- hidden_states = layer_module(hidden_states, output_attentions)
-
- if output_attentions:
- all_attention_outputs = all_attention_outputs + (hidden_states[1],)
-
- if output_attentions:
- outputs = (hidden_states[0],) + all_attention_outputs
- return outputs
-
- return hidden_states
-
-
-class EfficientFormerMeta4D(nn.Module):
- def __init__(self, config: EfficientFormerConfig, dim: int, drop_path: float = 0.0):
- super().__init__()
- pool_size = config.pool_size if config.pool_size is not None else 3
- self.token_mixer = EfficientFormerPooling(pool_size=pool_size)
- mlp_hidden_dim = int(dim * config.mlp_expansion_ratio)
- self.mlp = EfficientFormerConvMlp(
- config, in_features=dim, hidden_features=mlp_hidden_dim, drop=config.hidden_dropout_prob
- )
-
- self.drop_path = EfficientFormerDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
- self.use_layer_scale = config.use_layer_scale
- if config.use_layer_scale:
- self.layer_scale_1 = nn.Parameter(config.layer_scale_init_value * torch.ones((dim)), requires_grad=True)
- self.layer_scale_2 = nn.Parameter(config.layer_scale_init_value * torch.ones((dim)), requires_grad=True)
-
- def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor]:
- outputs = self.token_mixer(hidden_states)
-
- if self.use_layer_scale:
- layer_output = hidden_states + self.drop_path(self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * outputs)
-
- layer_output = layer_output + self.drop_path(
- self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(layer_output)
- )
- else:
- layer_output = hidden_states + self.drop_path(outputs)
- layer_output = layer_output + self.drop_path(self.mlp(layer_output))
-
- return layer_output
-
-
-class EfficientFormerMeta4DLayers(nn.Module):
- def __init__(self, config: EfficientFormerConfig, stage_idx: int):
- super().__init__()
- num_layers = (
- config.depths[stage_idx] if stage_idx != -1 else config.depths[stage_idx] - config.num_meta3d_blocks
- )
- drop_paths = [
- config.drop_path_rate * (block_idx + sum(config.depths[:stage_idx])) for block_idx in range(num_layers)
- ]
-
- self.blocks = nn.ModuleList(
- [
- EfficientFormerMeta4D(config, config.hidden_sizes[stage_idx], drop_path=drop_path)
- for drop_path in drop_paths
- ]
- )
-
- def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor]:
- for layer_module in self.blocks:
- hidden_states = layer_module(hidden_states)
- return hidden_states
-
-
-class EfficientFormerIntermediateStage(nn.Module):
- def __init__(self, config: EfficientFormerConfig, index: int):
- super().__init__()
- self.meta4D_layers = EfficientFormerMeta4DLayers(config, index)
-
- def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor]:
- hidden_states = self.meta4D_layers(hidden_states)
- return hidden_states
-
-
-class EfficientFormerLastStage(nn.Module):
- def __init__(self, config: EfficientFormerConfig):
- super().__init__()
- self.meta4D_layers = EfficientFormerMeta4DLayers(config, -1)
- self.flat = EfficientFormerFlat()
- self.meta3D_layers = EfficientFormerMeta3DLayers(config)
-
- def forward(self, hidden_states: torch.Tensor, output_attentions: bool = False) -> Tuple[torch.Tensor]:
- hidden_states = self.meta4D_layers(hidden_states)
- hidden_states = self.flat(hidden_states)
- hidden_states = self.meta3D_layers(hidden_states, output_attentions)
-
- return hidden_states
-
-
-class EfficientFormerEncoder(nn.Module):
- def __init__(self, config: EfficientFormerConfig):
- super().__init__()
- self.config = config
- num_intermediate_stages = len(config.depths) - 1
- downsamples = [
- config.downsamples[i] or config.hidden_sizes[i] != config.hidden_sizes[i + 1]
- for i in range(num_intermediate_stages)
- ]
- intermediate_stages = []
-
- for i in range(num_intermediate_stages):
- intermediate_stages.append(EfficientFormerIntermediateStage(config, i))
- if downsamples[i]:
- intermediate_stages.append(
- EfficientFormerPatchEmbeddings(config, config.hidden_sizes[i], config.hidden_sizes[i + 1])
- )
-
- self.intermediate_stages = nn.ModuleList(intermediate_stages)
- self.last_stage = EfficientFormerLastStage(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- output_hidden_states: bool = False,
- output_attentions: bool = False,
- return_dict: bool = True,
- ) -> BaseModelOutput:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- for layer_module in self.intermediate_stages:
- hidden_states = layer_module(hidden_states)
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_output = self.last_stage(hidden_states, output_attentions=output_attentions)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + layer_output[1:]
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (layer_output[0],)
-
- if not return_dict:
- return tuple(v for v in [layer_output[0], all_hidden_states, all_self_attentions] if v is not None)
-
- return BaseModelOutput(
- last_hidden_state=layer_output[0],
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-class EfficientFormerPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = EfficientFormerConfig
- base_model_prefix = "efficientformer"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = False
-
- def _init_weights(self, module: nn.Module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-EFFICIENTFORMER_START_DOCSTRING = r"""
- This model is a PyTorch [nn.Module](https://pytorch.org/docs/stable/nn.html#nn.Module) subclass. Use it as a
- regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior.
-
- Parameters:
- config ([`EfficientFormerConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-EFFICIENTFORMER_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`ViTImageProcessor`]. See
- [`ViTImageProcessor.preprocess`] for details.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare EfficientFormer Model transformer outputting raw hidden-states without any specific head on top.",
- EFFICIENTFORMER_START_DOCSTRING,
-)
-class EfficientFormerModel(EfficientFormerPreTrainedModel):
- def __init__(self, config: EfficientFormerConfig):
- super().__init__(config)
- self.config = config
-
- self.patch_embed = EfficientFormerConvStem(config, config.hidden_sizes[0])
- self.encoder = EfficientFormerEncoder(config)
- self.layernorm = nn.LayerNorm(config.hidden_sizes[-1], eps=config.layer_norm_eps)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(EFFICIENTFORMER_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPooling,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- embedding_output = self.patch_embed(pixel_values)
- encoder_outputs = self.encoder(
- embedding_output, output_attentions=output_attentions, output_hidden_states=output_hidden_states
- )
-
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
-
- if not return_dict:
- head_outputs = (sequence_output,)
- return head_outputs + encoder_outputs[1:]
-
- return BaseModelOutput(
- last_hidden_state=sequence_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- EfficientFormer Model transformer with an image classification head on top (a linear layer on top of the final
- hidden state of the [CLS] token) e.g. for ImageNet.
- """,
- EFFICIENTFORMER_START_DOCSTRING,
-)
-class EfficientFormerForImageClassification(EfficientFormerPreTrainedModel):
- def __init__(self, config: EfficientFormerConfig):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.efficientformer = EfficientFormerModel(config)
-
- # Classifier head
- self.classifier = (
- nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(EFFICIENTFORMER_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=ImageClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, ImageClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.efficientformer(
- pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.classifier(sequence_output.mean(-2))
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@dataclass
-class EfficientFormerForImageClassificationWithTeacherOutput(ModelOutput):
- """
- Output type of [`EfficientFormerForImageClassificationWithTeacher`].
-
- Args:
- logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores as the average of the cls_logits and distillation logits.
- cls_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores of the classification head (i.e. the linear layer on top of the final hidden state of the
- class token).
- distillation_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores of the distillation head (i.e. the linear layer on top of the final hidden state of the
- distillation token).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
- the self-attention heads.
- """
-
- logits: torch.FloatTensor = None
- cls_logits: torch.FloatTensor = None
- distillation_logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@add_start_docstrings(
- """
- EfficientFormer Model transformer with image classification heads on top (a linear layer on top of the final hidden
- state of the [CLS] token and a linear layer on top of the final hidden state of the distillation token) e.g. for
- ImageNet.
-
-
-
- This model supports inference-only. Fine-tuning with distillation (i.e. with a teacher) is not yet
- supported.
-
-
- """,
- EFFICIENTFORMER_START_DOCSTRING,
-)
-class EfficientFormerForImageClassificationWithTeacher(EfficientFormerPreTrainedModel):
- def __init__(self, config: EfficientFormerConfig):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.efficientformer = EfficientFormerModel(config)
-
- # Classifier head
- self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
- # Distillation head
- self.distillation_classifier = (
- nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(EFFICIENTFORMER_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=EfficientFormerForImageClassificationWithTeacherOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, EfficientFormerForImageClassificationWithTeacherOutput]:
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- outputs = self.efficientformer(
- pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- cls_logits = self.classifier(sequence_output.mean(-2))
- distillation_logits = self.distillation_classifier(sequence_output.mean(-2))
-
- # during inference, return the average of both classifier predictions
- logits = (cls_logits + distillation_logits) / 2
-
- if not return_dict:
- output = (logits, cls_logits, distillation_logits) + outputs[1:]
- return output
-
- return EfficientFormerForImageClassificationWithTeacherOutput(
- logits=logits,
- cls_logits=cls_logits,
- distillation_logits=distillation_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/efficientformer/modeling_tf_efficientformer.py b/transformers/models/efficientformer/modeling_tf_efficientformer.py
deleted file mode 100644
index 77b62999e772ecc3d3ebd119da7fa8e8d80ff4be..0000000000000000000000000000000000000000
--- a/transformers/models/efficientformer/modeling_tf_efficientformer.py
+++ /dev/null
@@ -1,1193 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Snapchat Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TensorFlow EfficientFormer model."""
-
-import itertools
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import tensorflow as tf
-
-from ...activations_tf import ACT2FN
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFBaseModelOutputWithPooling,
- TFImageClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import shape_list, stable_softmax
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_efficientformer import EfficientFormerConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "EfficientFormerConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "snap-research/efficientformer-l1-300"
-_EXPECTED_OUTPUT_SHAPE = [1, 49, 448]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "snap-research/efficientformer-l1-300"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "LABEL_281"
-
-
-from ..deprecated._archive_maps import TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class TFEfficientFormerPatchEmbeddings(keras.layers.Layer):
- """
- This class performs downsampling between two stages. For the input tensor with the shape [batch_size, num_channels,
- height, width] it produces output tensor with the shape [batch_size, num_channels, height/stride, width/stride]
- """
-
- def __init__(
- self, config: EfficientFormerConfig, num_channels: int, embed_dim: int, apply_norm: bool = True, **kwargs
- ) -> None:
- super().__init__(**kwargs)
- self.num_channels = num_channels
-
- self.padding = keras.layers.ZeroPadding2D(padding=config.downsample_pad)
- self.projection = keras.layers.Conv2D(
- filters=embed_dim,
- kernel_size=config.downsample_patch_size,
- strides=config.downsample_stride,
- padding="valid",
- name="projection",
- )
- # Use same default momentum and epsilon as PyTorch equivalent for BatchNormalization
- self.norm = (
- keras.layers.BatchNormalization(axis=-1, epsilon=config.batch_norm_eps, momentum=0.9, name="norm")
- if apply_norm
- else tf.identity
- )
- self.embed_dim = embed_dim
-
- def call(self, pixel_values: tf.Tensor, training: bool = False) -> tf.Tensor:
- tf.debugging.assert_shapes(
- [(pixel_values, (..., None, None, self.num_channels))],
- message="Make sure that the channel dimension of the pixel values match with the one set in the configuration.",
- )
- embeddings = self.projection(self.padding(pixel_values))
- embeddings = self.norm(embeddings, training=training)
- return embeddings
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "projection", None) is not None:
- with tf.name_scope(self.projection.name):
- self.projection.build([None, None, None, self.num_channels])
- if getattr(self, "norm", None) is not None:
- if hasattr(self.norm, "name"):
- with tf.name_scope(self.norm.name):
- self.norm.build([None, None, None, self.embed_dim])
-
-
-class TFEfficientFormerSelfAttention(keras.layers.Layer):
- def __init__(
- self,
- dim: int,
- key_dim: int,
- num_heads: int,
- attention_ratio: int,
- resolution: int,
- config: EfficientFormerConfig,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.num_heads = num_heads
- self.key_dim = key_dim
- self.attention_ratio = attention_ratio
- self.scale = key_dim**-0.5
- self.total_key_dim = key_dim * num_heads
- self.expanded_key_dim = int(attention_ratio * key_dim)
- self.total_expanded_key_dim = int(self.expanded_key_dim * num_heads)
- hidden_size = self.total_expanded_key_dim + self.total_key_dim * 2
-
- self.qkv = keras.layers.Dense(
- units=hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="qkv"
- )
- self.projection = keras.layers.Dense(
- units=dim, kernel_initializer=get_initializer(config.initializer_range), name="projection"
- )
- self.resolution = resolution
- self.dim = dim
-
- def build(self, input_shape: tf.TensorShape) -> None:
- points = list(itertools.product(range(self.resolution), range(self.resolution)))
- num_points = len(points)
- attention_offsets = {}
-
- idxs = []
-
- for point_1 in points:
- for point_2 in points:
- offset = (abs(point_1[0] - point_2[0]), abs(point_1[1] - point_2[1]))
- if offset not in attention_offsets:
- attention_offsets[offset] = len(attention_offsets)
- idxs.append(attention_offsets[offset])
-
- self.attention_biases = self.add_weight(
- shape=(self.num_heads, len(attention_offsets)),
- initializer=keras.initializers.zeros(),
- trainable=True,
- name="attention_biases",
- )
- self.attention_bias_idxs = self.add_weight(
- shape=(num_points, num_points),
- trainable=False,
- dtype=tf.int32,
- name="attention_bias_idxs",
- )
-
- self.attention_bias_idxs.assign(tf.reshape(tf.cast(idxs, dtype=tf.int32), (num_points, num_points)))
-
- if self.built:
- return
- self.built = True
- if getattr(self, "qkv", None) is not None:
- with tf.name_scope(self.qkv.name):
- self.qkv.build([None, None, self.dim])
- if getattr(self, "projection", None) is not None:
- with tf.name_scope(self.projection.name):
- self.projection.build([None, None, self.total_expanded_key_dim])
-
- def call(
- self, hidden_states: tf.Tensor, output_attentions: bool = False, training: bool = False
- ) -> Tuple[tf.Tensor]:
- batch_size, sequence_length, *_ = shape_list(hidden_states)
- qkv = self.qkv(inputs=hidden_states)
-
- query_layer, key_layer, value_layer = tf.split(
- tf.reshape(tensor=qkv, shape=(batch_size, sequence_length, self.num_heads, -1)),
- num_or_size_splits=[self.key_dim, self.key_dim, self.expanded_key_dim],
- axis=3,
- )
-
- query_layer = tf.transpose(query_layer, perm=[0, 2, 1, 3])
- key_layer = tf.transpose(key_layer, perm=[0, 2, 1, 3])
- value_layer = tf.transpose(value_layer, perm=[0, 2, 1, 3])
-
- attention_probs = tf.matmul(query_layer, tf.transpose(key_layer, perm=[0, 1, 3, 2]))
- scale = tf.cast(self.scale, dtype=attention_probs.dtype)
- attention_probs = tf.multiply(attention_probs, scale)
-
- attention_biases = tf.gather(params=self.attention_biases, indices=self.attention_bias_idxs, axis=1)
- attention_probs = attention_probs + attention_biases
- attention_probs = stable_softmax(logits=attention_probs, axis=-1)
-
- context_layer = tf.matmul(attention_probs, value_layer)
- context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
-
- context_layer = tf.reshape(
- tensor=context_layer, shape=(batch_size, sequence_length, self.total_expanded_key_dim)
- )
- context_layer = self.projection(context_layer)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-class TFEfficientFormerConvStem(keras.layers.Layer):
- def __init__(self, config: EfficientFormerConfig, out_channels: int, **kwargs):
- super().__init__(**kwargs)
-
- self.padding = keras.layers.ZeroPadding2D(padding=1)
- self.convolution1 = keras.layers.Conv2D(
- filters=out_channels // 2, kernel_size=3, strides=2, padding="valid", name="convolution1"
- )
- # Use same default momentum and epsilon as PyTorch equivalent for BatchNormalization
- self.batchnorm_before = keras.layers.BatchNormalization(
- axis=-1, epsilon=config.batch_norm_eps, momentum=0.9, name="batchnorm_before"
- )
-
- self.convolution2 = keras.layers.Conv2D(
- filters=out_channels,
- kernel_size=3,
- strides=2,
- padding="valid",
- name="convolution2",
- )
- # Use same default momentum and epsilon as PyTorch equivalent for BatchNormalization
- self.batchnorm_after = keras.layers.BatchNormalization(
- axis=-1, epsilon=config.batch_norm_eps, momentum=0.9, name="batchnorm_after"
- )
-
- self.activation = keras.layers.Activation(activation=keras.activations.relu, name="activation")
- self.out_channels = out_channels
- self.config = config
-
- def call(self, pixel_values: tf.Tensor, training: bool = False) -> tf.Tensor:
- features = self.batchnorm_before(self.convolution1(self.padding(pixel_values)), training=training)
- features = self.activation(features)
- features = self.batchnorm_after(self.convolution2(self.padding(features)), training=training)
- features = self.activation(features)
- return features
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convolution1", None) is not None:
- with tf.name_scope(self.convolution1.name):
- self.convolution1.build([None, None, None, self.config.num_channels])
- if getattr(self, "batchnorm_before", None) is not None:
- with tf.name_scope(self.batchnorm_before.name):
- self.batchnorm_before.build([None, None, None, self.out_channels // 2])
- if getattr(self, "convolution2", None) is not None:
- with tf.name_scope(self.convolution2.name):
- self.convolution2.build([None, None, None, self.out_channels // 2])
- if getattr(self, "batchnorm_after", None) is not None:
- with tf.name_scope(self.batchnorm_after.name):
- self.batchnorm_after.build([None, None, None, self.out_channels])
- if getattr(self, "activation", None) is not None:
- with tf.name_scope(self.activation.name):
- self.activation.build(None)
-
-
-class TFEfficientFormerPooling(keras.layers.Layer):
- def __init__(self, pool_size: int, **kwargs):
- super().__init__(**kwargs)
- self.pool = keras.layers.AveragePooling2D(pool_size=pool_size, strides=1, padding="same")
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- output = self.pool(hidden_states)
- output = output - hidden_states
- return output
-
-
-class TFEfficientFormerDenseMlp(keras.layers.Layer):
- def __init__(
- self,
- config: EfficientFormerConfig,
- in_features: int,
- hidden_features: Optional[int] = None,
- out_features: Optional[int] = None,
- **kwargs,
- ):
- super().__init__(**kwargs)
- out_features = out_features or in_features
- hidden_features = hidden_features or in_features
-
- self.linear_in = keras.layers.Dense(
- units=hidden_features, kernel_initializer=get_initializer(config.initializer_range), name="linear_in"
- )
- self.activation = ACT2FN[config.hidden_act]
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
-
- self.linear_out = keras.layers.Dense(
- units=out_features, kernel_initializer=get_initializer(config.initializer_range), name="linear_out"
- )
- self.hidden_features = hidden_features
- self.in_features = in_features
-
- def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = self.linear_in(inputs=hidden_states)
- hidden_states = self.activation(hidden_states)
- hidden_states = self.dropout(inputs=hidden_states, training=training)
- hidden_states = self.linear_out(inputs=hidden_states)
- hidden_states = self.dropout(inputs=hidden_states, training=training)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "linear_in", None) is not None:
- with tf.name_scope(self.linear_in.name):
- self.linear_in.build([None, None, self.in_features])
- if getattr(self, "linear_out", None) is not None:
- with tf.name_scope(self.linear_out.name):
- self.linear_out.build([None, None, self.hidden_features])
-
-
-class TFEfficientFormerConvMlp(keras.layers.Layer):
- def __init__(
- self,
- config: EfficientFormerConfig,
- in_features: int,
- hidden_features: Optional[int] = None,
- out_features: Optional[int] = None,
- drop: float = 0.0,
- **kwargs,
- ):
- super().__init__(**kwargs)
- out_features = out_features or in_features
- hidden_features = hidden_features or in_features
-
- self.convolution1 = keras.layers.Conv2D(
- filters=hidden_features,
- kernel_size=1,
- name="convolution1",
- padding="valid",
- )
-
- self.activation = ACT2FN[config.hidden_act]
-
- self.convolution2 = keras.layers.Conv2D(
- filters=out_features,
- kernel_size=1,
- name="convolution2",
- padding="valid",
- )
-
- self.dropout = keras.layers.Dropout(rate=drop)
-
- # Use same default momentum and epsilon as PyTorch equivalent for BatchNormalization
- self.batchnorm_before = keras.layers.BatchNormalization(
- axis=-1, epsilon=config.batch_norm_eps, momentum=0.9, name="batchnorm_before"
- )
- # Use same default momentum and epsilon as PyTorch equivalent for BatchNormalization
- self.batchnorm_after = keras.layers.BatchNormalization(
- axis=-1, epsilon=config.batch_norm_eps, momentum=0.9, name="batchnorm_after"
- )
- self.hidden_features = hidden_features
- self.in_features = in_features
- self.out_features = out_features
-
- def call(self, hidden_state: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_state = self.convolution1(hidden_state)
- hidden_state = self.batchnorm_before(hidden_state, training=training)
- hidden_state = self.activation(hidden_state)
- hidden_state = self.dropout(hidden_state, training=training)
- hidden_state = self.convolution2(hidden_state)
- hidden_state = self.batchnorm_after(hidden_state, training=training)
- hidden_state = self.dropout(hidden_state, training=training)
- return hidden_state
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "convolution1", None) is not None:
- with tf.name_scope(self.convolution1.name):
- self.convolution1.build([None, None, None, self.in_features])
- if getattr(self, "convolution2", None) is not None:
- with tf.name_scope(self.convolution2.name):
- self.convolution2.build([None, None, None, self.hidden_features])
- if getattr(self, "batchnorm_before", None) is not None:
- with tf.name_scope(self.batchnorm_before.name):
- self.batchnorm_before.build([None, None, None, self.hidden_features])
- if getattr(self, "batchnorm_after", None) is not None:
- with tf.name_scope(self.batchnorm_after.name):
- self.batchnorm_after.build([None, None, None, self.out_features])
-
-
-# Copied from transformers.models.convnext.modeling_tf_convnext.TFConvNextDropPath with ConvNext->EfficientFormer
-class TFEfficientFormerDropPath(keras.layers.Layer):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
- References:
- (1) github.com:rwightman/pytorch-image-models
- """
-
- def __init__(self, drop_path: float, **kwargs):
- super().__init__(**kwargs)
- self.drop_path = drop_path
-
- def call(self, x: tf.Tensor, training=None):
- if training:
- keep_prob = 1 - self.drop_path
- shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)
- random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)
- random_tensor = tf.floor(random_tensor)
- return (x / keep_prob) * random_tensor
- return x
-
-
-class TFEfficientFormerFlat(keras.layers.Layer):
- def __init__(self, **kwargs):
- super().__init__(**kwargs)
-
- def call(self, hidden_states: tf.Tensor) -> Tuple[tf.Tensor]:
- batch_size, _, _, in_channels = shape_list(hidden_states)
- hidden_states = tf.reshape(hidden_states, shape=[batch_size, -1, in_channels])
- return hidden_states
-
-
-class TFEfficientFormerMeta3D(keras.layers.Layer):
- def __init__(self, config: EfficientFormerConfig, dim: int, drop_path: float = 0.0, **kwargs):
- super().__init__(**kwargs)
-
- self.token_mixer = TFEfficientFormerSelfAttention(
- dim=config.dim,
- key_dim=config.key_dim,
- num_heads=config.num_attention_heads,
- attention_ratio=config.attention_ratio,
- resolution=config.resolution,
- name="token_mixer",
- config=config,
- )
- self.dim = dim
- self.config = config
-
- self.layernorm1 = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm1")
- self.layernorm2 = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm2")
- mlp_hidden_dim = int(dim * config.mlp_expansion_ratio)
- self.mlp = TFEfficientFormerDenseMlp(config, in_features=dim, hidden_features=mlp_hidden_dim, name="mlp")
-
- # Using `layers.Activation` instead of `tf.identity` to better control `training' behavior.
- self.drop_path = (
- TFEfficientFormerDropPath(drop_path)
- if drop_path > 0.0
- else keras.layers.Activation("linear", name="drop_path")
- )
- self.config = config
-
- def build(self, input_shape=None):
- self.layer_scale_1 = None
- self.layer_scale_2 = None
-
- if self.config.use_layer_scale:
- self.layer_scale_1 = self.add_weight(
- shape=(self.dim,),
- initializer=keras.initializers.Constant(value=self.config.layer_scale_init_value),
- trainable=True,
- name="layer_scale_1",
- )
- self.layer_scale_2 = self.add_weight(
- shape=(self.dim,),
- initializer=keras.initializers.Constant(value=self.config.layer_scale_init_value),
- trainable=True,
- name="layer_scale_2",
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "token_mixer", None) is not None:
- with tf.name_scope(self.token_mixer.name):
- self.token_mixer.build(None)
- if getattr(self, "layernorm1", None) is not None:
- with tf.name_scope(self.layernorm1.name):
- self.layernorm1.build([None, None, self.dim])
- if getattr(self, "layernorm2", None) is not None:
- with tf.name_scope(self.layernorm2.name):
- self.layernorm2.build([None, None, self.dim])
- if getattr(self, "mlp", None) is not None:
- with tf.name_scope(self.mlp.name):
- self.mlp.build(None)
- if getattr(self, "drop_path", None) is not None:
- with tf.name_scope(self.drop_path.name):
- self.drop_path.build(None)
-
- def call(
- self, hidden_states: tf.Tensor, output_attentions: bool = False, training: bool = False
- ) -> Tuple[tf.Tensor]:
- self_attention_outputs = self.token_mixer(
- hidden_states=self.layernorm1(hidden_states, training=training),
- output_attentions=output_attentions,
- training=training,
- )
-
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- if self.config.use_layer_scale:
- layer_output = hidden_states + self.drop_path(
- tf.expand_dims(tf.expand_dims(self.layer_scale_1, 0), 0) * attention_output,
- training=training,
- )
- layer_output = layer_output + self.drop_path(
- tf.expand_dims(tf.expand_dims(self.layer_scale_2, 0), 0)
- * self.mlp(hidden_states=self.layernorm2(inputs=layer_output, training=training), training=training),
- training=training,
- )
- else:
- layer_output = hidden_states + self.drop_path(attention_output, training=training)
- layer_output = layer_output + self.drop_path(
- self.mlp(hidden_states=self.layernorm2(inputs=layer_output, training=training), training=training),
- training=training,
- )
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-class TFEfficientFormerMeta3DLayers(keras.layers.Layer):
- def __init__(self, config: EfficientFormerConfig, **kwargs):
- super().__init__(**kwargs)
- drop_paths = [
- config.drop_path_rate * (block_idx + sum(config.depths[:-1]))
- for block_idx in range(config.num_meta3d_blocks)
- ]
- self.blocks = [
- TFEfficientFormerMeta3D(config, config.hidden_sizes[-1], drop_path=drop_path, name=f"blocks.{i}")
- for i, drop_path in enumerate(drop_paths)
- ]
-
- def call(
- self, hidden_states: tf.Tensor, output_attentions: bool = False, training: bool = False
- ) -> Tuple[tf.Tensor]:
- all_attention_outputs = () if output_attentions else None
-
- for i, layer_module in enumerate(self.blocks):
- if isinstance(hidden_states, tuple):
- hidden_states = hidden_states[0]
-
- hidden_states = layer_module(
- hidden_states=hidden_states, output_attentions=output_attentions, training=training
- )
- if output_attentions:
- all_attention_outputs = all_attention_outputs + (hidden_states[1],)
-
- if output_attentions:
- outputs = (hidden_states[0],) + all_attention_outputs
- return outputs
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "blocks", None) is not None:
- for layer in self.blocks:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-class TFEfficientFormerMeta4D(keras.layers.Layer):
- def __init__(self, config: EfficientFormerConfig, dim: int, drop_path: float = 0.0, **kwargs):
- super().__init__(**kwargs)
- pool_size = config.pool_size if config.pool_size is not None else 3
- self.token_mixer = TFEfficientFormerPooling(pool_size=pool_size, name="token_mixer")
- self.dim = dim
- mlp_hidden_dim = int(dim * config.mlp_expansion_ratio)
- self.mlp = TFEfficientFormerConvMlp(
- config=config, in_features=dim, hidden_features=mlp_hidden_dim, drop=config.hidden_dropout_prob, name="mlp"
- )
-
- self.drop_path = (
- TFEfficientFormerDropPath(drop_path, name="drop_path")
- if drop_path > 0.0
- else keras.layers.Activation("linear", name="drop_path")
- )
- self.config = config
-
- def build(self, input_shape=None):
- self.layer_scale_1 = None
- self.layer_scale_2 = None
-
- if self.config.use_layer_scale:
- self.layer_scale_1 = self.add_weight(
- shape=(self.dim),
- initializer=keras.initializers.Constant(value=self.config.layer_scale_init_value),
- trainable=True,
- name="layer_scale_1",
- )
- self.layer_scale_2 = self.add_weight(
- shape=(self.dim),
- initializer=keras.initializers.Constant(value=self.config.layer_scale_init_value),
- trainable=True,
- name="layer_scale_2",
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "token_mixer", None) is not None:
- with tf.name_scope(self.token_mixer.name):
- self.token_mixer.build(None)
- if getattr(self, "mlp", None) is not None:
- with tf.name_scope(self.mlp.name):
- self.mlp.build(None)
- if getattr(self, "drop_path", None) is not None:
- with tf.name_scope(self.drop_path.name):
- self.drop_path.build(None)
-
- def call(self, hidden_states: tf.Tensor, training: bool = False) -> Tuple[tf.Tensor]:
- outputs = self.token_mixer(hidden_states)
-
- if self.config.use_layer_scale:
- layer_output = hidden_states + self.drop_path(
- tf.expand_dims(tf.expand_dims(self.layer_scale_1, 0), 0) * outputs,
- training=training,
- )
-
- layer_output = layer_output + self.drop_path(
- tf.expand_dims(tf.expand_dims(self.layer_scale_2, 0), 0)
- * self.mlp(hidden_state=layer_output, training=training),
- training=training,
- )
-
- else:
- layer_output = hidden_states + self.drop_path(outputs, training=training)
- layer_output = layer_output + self.drop_path(
- self.mlp(hidden_state=layer_output, training=training), training=training
- )
-
- return layer_output
-
-
-class TFEfficientFormerMeta4DLayers(keras.layers.Layer):
- def __init__(self, config: EfficientFormerConfig, stage_idx: int, **kwargs):
- super().__init__(**kwargs)
- num_layers = (
- config.depths[stage_idx] if stage_idx != -1 else config.depths[stage_idx] - config.num_meta3d_blocks
- )
- drop_paths = [
- config.drop_path_rate * (block_idx + sum(config.depths[:stage_idx])) for block_idx in range(num_layers)
- ]
-
- self.blocks = [
- TFEfficientFormerMeta4D(
- config=config, dim=config.hidden_sizes[stage_idx], drop_path=drop_paths[i], name=f"blocks.{i}"
- )
- for i in range(len(drop_paths))
- ]
-
- def call(self, hidden_states: tf.Tensor, training: bool = False) -> Tuple[tf.Tensor]:
- for layer_module in self.blocks:
- hidden_states = layer_module(hidden_states=hidden_states, training=training)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "blocks", None) is not None:
- for layer in self.blocks:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-class TFEfficientFormerIntermediateStage(keras.layers.Layer):
- def __init__(self, config: EfficientFormerConfig, index: int, **kwargs):
- super().__init__(**kwargs)
- self.meta4D_layers = TFEfficientFormerMeta4DLayers(config=config, stage_idx=index, name="meta4D_layers")
-
- def call(self, hidden_states: tf.Tensor, training: bool = False) -> Tuple[tf.Tensor]:
- hidden_states = self.meta4D_layers(hidden_states=hidden_states, training=training)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "meta4D_layers", None) is not None:
- with tf.name_scope(self.meta4D_layers.name):
- self.meta4D_layers.build(None)
-
-
-class TFEfficientFormerLastStage(keras.layers.Layer):
- def __init__(self, config: EfficientFormerConfig, **kwargs):
- super().__init__(**kwargs)
- self.meta4D_layers = TFEfficientFormerMeta4DLayers(config=config, stage_idx=-1, name="meta4D_layers")
- self.flat = TFEfficientFormerFlat(name="flat")
- self.meta3D_layers = TFEfficientFormerMeta3DLayers(config, name="meta3D_layers")
-
- def call(
- self, hidden_states: tf.Tensor, output_attentions: bool = False, training: bool = False
- ) -> Tuple[tf.Tensor]:
- hidden_states = self.meta4D_layers(hidden_states=hidden_states, training=training)
- hidden_states = self.flat(hidden_states=hidden_states)
- hidden_states = self.meta3D_layers(
- hidden_states=hidden_states, output_attentions=output_attentions, training=training
- )
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "meta4D_layers", None) is not None:
- with tf.name_scope(self.meta4D_layers.name):
- self.meta4D_layers.build(None)
- if getattr(self, "flat", None) is not None:
- with tf.name_scope(self.flat.name):
- self.flat.build(None)
- if getattr(self, "meta3D_layers", None) is not None:
- with tf.name_scope(self.meta3D_layers.name):
- self.meta3D_layers.build(None)
-
-
-class TFEfficientFormerEncoder(keras.layers.Layer):
- def __init__(self, config: EfficientFormerConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- num_intermediate_stages = len(config.depths) - 1
- downsamples = [
- config.downsamples[i] or config.hidden_sizes[i] != config.hidden_sizes[i + 1]
- for i in range(num_intermediate_stages)
- ]
-
- intermediate_stages = []
- layer_count = -1
- for i in range(num_intermediate_stages):
- layer_count += 1
- intermediate_stages.append(
- TFEfficientFormerIntermediateStage(config, i, name=f"intermediate_stages.{layer_count}")
- )
- if downsamples[i]:
- layer_count += 1
- intermediate_stages.append(
- TFEfficientFormerPatchEmbeddings(
- config,
- config.hidden_sizes[i],
- config.hidden_sizes[i + 1],
- name=f"intermediate_stages.{layer_count}",
- )
- )
- self.intermediate_stages = intermediate_stages
- self.last_stage = TFEfficientFormerLastStage(config, name="last_stage")
-
- def call(
- self,
- hidden_states: tf.Tensor,
- output_hidden_states: bool,
- output_attentions: bool,
- return_dict: bool,
- training: bool = False,
- ) -> TFBaseModelOutput:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- for layer_module in self.intermediate_stages:
- hidden_states = layer_module(hidden_states, training=training)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_output = self.last_stage(hidden_states, output_attentions=output_attentions, training=training)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + layer_output[1:]
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (layer_output[0],)
-
- if not return_dict:
- return tuple(v for v in [layer_output[0], all_hidden_states, all_self_attentions] if v is not None)
-
- return TFBaseModelOutput(
- last_hidden_state=layer_output[0],
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "last_stage", None) is not None:
- with tf.name_scope(self.last_stage.name):
- self.last_stage.build(None)
- for layer in self.intermediate_stages:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-@keras_serializable
-class TFEfficientFormerMainLayer(keras.layers.Layer):
- config_class = EfficientFormerConfig
-
- def __init__(self, config: EfficientFormerConfig, **kwargs) -> None:
- super().__init__(**kwargs)
- self.config = config
-
- self.patch_embed = TFEfficientFormerConvStem(config, config.hidden_sizes[0], name="patch_embed")
- self.encoder = TFEfficientFormerEncoder(config, name="encoder")
- self.layernorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")
-
- @unpack_inputs
- def call(
- self,
- pixel_values: Optional[tf.Tensor] = None,
- output_attentions: Optional[tf.Tensor] = None,
- output_hidden_states: Optional[tf.Tensor] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor, ...]]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- # When running on CPU, keras.layers.Conv2D and keras.layers.AveragePool2D do not
- # support channels first NCHW format. A number of blocks contain both.
- # So change the input format from (batch_size, num_channels, height, width) to
- # (batch_size, height, width, num_channels) here.
- # shape = (batch_size, in_height, in_width, in_channels=num_channels)
- pixel_values = tf.transpose(pixel_values, perm=(0, 2, 3, 1))
- embedding_output = self.patch_embed(pixel_values, training=training)
-
- encoder_outputs = self.encoder(
- hidden_states=embedding_output,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output, training=training)
-
- # Change the hidden states from (batch_size, height, width, num_channels) to
- # (batch_size, num_channels, height, width).
- # The hidden states are in (batch_size, height, width, num_channels)
- # shape after all stages except the MB3D blocks.
- if output_hidden_states:
- hidden_states = tuple([tf.transpose(h, perm=(0, 3, 1, 2)) for h in encoder_outputs[1][:-1]]) + (
- encoder_outputs[1][-1],
- )
-
- if not return_dict:
- head_outputs = (sequence_output,)
- return head_outputs + encoder_outputs[1:]
-
- return TFBaseModelOutput(
- last_hidden_state=sequence_output,
- hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "patch_embed", None) is not None:
- with tf.name_scope(self.patch_embed.name):
- self.patch_embed.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "layernorm", None) is not None:
- with tf.name_scope(self.layernorm.name):
- self.layernorm.build([None, None, self.config.hidden_sizes[-1]])
-
-
-class TFEfficientFormerPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = EfficientFormerConfig
- base_model_prefix = "efficientformer"
- main_input_name = "pixel_values"
-
-
-EFFICIENTFORMER_START_DOCSTRING = r"""
- This model is a TensorFlow
- [keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer). Use it as a regular
- TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and behavior.
-
-
- Parameters:
- config ([`EfficientFormerConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-EFFICIENTFORMER_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values ((`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`EfficientFormerImageProcessor.__call__`] for details.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare EfficientFormer Model transformer outputting raw hidden-states without any specific head on top.",
- EFFICIENTFORMER_START_DOCSTRING,
-)
-class TFEfficientFormerModel(TFEfficientFormerPreTrainedModel):
- def __init__(self, config: EfficientFormerConfig, **kwargs) -> None:
- super().__init__(config, **kwargs)
-
- self.efficientformer = TFEfficientFormerMainLayer(config, name="efficientformer")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(EFFICIENTFORMER_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutputWithPooling,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def call(
- self,
- pixel_values: Optional[tf.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[Tuple, TFBaseModelOutput]:
- outputs = self.efficientformer(
- pixel_values=pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "efficientformer", None) is not None:
- with tf.name_scope(self.efficientformer.name):
- self.efficientformer.build(None)
-
-
-@add_start_docstrings(
- """
- EfficientFormer Model transformer with an image classification head on top of pooled last hidden state, e.g. for
- ImageNet.
- """,
- EFFICIENTFORMER_START_DOCSTRING,
-)
-class TFEfficientFormerForImageClassification(TFEfficientFormerPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: EfficientFormerConfig):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.efficientformer = TFEfficientFormerMainLayer(config, name="efficientformer")
-
- # Classifier head
- self.classifier = (
- keras.layers.Dense(config.num_labels, name="classifier")
- if config.num_labels > 0
- else keras.layers.Activation("linear", name="classifier")
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(EFFICIENTFORMER_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=TFImageClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def call(
- self,
- pixel_values: Optional[tf.Tensor] = None,
- labels: Optional[tf.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[tf.Tensor, TFImageClassifierOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.efficientformer(
- pixel_values=pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = outputs[0]
-
- logits = self.classifier(tf.reduce_mean(sequence_output, axis=-2))
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFImageClassifierOutput(
- loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "efficientformer", None) is not None:
- with tf.name_scope(self.efficientformer.name):
- self.efficientformer.build(None)
- if getattr(self, "classifier", None) is not None:
- if hasattr(self.classifier, "name"):
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_sizes[-1]])
-
-
-@dataclass
-class TFEfficientFormerForImageClassificationWithTeacherOutput(ModelOutput):
- """
- Args:
- Output type of [`EfficientFormerForImageClassificationWithTeacher`].
- logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores as the average of the cls_logits and distillation logits.
- cls_logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores of the classification head (i.e. the linear layer on top of the final hidden state of the
- class token).
- distillation_logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):
- Prediction scores of the distillation head (i.e. the linear layer on top of the final hidden state of the
- distillation token).
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when
- `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus
- the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when
- `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
- the self-attention heads.
- """
-
- logits: tf.Tensor = None
- cls_logits: tf.Tensor = None
- distillation_logits: tf.Tensor = None
- hidden_states: Optional[Tuple[tf.Tensor]] = None
- attentions: Optional[Tuple[tf.Tensor]] = None
-
-
-@add_start_docstrings(
- """
- EfficientFormer Model transformer with image classification heads on top (a linear layer on top of the final hidden
- state and a linear layer on top of the final hidden state of the distillation token) e.g. for ImageNet.
-
- .. warning::
- This model supports inference-only. Fine-tuning with distillation (i.e. with a teacher) is not yet
- supported.
- """,
- EFFICIENTFORMER_START_DOCSTRING,
-)
-class TFEfficientFormerForImageClassificationWithTeacher(TFEfficientFormerPreTrainedModel):
- def __init__(self, config: EfficientFormerConfig) -> None:
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.efficientformer = TFEfficientFormerMainLayer(config, name="efficientformer")
-
- # Classifier heads
- self.classifier = (
- keras.layers.Dense(config.num_labels, name="classifier")
- if config.num_labels > 0
- else keras.layers.Activation("linear", name="classifier")
- )
- self.distillation_classifier = (
- keras.layers.Dense(config.num_labels, name="distillation_classifier")
- if config.num_labels > 0
- else keras.layers.Activation("linear", name="distillation_classifier")
- )
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(EFFICIENTFORMER_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=TFEfficientFormerForImageClassificationWithTeacherOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def call(
- self,
- pixel_values: Optional[tf.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[tuple, TFEfficientFormerForImageClassificationWithTeacherOutput]:
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if training:
- raise Exception(
- "This model supports inference-only. Fine-tuning with distillation (i.e. with a teacher) is not yet supported."
- )
-
- outputs = self.efficientformer(
- pixel_values=pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = outputs[0]
-
- cls_logits = self.classifier(tf.reduce_mean(sequence_output, axis=-2))
- distillation_logits = self.distillation_classifier(tf.reduce_mean(sequence_output, axis=-2))
- logits = (cls_logits + distillation_logits) / 2
-
- if not return_dict:
- output = (logits, cls_logits, distillation_logits) + outputs[1:]
- return output
-
- return TFEfficientFormerForImageClassificationWithTeacherOutput(
- logits=logits,
- cls_logits=cls_logits,
- distillation_logits=distillation_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "efficientformer", None) is not None:
- with tf.name_scope(self.efficientformer.name):
- self.efficientformer.build(None)
- if getattr(self, "classifier", None) is not None:
- if hasattr(self.classifier, "name"):
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_sizes[-1]])
- if getattr(self, "distillation_classifier", None) is not None:
- if hasattr(self.distillation_classifier, "name"):
- with tf.name_scope(self.distillation_classifier.name):
- self.distillation_classifier.build([None, None, self.config.hidden_sizes[-1]])
diff --git a/transformers/models/efficientnet/__init__.py b/transformers/models/efficientnet/__init__.py
deleted file mode 100644
index 6df523721aefc55cf70bf627d935bd359acdeaab..0000000000000000000000000000000000000000
--- a/transformers/models/efficientnet/__init__.py
+++ /dev/null
@@ -1,84 +0,0 @@
-# flake8: noqa
-# There's no way to ignore "F401 '...' imported but unused" warnings in this
-# module, but to preserve other warnings. So, don't check this module at all.
-
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-# rely on isort to merge the imports
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {
- "configuration_efficientnet": [
- "EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "EfficientNetConfig",
- "EfficientNetOnnxConfig",
- ]
-}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["image_processing_efficientnet"] = ["EfficientNetImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_efficientnet"] = [
- "EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST",
- "EfficientNetForImageClassification",
- "EfficientNetModel",
- "EfficientNetPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_efficientnet import (
- EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
- EfficientNetConfig,
- EfficientNetOnnxConfig,
- )
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .image_processing_efficientnet import EfficientNetImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_efficientnet import (
- EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST,
- EfficientNetForImageClassification,
- EfficientNetModel,
- EfficientNetPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
diff --git a/transformers/models/efficientnet/__pycache__/__init__.cpython-310.pyc b/transformers/models/efficientnet/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 26f5750f81ab8c9fa2cbed6d5215748e3a1efc43..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientnet/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientnet/__pycache__/configuration_efficientnet.cpython-310.pyc b/transformers/models/efficientnet/__pycache__/configuration_efficientnet.cpython-310.pyc
deleted file mode 100644
index deb449c05b9c19f0b8e39c31a8928cbdd8df3b40..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientnet/__pycache__/configuration_efficientnet.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientnet/__pycache__/convert_efficientnet_to_pytorch.cpython-310.pyc b/transformers/models/efficientnet/__pycache__/convert_efficientnet_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 9ccd91845d00deb11e5e74ca718896c42f74bfaf..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientnet/__pycache__/convert_efficientnet_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientnet/__pycache__/image_processing_efficientnet.cpython-310.pyc b/transformers/models/efficientnet/__pycache__/image_processing_efficientnet.cpython-310.pyc
deleted file mode 100644
index c64be15c97edfa5226786b38f9e923e12eb3bf77..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientnet/__pycache__/image_processing_efficientnet.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientnet/__pycache__/modeling_efficientnet.cpython-310.pyc b/transformers/models/efficientnet/__pycache__/modeling_efficientnet.cpython-310.pyc
deleted file mode 100644
index 92b1c547ab545fea675297ed060ee616eab84ba6..0000000000000000000000000000000000000000
Binary files a/transformers/models/efficientnet/__pycache__/modeling_efficientnet.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/efficientnet/configuration_efficientnet.py b/transformers/models/efficientnet/configuration_efficientnet.py
deleted file mode 100644
index 77106c70d7d553fe73c6f4682f297f3037344007..0000000000000000000000000000000000000000
--- a/transformers/models/efficientnet/configuration_efficientnet.py
+++ /dev/null
@@ -1,169 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Google Research, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" EfficientNet model configuration"""
-
-from collections import OrderedDict
-from typing import List, Mapping
-
-from packaging import version
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class EfficientNetConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`EfficientNetModel`]. It is used to instantiate an
- EfficientNet model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the EfficientNet
- [google/efficientnet-b7](https://huggingface.co/google/efficientnet-b7) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- image_size (`int`, *optional*, defaults to 600):
- The input image size.
- width_coefficient (`float`, *optional*, defaults to 2.0):
- Scaling coefficient for network width at each stage.
- depth_coefficient (`float`, *optional*, defaults to 3.1):
- Scaling coefficient for network depth at each stage.
- depth_divisor `int`, *optional*, defaults to 8):
- A unit of network width.
- kernel_sizes (`List[int]`, *optional*, defaults to `[3, 3, 5, 3, 5, 5, 3]`):
- List of kernel sizes to be used in each block.
- in_channels (`List[int]`, *optional*, defaults to `[32, 16, 24, 40, 80, 112, 192]`):
- List of input channel sizes to be used in each block for convolutional layers.
- out_channels (`List[int]`, *optional*, defaults to `[16, 24, 40, 80, 112, 192, 320]`):
- List of output channel sizes to be used in each block for convolutional layers.
- depthwise_padding (`List[int]`, *optional*, defaults to `[]`):
- List of block indices with square padding.
- strides (`List[int]`, *optional*, defaults to `[1, 2, 2, 2, 1, 2, 1]`):
- List of stride sizes to be used in each block for convolutional layers.
- num_block_repeats (`List[int]`, *optional*, defaults to `[1, 2, 2, 3, 3, 4, 1]`):
- List of the number of times each block is to repeated.
- expand_ratios (`List[int]`, *optional*, defaults to `[1, 6, 6, 6, 6, 6, 6]`):
- List of scaling coefficient of each block.
- squeeze_expansion_ratio (`float`, *optional*, defaults to 0.25):
- Squeeze expansion ratio.
- hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
- The non-linear activation function (function or string) in each block. If string, `"gelu"`, `"relu"`,
- `"selu", `"gelu_new"`, `"silu"` and `"mish"` are supported.
- hiddem_dim (`int`, *optional*, defaults to 1280):
- The hidden dimension of the layer before the classification head.
- pooling_type (`str` or `function`, *optional*, defaults to `"mean"`):
- Type of final pooling to be applied before the dense classification head. Available options are [`"mean"`,
- `"max"`]
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- batch_norm_eps (`float`, *optional*, defaults to 1e-3):
- The epsilon used by the batch normalization layers.
- batch_norm_momentum (`float`, *optional*, defaults to 0.99):
- The momentum used by the batch normalization layers.
- dropout_rate (`float`, *optional*, defaults to 0.5):
- The dropout rate to be applied before final classifier layer.
- drop_connect_rate (`float`, *optional*, defaults to 0.2):
- The drop rate for skip connections.
-
- Example:
- ```python
- >>> from transformers import EfficientNetConfig, EfficientNetModel
-
- >>> # Initializing a EfficientNet efficientnet-b7 style configuration
- >>> configuration = EfficientNetConfig()
-
- >>> # Initializing a model (with random weights) from the efficientnet-b7 style configuration
- >>> model = EfficientNetModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "efficientnet"
-
- def __init__(
- self,
- num_channels: int = 3,
- image_size: int = 600,
- width_coefficient: float = 2.0,
- depth_coefficient: float = 3.1,
- depth_divisor: int = 8,
- kernel_sizes: List[int] = [3, 3, 5, 3, 5, 5, 3],
- in_channels: List[int] = [32, 16, 24, 40, 80, 112, 192],
- out_channels: List[int] = [16, 24, 40, 80, 112, 192, 320],
- depthwise_padding: List[int] = [],
- strides: List[int] = [1, 2, 2, 2, 1, 2, 1],
- num_block_repeats: List[int] = [1, 2, 2, 3, 3, 4, 1],
- expand_ratios: List[int] = [1, 6, 6, 6, 6, 6, 6],
- squeeze_expansion_ratio: float = 0.25,
- hidden_act: str = "swish",
- hidden_dim: int = 2560,
- pooling_type: str = "mean",
- initializer_range: float = 0.02,
- batch_norm_eps: float = 0.001,
- batch_norm_momentum: float = 0.99,
- dropout_rate: float = 0.5,
- drop_connect_rate: float = 0.2,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.num_channels = num_channels
- self.image_size = image_size
- self.width_coefficient = width_coefficient
- self.depth_coefficient = depth_coefficient
- self.depth_divisor = depth_divisor
- self.kernel_sizes = kernel_sizes
- self.in_channels = in_channels
- self.out_channels = out_channels
- self.depthwise_padding = depthwise_padding
- self.strides = strides
- self.num_block_repeats = num_block_repeats
- self.expand_ratios = expand_ratios
- self.squeeze_expansion_ratio = squeeze_expansion_ratio
- self.hidden_act = hidden_act
- self.hidden_dim = hidden_dim
- self.pooling_type = pooling_type
- self.initializer_range = initializer_range
- self.batch_norm_eps = batch_norm_eps
- self.batch_norm_momentum = batch_norm_momentum
- self.dropout_rate = dropout_rate
- self.drop_connect_rate = drop_connect_rate
- self.num_hidden_layers = sum(num_block_repeats) * 4
-
-
-class EfficientNetOnnxConfig(OnnxConfig):
- torch_onnx_minimum_version = version.parse("1.11")
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- return OrderedDict(
- [
- ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
- ]
- )
-
- @property
- def atol_for_validation(self) -> float:
- return 1e-5
diff --git a/transformers/models/efficientnet/convert_efficientnet_to_pytorch.py b/transformers/models/efficientnet/convert_efficientnet_to_pytorch.py
deleted file mode 100644
index e9988524aca04de2a1d600586ff01d9b9a3ea6c2..0000000000000000000000000000000000000000
--- a/transformers/models/efficientnet/convert_efficientnet_to_pytorch.py
+++ /dev/null
@@ -1,339 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert EfficientNet checkpoints from the original repository.
-
-URL: https://github.com/keras-team/keras/blob/v2.11.0/keras/applications/efficientnet.py"""
-
-import argparse
-import json
-import os
-
-import numpy as np
-import PIL
-import requests
-import tensorflow.keras.applications.efficientnet as efficientnet
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-from tensorflow.keras.preprocessing import image
-
-from transformers import (
- EfficientNetConfig,
- EfficientNetForImageClassification,
- EfficientNetImageProcessor,
-)
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-model_classes = {
- "b0": efficientnet.EfficientNetB0,
- "b1": efficientnet.EfficientNetB1,
- "b2": efficientnet.EfficientNetB2,
- "b3": efficientnet.EfficientNetB3,
- "b4": efficientnet.EfficientNetB4,
- "b5": efficientnet.EfficientNetB5,
- "b6": efficientnet.EfficientNetB6,
- "b7": efficientnet.EfficientNetB7,
-}
-
-CONFIG_MAP = {
- "b0": {
- "hidden_dim": 1280,
- "width_coef": 1.0,
- "depth_coef": 1.0,
- "image_size": 224,
- "dropout_rate": 0.2,
- "dw_padding": [],
- },
- "b1": {
- "hidden_dim": 1280,
- "width_coef": 1.0,
- "depth_coef": 1.1,
- "image_size": 240,
- "dropout_rate": 0.2,
- "dw_padding": [16],
- },
- "b2": {
- "hidden_dim": 1408,
- "width_coef": 1.1,
- "depth_coef": 1.2,
- "image_size": 260,
- "dropout_rate": 0.3,
- "dw_padding": [5, 8, 16],
- },
- "b3": {
- "hidden_dim": 1536,
- "width_coef": 1.2,
- "depth_coef": 1.4,
- "image_size": 300,
- "dropout_rate": 0.3,
- "dw_padding": [5, 18],
- },
- "b4": {
- "hidden_dim": 1792,
- "width_coef": 1.4,
- "depth_coef": 1.8,
- "image_size": 380,
- "dropout_rate": 0.4,
- "dw_padding": [6],
- },
- "b5": {
- "hidden_dim": 2048,
- "width_coef": 1.6,
- "depth_coef": 2.2,
- "image_size": 456,
- "dropout_rate": 0.4,
- "dw_padding": [13, 27],
- },
- "b6": {
- "hidden_dim": 2304,
- "width_coef": 1.8,
- "depth_coef": 2.6,
- "image_size": 528,
- "dropout_rate": 0.5,
- "dw_padding": [31],
- },
- "b7": {
- "hidden_dim": 2560,
- "width_coef": 2.0,
- "depth_coef": 3.1,
- "image_size": 600,
- "dropout_rate": 0.5,
- "dw_padding": [18],
- },
-}
-
-
-def get_efficientnet_config(model_name):
- config = EfficientNetConfig()
- config.hidden_dim = CONFIG_MAP[model_name]["hidden_dim"]
- config.width_coefficient = CONFIG_MAP[model_name]["width_coef"]
- config.depth_coefficient = CONFIG_MAP[model_name]["depth_coef"]
- config.image_size = CONFIG_MAP[model_name]["image_size"]
- config.dropout_rate = CONFIG_MAP[model_name]["dropout_rate"]
- config.depthwise_padding = CONFIG_MAP[model_name]["dw_padding"]
-
- repo_id = "huggingface/label-files"
- filename = "imagenet-1k-id2label.json"
- config.num_labels = 1000
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
-
- config.id2label = id2label
- config.label2id = {v: k for k, v in id2label.items()}
- return config
-
-
-# We will verify our results on an image of cute cats
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- im = Image.open(requests.get(url, stream=True).raw)
- return im
-
-
-def convert_image_processor(model_name):
- size = CONFIG_MAP[model_name]["image_size"]
- preprocessor = EfficientNetImageProcessor(
- size={"height": size, "width": size},
- image_mean=[0.485, 0.456, 0.406],
- image_std=[0.47853944, 0.4732864, 0.47434163],
- do_center_crop=False,
- )
- return preprocessor
-
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-def rename_keys(original_param_names):
- block_names = [v.split("_")[0].split("block")[1] for v in original_param_names if v.startswith("block")]
- block_names = sorted(set(block_names))
- num_blocks = len(block_names)
- block_name_mapping = {b: str(i) for b, i in zip(block_names, range(num_blocks))}
-
- rename_keys = []
- rename_keys.append(("stem_conv/kernel:0", "embeddings.convolution.weight"))
- rename_keys.append(("stem_bn/gamma:0", "embeddings.batchnorm.weight"))
- rename_keys.append(("stem_bn/beta:0", "embeddings.batchnorm.bias"))
- rename_keys.append(("stem_bn/moving_mean:0", "embeddings.batchnorm.running_mean"))
- rename_keys.append(("stem_bn/moving_variance:0", "embeddings.batchnorm.running_var"))
-
- for b in block_names:
- hf_b = block_name_mapping[b]
- rename_keys.append((f"block{b}_expand_conv/kernel:0", f"encoder.blocks.{hf_b}.expansion.expand_conv.weight"))
- rename_keys.append((f"block{b}_expand_bn/gamma:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.weight"))
- rename_keys.append((f"block{b}_expand_bn/beta:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.bias"))
- rename_keys.append(
- (f"block{b}_expand_bn/moving_mean:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.running_mean")
- )
- rename_keys.append(
- (f"block{b}_expand_bn/moving_variance:0", f"encoder.blocks.{hf_b}.expansion.expand_bn.running_var")
- )
- rename_keys.append(
- (f"block{b}_dwconv/depthwise_kernel:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight")
- )
- rename_keys.append((f"block{b}_bn/gamma:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight"))
- rename_keys.append((f"block{b}_bn/beta:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias"))
- rename_keys.append(
- (f"block{b}_bn/moving_mean:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean")
- )
- rename_keys.append(
- (f"block{b}_bn/moving_variance:0", f"encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var")
- )
-
- rename_keys.append((f"block{b}_se_reduce/kernel:0", f"encoder.blocks.{hf_b}.squeeze_excite.reduce.weight"))
- rename_keys.append((f"block{b}_se_reduce/bias:0", f"encoder.blocks.{hf_b}.squeeze_excite.reduce.bias"))
- rename_keys.append((f"block{b}_se_expand/kernel:0", f"encoder.blocks.{hf_b}.squeeze_excite.expand.weight"))
- rename_keys.append((f"block{b}_se_expand/bias:0", f"encoder.blocks.{hf_b}.squeeze_excite.expand.bias"))
- rename_keys.append(
- (f"block{b}_project_conv/kernel:0", f"encoder.blocks.{hf_b}.projection.project_conv.weight")
- )
- rename_keys.append((f"block{b}_project_bn/gamma:0", f"encoder.blocks.{hf_b}.projection.project_bn.weight"))
- rename_keys.append((f"block{b}_project_bn/beta:0", f"encoder.blocks.{hf_b}.projection.project_bn.bias"))
- rename_keys.append(
- (f"block{b}_project_bn/moving_mean:0", f"encoder.blocks.{hf_b}.projection.project_bn.running_mean")
- )
- rename_keys.append(
- (f"block{b}_project_bn/moving_variance:0", f"encoder.blocks.{hf_b}.projection.project_bn.running_var")
- )
-
- rename_keys.append(("top_conv/kernel:0", "encoder.top_conv.weight"))
- rename_keys.append(("top_bn/gamma:0", "encoder.top_bn.weight"))
- rename_keys.append(("top_bn/beta:0", "encoder.top_bn.bias"))
- rename_keys.append(("top_bn/moving_mean:0", "encoder.top_bn.running_mean"))
- rename_keys.append(("top_bn/moving_variance:0", "encoder.top_bn.running_var"))
-
- key_mapping = {}
- for item in rename_keys:
- if item[0] in original_param_names:
- key_mapping[item[0]] = "efficientnet." + item[1]
-
- key_mapping["predictions/kernel:0"] = "classifier.weight"
- key_mapping["predictions/bias:0"] = "classifier.bias"
- return key_mapping
-
-
-def replace_params(hf_params, tf_params, key_mapping):
- for key, value in tf_params.items():
- if "normalization" in key:
- continue
-
- hf_key = key_mapping[key]
- if "_conv" in key and "kernel" in key:
- new_hf_value = torch.from_numpy(value).permute(3, 2, 0, 1)
- elif "depthwise_kernel" in key:
- new_hf_value = torch.from_numpy(value).permute(2, 3, 0, 1)
- elif "kernel" in key:
- new_hf_value = torch.from_numpy(np.transpose(value))
- else:
- new_hf_value = torch.from_numpy(value)
-
- # Replace HF parameters with original TF model parameters
- assert hf_params[hf_key].shape == new_hf_value.shape
- hf_params[hf_key].copy_(new_hf_value)
-
-
-@torch.no_grad()
-def convert_efficientnet_checkpoint(model_name, pytorch_dump_folder_path, save_model, push_to_hub):
- """
- Copy/paste/tweak model's weights to our EfficientNet structure.
- """
- # Load original model
- original_model = model_classes[model_name](
- include_top=True,
- weights="imagenet",
- input_tensor=None,
- input_shape=None,
- pooling=None,
- classes=1000,
- classifier_activation="softmax",
- )
-
- tf_params = original_model.trainable_variables
- tf_non_train_params = original_model.non_trainable_variables
- tf_params = {param.name: param.numpy() for param in tf_params}
- for param in tf_non_train_params:
- tf_params[param.name] = param.numpy()
- tf_param_names = list(tf_params.keys())
-
- # Load HuggingFace model
- config = get_efficientnet_config(model_name)
- hf_model = EfficientNetForImageClassification(config).eval()
- hf_params = hf_model.state_dict()
-
- # Create src-to-dst parameter name mapping dictionary
- print("Converting parameters...")
- key_mapping = rename_keys(tf_param_names)
- replace_params(hf_params, tf_params, key_mapping)
-
- # Initialize preprocessor and preprocess input image
- preprocessor = convert_image_processor(model_name)
- inputs = preprocessor(images=prepare_img(), return_tensors="pt")
-
- # HF model inference
- hf_model.eval()
- with torch.no_grad():
- outputs = hf_model(**inputs)
- hf_logits = outputs.logits.detach().numpy()
-
- # Original model inference
- original_model.trainable = False
- image_size = CONFIG_MAP[model_name]["image_size"]
- img = prepare_img().resize((image_size, image_size), resample=PIL.Image.NEAREST)
- x = image.img_to_array(img)
- x = np.expand_dims(x, axis=0)
- original_logits = original_model.predict(x)
-
- # Check whether original and HF model outputs match -> np.allclose
- assert np.allclose(original_logits, hf_logits, atol=1e-3), "The predicted logits are not the same."
- print("Model outputs match!")
-
- if save_model:
- # Create folder to save model
- if not os.path.isdir(pytorch_dump_folder_path):
- os.mkdir(pytorch_dump_folder_path)
- # Save converted model and image processor
- hf_model.save_pretrained(pytorch_dump_folder_path)
- preprocessor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- # Push model and image processor to hub
- print(f"Pushing converted {model_name} to the hub...")
- model_name = f"efficientnet-{model_name}"
- preprocessor.push_to_hub(model_name)
- hf_model.push_to_hub(model_name)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="b0",
- type=str,
- help="Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default="hf_model",
- type=str,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument("--save_model", action="store_true", help="Save model to local")
- parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")
-
- args = parser.parse_args()
- convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)
diff --git a/transformers/models/efficientnet/image_processing_efficientnet.py b/transformers/models/efficientnet/image_processing_efficientnet.py
deleted file mode 100644
index 4fd2364a3020c5ee48b387242df8a0dc24122a24..0000000000000000000000000000000000000000
--- a/transformers/models/efficientnet/image_processing_efficientnet.py
+++ /dev/null
@@ -1,387 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for EfficientNet."""
-
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
-from ...image_transforms import rescale, resize, to_channel_dimension_format
-from ...image_utils import (
- IMAGENET_STANDARD_MEAN,
- IMAGENET_STANDARD_STD,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import TensorType, is_vision_available, logging
-
-
-if is_vision_available():
- import PIL
-
-
-logger = logging.get_logger(__name__)
-
-
-class EfficientNetImageProcessor(BaseImageProcessor):
- r"""
- Constructs a EfficientNet image processor.
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by
- `do_resize` in `preprocess`.
- size (`Dict[str, int]` *optional*, defaults to `{"height": 346, "width": 346}`):
- Size of the image after `resize`. Can be overridden by `size` in `preprocess`.
- resample (`PILImageResampling` filter, *optional*, defaults to 0):
- Resampling filter to use if resizing the image. Can be overridden by `resample` in `preprocess`.
- do_center_crop (`bool`, *optional*, defaults to `False`):
- Whether to center crop the image. If the input size is smaller than `crop_size` along any edge, the image
- is padded with 0's and then center cropped. Can be overridden by `do_center_crop` in `preprocess`.
- crop_size (`Dict[str, int]`, *optional*, defaults to `{"height": 289, "width": 289}`):
- Desired output size when applying center-cropping. Can be overridden by `crop_size` in `preprocess`.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the
- `preprocess` method.
- rescale_offset (`bool`, *optional*, defaults to `False`):
- Whether to rescale the image between [-scale_range, scale_range] instead of [0, scale_range]. Can be
- overridden by the `rescale_factor` parameter in the `preprocess` method.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`
- parameter in the `preprocess` method.
- do_normalize (`bool`, *optional*, defaults to `True`):
- Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
- method.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
- Mean to use if normalizing the image. This is a float or list of floats the length of the number of
- channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
- Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
- number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
- include_top (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image again. Should be set to True if the inputs are used for image classification.
- """
-
- model_input_names = ["pixel_values"]
-
- def __init__(
- self,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PIL.Image.NEAREST,
- do_center_crop: bool = False,
- crop_size: Dict[str, int] = None,
- rescale_factor: Union[int, float] = 1 / 255,
- rescale_offset: bool = False,
- do_rescale: bool = True,
- do_normalize: bool = True,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- include_top: bool = True,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- size = size if size is not None else {"height": 346, "width": 346}
- size = get_size_dict(size)
- crop_size = crop_size if crop_size is not None else {"height": 289, "width": 289}
- crop_size = get_size_dict(crop_size, param_name="crop_size")
-
- self.do_resize = do_resize
- self.size = size
- self.resample = resample
- self.do_center_crop = do_center_crop
- self.crop_size = crop_size
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.rescale_offset = rescale_offset
- self.do_normalize = do_normalize
- self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
- self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
- self.include_top = include_top
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size",
- "resample",
- "do_center_crop",
- "crop_size",
- "do_rescale",
- "rescale_factor",
- "rescale_offset",
- "do_normalize",
- "image_mean",
- "image_std",
- "include_top",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- # Copied from transformers.models.vit.image_processing_vit.ViTImageProcessor.resize with PILImageResampling.BILINEAR->PILImageResampling.NEAREST
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.NEAREST,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize an image to `(size["height"], size["width"])`.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.NEAREST`):
- `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.NEAREST`.
- data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
-
- Returns:
- `np.ndarray`: The resized image.
- """
- size = get_size_dict(size)
- if "height" not in size or "width" not in size:
- raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}")
- output_size = (size["height"], size["width"])
- return resize(
- image,
- size=output_size,
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
-
- def rescale(
- self,
- image: np.ndarray,
- scale: Union[int, float],
- offset: bool = True,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ):
- """
- Rescale an image by a scale factor.
-
- If `offset` is `True`, the image has its values rescaled by `scale` and then offset by 1. If `scale` is
- 1/127.5, the image is rescaled between [-1, 1].
- image = image * scale - 1
-
- If `offset` is `False`, and `scale` is 1/255, the image is rescaled between [0, 1].
- image = image * scale
-
- Args:
- image (`np.ndarray`):
- Image to rescale.
- scale (`int` or `float`):
- Scale to apply to the image.
- offset (`bool`, *optional*):
- Whether to scale the image in both negative and positive directions.
- data_format (`str` or `ChannelDimension`, *optional*):
- The channel dimension format of the image. If not provided, it will be the same as the input image.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- rescaled_image = rescale(
- image, scale=scale, data_format=data_format, input_data_format=input_data_format, **kwargs
- )
-
- if offset:
- rescaled_image = rescaled_image - 1
-
- return rescaled_image
-
- def preprocess(
- self,
- images: ImageInput,
- do_resize: bool = None,
- size: Dict[str, int] = None,
- resample=None,
- do_center_crop: bool = None,
- crop_size: Dict[str, int] = None,
- do_rescale: bool = None,
- rescale_factor: float = None,
- rescale_offset: bool = None,
- do_normalize: bool = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- include_top: bool = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: ChannelDimension = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> PIL.Image.Image:
- """
- Preprocess an image or batch of images.
-
- Args:
- images (`ImageInput`):
- Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
- passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
- Size of the image after `resize`.
- resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
- PILImageResampling filter to use if resizing the image Only has an effect if `do_resize` is set to
- `True`.
- do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`):
- Whether to center crop the image.
- crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
- Size of the image after center crop. If one edge the image is smaller than `crop_size`, it will be
- padded with zeros and then cropped
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether to rescale the image values between [0 - 1].
- rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
- Rescale factor to rescale the image by if `do_rescale` is set to `True`.
- rescale_offset (`bool`, *optional*, defaults to `self.rescale_offset`):
- Whether to rescale the image between [-scale_range, scale_range] instead of [0, scale_range].
- do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
- Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
- Image mean.
- image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
- Image standard deviation.
- include_top (`bool`, *optional*, defaults to `self.include_top`):
- Rescales the image again for image classification if set to True.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - `None`: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- do_resize = do_resize if do_resize is not None else self.do_resize
- resample = resample if resample is not None else self.resample
- do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- rescale_offset = rescale_offset if rescale_offset is not None else self.rescale_offset
- do_normalize = do_normalize if do_normalize is not None else self.do_normalize
- image_mean = image_mean if image_mean is not None else self.image_mean
- image_std = image_std if image_std is not None else self.image_std
- include_top = include_top if include_top is not None else self.include_top
-
- size = size if size is not None else self.size
- size = get_size_dict(size)
- crop_size = crop_size if crop_size is not None else self.crop_size
- crop_size = get_size_dict(crop_size, param_name="crop_size")
-
- images = make_list_of_images(images)
-
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_center_crop=do_center_crop,
- crop_size=crop_size,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
- # All transformations expect numpy arrays.
- images = [to_numpy_array(image) for image in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- if do_resize:
- images = [
- self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_center_crop:
- images = [
- self.center_crop(image=image, size=crop_size, input_data_format=input_data_format) for image in images
- ]
-
- if do_rescale:
- images = [
- self.rescale(
- image=image, scale=rescale_factor, offset=rescale_offset, input_data_format=input_data_format
- )
- for image in images
- ]
-
- if do_normalize:
- images = [
- self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
- for image in images
- ]
-
- if include_top:
- images = [
- self.normalize(image=image, mean=0, std=image_std, input_data_format=input_data_format)
- for image in images
- ]
-
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
- ]
-
- data = {"pixel_values": images}
- return BatchFeature(data=data, tensor_type=return_tensors)
diff --git a/transformers/models/efficientnet/modeling_efficientnet.py b/transformers/models/efficientnet/modeling_efficientnet.py
deleted file mode 100644
index e415d7f1b46a1ed61fca75955423a88e48a86e94..0000000000000000000000000000000000000000
--- a/transformers/models/efficientnet/modeling_efficientnet.py
+++ /dev/null
@@ -1,648 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Google Research, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch EfficientNet model."""
-
-
-import math
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutputWithNoAttention,
- BaseModelOutputWithPoolingAndNoAttention,
- ImageClassifierOutputWithNoAttention,
-)
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_efficientnet import EfficientNetConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "EfficientNetConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "google/efficientnet-b7"
-_EXPECTED_OUTPUT_SHAPE = [1, 768, 7, 7]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "google/efficientnet-b7"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-EFFICIENTNET_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`EfficientNetConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-EFFICIENTNET_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`AutoImageProcessor.__call__`] for details.
-
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-def round_filters(config: EfficientNetConfig, num_channels: int):
- r"""
- Round number of filters based on depth multiplier.
- """
- divisor = config.depth_divisor
- num_channels *= config.width_coefficient
- new_dim = max(divisor, int(num_channels + divisor / 2) // divisor * divisor)
-
- # Make sure that round down does not go down by more than 10%.
- if new_dim < 0.9 * num_channels:
- new_dim += divisor
-
- return int(new_dim)
-
-
-def correct_pad(kernel_size: Union[int, Tuple], adjust: bool = True):
- r"""
- Utility function to get the tuple padding value for the depthwise convolution.
-
- Args:
- kernel_size (`int` or `tuple`):
- Kernel size of the convolution layers.
- adjust (`bool`, *optional*, defaults to `True`):
- Adjusts padding value to apply to right and bottom sides of the input.
- """
- if isinstance(kernel_size, int):
- kernel_size = (kernel_size, kernel_size)
-
- correct = (kernel_size[0] // 2, kernel_size[1] // 2)
- if adjust:
- return (correct[1] - 1, correct[1], correct[0] - 1, correct[0])
- else:
- return (correct[1], correct[1], correct[0], correct[0])
-
-
-class EfficientNetEmbeddings(nn.Module):
- r"""
- A module that corresponds to the stem module of the original work.
- """
-
- def __init__(self, config: EfficientNetConfig):
- super().__init__()
-
- self.out_dim = round_filters(config, 32)
- self.padding = nn.ZeroPad2d(padding=(0, 1, 0, 1))
- self.convolution = nn.Conv2d(
- config.num_channels, self.out_dim, kernel_size=3, stride=2, padding="valid", bias=False
- )
- self.batchnorm = nn.BatchNorm2d(self.out_dim, eps=config.batch_norm_eps, momentum=config.batch_norm_momentum)
- self.activation = ACT2FN[config.hidden_act]
-
- def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
- features = self.padding(pixel_values)
- features = self.convolution(features)
- features = self.batchnorm(features)
- features = self.activation(features)
-
- return features
-
-
-class EfficientNetDepthwiseConv2d(nn.Conv2d):
- def __init__(
- self,
- in_channels,
- depth_multiplier=1,
- kernel_size=3,
- stride=1,
- padding=0,
- dilation=1,
- bias=True,
- padding_mode="zeros",
- ):
- out_channels = in_channels * depth_multiplier
- super().__init__(
- in_channels=in_channels,
- out_channels=out_channels,
- kernel_size=kernel_size,
- stride=stride,
- padding=padding,
- dilation=dilation,
- groups=in_channels,
- bias=bias,
- padding_mode=padding_mode,
- )
-
-
-class EfficientNetExpansionLayer(nn.Module):
- r"""
- This corresponds to the expansion phase of each block in the original implementation.
- """
-
- def __init__(self, config: EfficientNetConfig, in_dim: int, out_dim: int, stride: int):
- super().__init__()
- self.expand_conv = nn.Conv2d(
- in_channels=in_dim,
- out_channels=out_dim,
- kernel_size=1,
- padding="same",
- bias=False,
- )
- self.expand_bn = nn.BatchNorm2d(num_features=out_dim, eps=config.batch_norm_eps)
- self.expand_act = ACT2FN[config.hidden_act]
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
- # Expand phase
- hidden_states = self.expand_conv(hidden_states)
- hidden_states = self.expand_bn(hidden_states)
- hidden_states = self.expand_act(hidden_states)
-
- return hidden_states
-
-
-class EfficientNetDepthwiseLayer(nn.Module):
- r"""
- This corresponds to the depthwise convolution phase of each block in the original implementation.
- """
-
- def __init__(
- self,
- config: EfficientNetConfig,
- in_dim: int,
- stride: int,
- kernel_size: int,
- adjust_padding: bool,
- ):
- super().__init__()
- self.stride = stride
- conv_pad = "valid" if self.stride == 2 else "same"
- padding = correct_pad(kernel_size, adjust=adjust_padding)
-
- self.depthwise_conv_pad = nn.ZeroPad2d(padding=padding)
- self.depthwise_conv = EfficientNetDepthwiseConv2d(
- in_dim, kernel_size=kernel_size, stride=stride, padding=conv_pad, bias=False
- )
- self.depthwise_norm = nn.BatchNorm2d(
- num_features=in_dim, eps=config.batch_norm_eps, momentum=config.batch_norm_momentum
- )
- self.depthwise_act = ACT2FN[config.hidden_act]
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
- # Depthwise convolution
- if self.stride == 2:
- hidden_states = self.depthwise_conv_pad(hidden_states)
-
- hidden_states = self.depthwise_conv(hidden_states)
- hidden_states = self.depthwise_norm(hidden_states)
- hidden_states = self.depthwise_act(hidden_states)
-
- return hidden_states
-
-
-class EfficientNetSqueezeExciteLayer(nn.Module):
- r"""
- This corresponds to the Squeeze and Excitement phase of each block in the original implementation.
- """
-
- def __init__(self, config: EfficientNetConfig, in_dim: int, expand_dim: int, expand: bool = False):
- super().__init__()
- self.dim = expand_dim if expand else in_dim
- self.dim_se = max(1, int(in_dim * config.squeeze_expansion_ratio))
-
- self.squeeze = nn.AdaptiveAvgPool2d(output_size=1)
- self.reduce = nn.Conv2d(
- in_channels=self.dim,
- out_channels=self.dim_se,
- kernel_size=1,
- padding="same",
- )
- self.expand = nn.Conv2d(
- in_channels=self.dim_se,
- out_channels=self.dim,
- kernel_size=1,
- padding="same",
- )
- self.act_reduce = ACT2FN[config.hidden_act]
- self.act_expand = nn.Sigmoid()
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
- inputs = hidden_states
- hidden_states = self.squeeze(hidden_states)
- hidden_states = self.reduce(hidden_states)
- hidden_states = self.act_reduce(hidden_states)
-
- hidden_states = self.expand(hidden_states)
- hidden_states = self.act_expand(hidden_states)
- hidden_states = torch.mul(inputs, hidden_states)
-
- return hidden_states
-
-
-class EfficientNetFinalBlockLayer(nn.Module):
- r"""
- This corresponds to the final phase of each block in the original implementation.
- """
-
- def __init__(
- self, config: EfficientNetConfig, in_dim: int, out_dim: int, stride: int, drop_rate: float, id_skip: bool
- ):
- super().__init__()
- self.apply_dropout = stride == 1 and not id_skip
- self.project_conv = nn.Conv2d(
- in_channels=in_dim,
- out_channels=out_dim,
- kernel_size=1,
- padding="same",
- bias=False,
- )
- self.project_bn = nn.BatchNorm2d(
- num_features=out_dim, eps=config.batch_norm_eps, momentum=config.batch_norm_momentum
- )
- self.dropout = nn.Dropout(p=drop_rate)
-
- def forward(self, embeddings: torch.FloatTensor, hidden_states: torch.FloatTensor) -> torch.Tensor:
- hidden_states = self.project_conv(hidden_states)
- hidden_states = self.project_bn(hidden_states)
-
- if self.apply_dropout:
- hidden_states = self.dropout(hidden_states)
- hidden_states = hidden_states + embeddings
-
- return hidden_states
-
-
-class EfficientNetBlock(nn.Module):
- r"""
- This corresponds to the expansion and depthwise convolution phase of each block in the original implementation.
-
- Args:
- config ([`EfficientNetConfig`]):
- Model configuration class.
- in_dim (`int`):
- Number of input channels.
- out_dim (`int`):
- Number of output channels.
- stride (`int`):
- Stride size to be used in convolution layers.
- expand_ratio (`int`):
- Expand ratio to set the output dimensions for the expansion and squeeze-excite layers.
- kernel_size (`int`):
- Kernel size for the depthwise convolution layer.
- drop_rate (`float`):
- Dropout rate to be used in the final phase of each block.
- id_skip (`bool`):
- Whether to apply dropout and sum the final hidden states with the input embeddings during the final phase
- of each block. Set to `True` for the first block of each stage.
- adjust_padding (`bool`):
- Whether to apply padding to only right and bottom side of the input kernel before the depthwise convolution
- operation, set to `True` for inputs with odd input sizes.
- """
-
- def __init__(
- self,
- config: EfficientNetConfig,
- in_dim: int,
- out_dim: int,
- stride: int,
- expand_ratio: int,
- kernel_size: int,
- drop_rate: float,
- id_skip: bool,
- adjust_padding: bool,
- ):
- super().__init__()
- self.expand_ratio = expand_ratio
- self.expand = True if self.expand_ratio != 1 else False
- expand_in_dim = in_dim * expand_ratio
-
- if self.expand:
- self.expansion = EfficientNetExpansionLayer(
- config=config, in_dim=in_dim, out_dim=expand_in_dim, stride=stride
- )
-
- self.depthwise_conv = EfficientNetDepthwiseLayer(
- config=config,
- in_dim=expand_in_dim if self.expand else in_dim,
- stride=stride,
- kernel_size=kernel_size,
- adjust_padding=adjust_padding,
- )
- self.squeeze_excite = EfficientNetSqueezeExciteLayer(
- config=config, in_dim=in_dim, expand_dim=expand_in_dim, expand=self.expand
- )
- self.projection = EfficientNetFinalBlockLayer(
- config=config,
- in_dim=expand_in_dim if self.expand else in_dim,
- out_dim=out_dim,
- stride=stride,
- drop_rate=drop_rate,
- id_skip=id_skip,
- )
-
- def forward(self, hidden_states: torch.FloatTensor) -> torch.Tensor:
- embeddings = hidden_states
- # Expansion and depthwise convolution phase
- if self.expand_ratio != 1:
- hidden_states = self.expansion(hidden_states)
- hidden_states = self.depthwise_conv(hidden_states)
-
- # Squeeze and excite phase
- hidden_states = self.squeeze_excite(hidden_states)
- hidden_states = self.projection(embeddings, hidden_states)
- return hidden_states
-
-
-class EfficientNetEncoder(nn.Module):
- r"""
- Forward propogates the embeddings through each EfficientNet block.
-
- Args:
- config ([`EfficientNetConfig`]):
- Model configuration class.
- """
-
- def __init__(self, config: EfficientNetConfig):
- super().__init__()
- self.config = config
- self.depth_coefficient = config.depth_coefficient
-
- def round_repeats(repeats):
- # Round number of block repeats based on depth multiplier.
- return int(math.ceil(self.depth_coefficient * repeats))
-
- num_base_blocks = len(config.in_channels)
- num_blocks = sum(round_repeats(n) for n in config.num_block_repeats)
-
- curr_block_num = 0
- blocks = []
- for i in range(num_base_blocks):
- in_dim = round_filters(config, config.in_channels[i])
- out_dim = round_filters(config, config.out_channels[i])
- stride = config.strides[i]
- kernel_size = config.kernel_sizes[i]
- expand_ratio = config.expand_ratios[i]
-
- for j in range(round_repeats(config.num_block_repeats[i])):
- id_skip = True if j == 0 else False
- stride = 1 if j > 0 else stride
- in_dim = out_dim if j > 0 else in_dim
- adjust_padding = False if curr_block_num in config.depthwise_padding else True
- drop_rate = config.drop_connect_rate * curr_block_num / num_blocks
-
- block = EfficientNetBlock(
- config=config,
- in_dim=in_dim,
- out_dim=out_dim,
- stride=stride,
- kernel_size=kernel_size,
- expand_ratio=expand_ratio,
- drop_rate=drop_rate,
- id_skip=id_skip,
- adjust_padding=adjust_padding,
- )
- blocks.append(block)
- curr_block_num += 1
-
- self.blocks = nn.ModuleList(blocks)
- self.top_conv = nn.Conv2d(
- in_channels=out_dim,
- out_channels=round_filters(config, 1280),
- kernel_size=1,
- padding="same",
- bias=False,
- )
- self.top_bn = nn.BatchNorm2d(
- num_features=config.hidden_dim, eps=config.batch_norm_eps, momentum=config.batch_norm_momentum
- )
- self.top_activation = ACT2FN[config.hidden_act]
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> BaseModelOutputWithNoAttention:
- all_hidden_states = (hidden_states,) if output_hidden_states else None
-
- for block in self.blocks:
- hidden_states = block(hidden_states)
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- hidden_states = self.top_conv(hidden_states)
- hidden_states = self.top_bn(hidden_states)
- hidden_states = self.top_activation(hidden_states)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
-
- return BaseModelOutputWithNoAttention(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- )
-
-
-class EfficientNetPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = EfficientNetConfig
- base_model_prefix = "efficientnet"
- main_input_name = "pixel_values"
- _no_split_modules = []
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-@add_start_docstrings(
- "The bare EfficientNet model outputting raw features without any specific head on top.",
- EFFICIENTNET_START_DOCSTRING,
-)
-class EfficientNetModel(EfficientNetPreTrainedModel):
- def __init__(self, config: EfficientNetConfig):
- super().__init__(config)
- self.config = config
- self.embeddings = EfficientNetEmbeddings(config)
- self.encoder = EfficientNetEncoder(config)
-
- # Final pooling layer
- if config.pooling_type == "mean":
- self.pooler = nn.AvgPool2d(config.hidden_dim, ceil_mode=True)
- elif config.pooling_type == "max":
- self.pooler = nn.MaxPool2d(config.hidden_dim, ceil_mode=True)
- else:
- raise ValueError(f"config.pooling must be one of ['mean', 'max'] got {config.pooling}")
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(EFFICIENTNET_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPoolingAndNoAttention,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: torch.FloatTensor = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPoolingAndNoAttention]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- embedding_output = self.embeddings(pixel_values)
-
- encoder_outputs = self.encoder(
- embedding_output,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- # Apply pooling
- last_hidden_state = encoder_outputs[0]
- pooled_output = self.pooler(last_hidden_state)
- # Reshape (batch_size, 1280, 1 , 1) -> (batch_size, 1280)
- pooled_output = pooled_output.reshape(pooled_output.shape[:2])
-
- if not return_dict:
- return (last_hidden_state, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPoolingAndNoAttention(
- last_hidden_state=last_hidden_state,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- EfficientNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g.
- for ImageNet.
- """,
- EFFICIENTNET_START_DOCSTRING,
-)
-class EfficientNetForImageClassification(EfficientNetPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
- self.efficientnet = EfficientNetModel(config)
- # Classifier head
- self.dropout = nn.Dropout(p=config.dropout_rate)
- self.classifier = nn.Linear(config.hidden_dim, self.num_labels) if self.num_labels > 0 else nn.Identity()
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(EFFICIENTNET_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=ImageClassifierOutputWithNoAttention,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: torch.FloatTensor = None,
- labels: Optional[torch.LongTensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.efficientnet(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
-
- pooled_output = outputs.pooler_output if return_dict else outputs[1]
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return ImageClassifierOutputWithNoAttention(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- )
diff --git a/transformers/models/electra/__init__.py b/transformers/models/electra/__init__.py
deleted file mode 100644
index 09ce039d25fd057608693a8d6c9d79358d970225..0000000000000000000000000000000000000000
--- a/transformers/models/electra/__init__.py
+++ /dev/null
@@ -1,168 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_flax_available,
- is_tf_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_electra": ["ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP", "ElectraConfig", "ElectraOnnxConfig"],
- "tokenization_electra": ["ElectraTokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_electra_fast"] = ["ElectraTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_electra"] = [
- "ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST",
- "ElectraForCausalLM",
- "ElectraForMaskedLM",
- "ElectraForMultipleChoice",
- "ElectraForPreTraining",
- "ElectraForQuestionAnswering",
- "ElectraForSequenceClassification",
- "ElectraForTokenClassification",
- "ElectraModel",
- "ElectraPreTrainedModel",
- "load_tf_weights_in_electra",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_electra"] = [
- "TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFElectraForMaskedLM",
- "TFElectraForMultipleChoice",
- "TFElectraForPreTraining",
- "TFElectraForQuestionAnswering",
- "TFElectraForSequenceClassification",
- "TFElectraForTokenClassification",
- "TFElectraModel",
- "TFElectraPreTrainedModel",
- ]
-
-try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flax_electra"] = [
- "FlaxElectraForCausalLM",
- "FlaxElectraForMaskedLM",
- "FlaxElectraForMultipleChoice",
- "FlaxElectraForPreTraining",
- "FlaxElectraForQuestionAnswering",
- "FlaxElectraForSequenceClassification",
- "FlaxElectraForTokenClassification",
- "FlaxElectraModel",
- "FlaxElectraPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraOnnxConfig
- from .tokenization_electra import ElectraTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_electra_fast import ElectraTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_electra import (
- ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
- ElectraForCausalLM,
- ElectraForMaskedLM,
- ElectraForMultipleChoice,
- ElectraForPreTraining,
- ElectraForQuestionAnswering,
- ElectraForSequenceClassification,
- ElectraForTokenClassification,
- ElectraModel,
- ElectraPreTrainedModel,
- load_tf_weights_in_electra,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_electra import (
- TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFElectraForMaskedLM,
- TFElectraForMultipleChoice,
- TFElectraForPreTraining,
- TFElectraForQuestionAnswering,
- TFElectraForSequenceClassification,
- TFElectraForTokenClassification,
- TFElectraModel,
- TFElectraPreTrainedModel,
- )
-
- try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flax_electra import (
- FlaxElectraForCausalLM,
- FlaxElectraForMaskedLM,
- FlaxElectraForMultipleChoice,
- FlaxElectraForPreTraining,
- FlaxElectraForQuestionAnswering,
- FlaxElectraForSequenceClassification,
- FlaxElectraForTokenClassification,
- FlaxElectraModel,
- FlaxElectraPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/electra/__pycache__/__init__.cpython-310.pyc b/transformers/models/electra/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 2c175d8114338d05b36a4f9e00458dc8e37ad7d6..0000000000000000000000000000000000000000
Binary files a/transformers/models/electra/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/electra/__pycache__/configuration_electra.cpython-310.pyc b/transformers/models/electra/__pycache__/configuration_electra.cpython-310.pyc
deleted file mode 100644
index 74e8b1f6fea2384a3ec873caf1d29915a7a6876e..0000000000000000000000000000000000000000
Binary files a/transformers/models/electra/__pycache__/configuration_electra.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/electra/__pycache__/convert_electra_original_tf_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/electra/__pycache__/convert_electra_original_tf_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 54636bddc7ef7fde7b5728a532ced211fda7caa7..0000000000000000000000000000000000000000
Binary files a/transformers/models/electra/__pycache__/convert_electra_original_tf_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/electra/__pycache__/modeling_electra.cpython-310.pyc b/transformers/models/electra/__pycache__/modeling_electra.cpython-310.pyc
deleted file mode 100644
index 045e92cc613aa3d0efbab45d42f94a65ea899d62..0000000000000000000000000000000000000000
Binary files a/transformers/models/electra/__pycache__/modeling_electra.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/electra/__pycache__/modeling_flax_electra.cpython-310.pyc b/transformers/models/electra/__pycache__/modeling_flax_electra.cpython-310.pyc
deleted file mode 100644
index 827a07136aeb20405a7140812df8fc2566193b10..0000000000000000000000000000000000000000
Binary files a/transformers/models/electra/__pycache__/modeling_flax_electra.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/electra/__pycache__/modeling_tf_electra.cpython-310.pyc b/transformers/models/electra/__pycache__/modeling_tf_electra.cpython-310.pyc
deleted file mode 100644
index 7fdf28e9c18ccd339332de70d5dded2de66a7b0b..0000000000000000000000000000000000000000
Binary files a/transformers/models/electra/__pycache__/modeling_tf_electra.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/electra/__pycache__/tokenization_electra.cpython-310.pyc b/transformers/models/electra/__pycache__/tokenization_electra.cpython-310.pyc
deleted file mode 100644
index 1e4b892df52d2890a7eb6b03c6c81a3801be2425..0000000000000000000000000000000000000000
Binary files a/transformers/models/electra/__pycache__/tokenization_electra.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/electra/__pycache__/tokenization_electra_fast.cpython-310.pyc b/transformers/models/electra/__pycache__/tokenization_electra_fast.cpython-310.pyc
deleted file mode 100644
index 5b95f43e34c6940622431fe44740f35d7bf4009f..0000000000000000000000000000000000000000
Binary files a/transformers/models/electra/__pycache__/tokenization_electra_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/electra/configuration_electra.py b/transformers/models/electra/configuration_electra.py
deleted file mode 100644
index b6d1368a9d22d2278b08f080a8b08fd6df5498de..0000000000000000000000000000000000000000
--- a/transformers/models/electra/configuration_electra.py
+++ /dev/null
@@ -1,187 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" ELECTRA model configuration"""
-
-from collections import OrderedDict
-from typing import Mapping
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class ElectraConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`ElectraModel`] or a [`TFElectraModel`]. It is
- used to instantiate a ELECTRA model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the ELECTRA
- [google/electra-small-discriminator](https://huggingface.co/google/electra-small-discriminator) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the ELECTRA model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`ElectraModel`] or [`TFElectraModel`].
- embedding_size (`int`, *optional*, defaults to 128):
- Dimensionality of the encoder layers and the pooler layer.
- hidden_size (`int`, *optional*, defaults to 256):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 4):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 1024):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 2):
- The vocabulary size of the `token_type_ids` passed when calling [`ElectraModel`] or [`TFElectraModel`].
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- summary_type (`str`, *optional*, defaults to `"first"`):
- Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
-
- Has to be one of the following options:
-
- - `"last"`: Take the last token hidden state (like XLNet).
- - `"first"`: Take the first token hidden state (like BERT).
- - `"mean"`: Take the mean of all tokens hidden states.
- - `"cls_index"`: Supply a Tensor of classification token position (like GPT/GPT-2).
- - `"attn"`: Not implemented now, use multi-head attention.
- summary_use_proj (`bool`, *optional*, defaults to `True`):
- Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
-
- Whether or not to add a projection after the vector extraction.
- summary_activation (`str`, *optional*):
- Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
-
- Pass `"gelu"` for a gelu activation to the output, any other value will result in no activation.
- summary_last_dropout (`float`, *optional*, defaults to 0.0):
- Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
-
- The dropout ratio to be used after the projection and activation.
- position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
- Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
- positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
- [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
- For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
- with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- classifier_dropout (`float`, *optional*):
- The dropout ratio for the classification head.
-
- Examples:
-
- ```python
- >>> from transformers import ElectraConfig, ElectraModel
-
- >>> # Initializing a ELECTRA electra-base-uncased style configuration
- >>> configuration = ElectraConfig()
-
- >>> # Initializing a model (with random weights) from the electra-base-uncased style configuration
- >>> model = ElectraModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "electra"
-
- def __init__(
- self,
- vocab_size=30522,
- embedding_size=128,
- hidden_size=256,
- num_hidden_layers=12,
- num_attention_heads=4,
- intermediate_size=1024,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=2,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- summary_type="first",
- summary_use_proj=True,
- summary_activation="gelu",
- summary_last_dropout=0.1,
- pad_token_id=0,
- position_embedding_type="absolute",
- use_cache=True,
- classifier_dropout=None,
- **kwargs,
- ):
- super().__init__(pad_token_id=pad_token_id, **kwargs)
-
- self.vocab_size = vocab_size
- self.embedding_size = embedding_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
-
- self.summary_type = summary_type
- self.summary_use_proj = summary_use_proj
- self.summary_activation = summary_activation
- self.summary_last_dropout = summary_last_dropout
- self.position_embedding_type = position_embedding_type
- self.use_cache = use_cache
- self.classifier_dropout = classifier_dropout
-
-
-class ElectraOnnxConfig(OnnxConfig):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- if self.task == "multiple-choice":
- dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
- else:
- dynamic_axis = {0: "batch", 1: "sequence"}
- return OrderedDict(
- [
- ("input_ids", dynamic_axis),
- ("attention_mask", dynamic_axis),
- ("token_type_ids", dynamic_axis),
- ]
- )
diff --git a/transformers/models/electra/convert_electra_original_tf_checkpoint_to_pytorch.py b/transformers/models/electra/convert_electra_original_tf_checkpoint_to_pytorch.py
deleted file mode 100644
index d5d6376d7b994281b8743d54baa8c4c23db9c05b..0000000000000000000000000000000000000000
--- a/transformers/models/electra/convert_electra_original_tf_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,80 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert ELECTRA checkpoint."""
-
-
-import argparse
-
-import torch
-
-from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining, load_tf_weights_in_electra
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-
-
-def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path, discriminator_or_generator):
- # Initialise PyTorch model
- config = ElectraConfig.from_json_file(config_file)
- print(f"Building PyTorch model from configuration: {config}")
-
- if discriminator_or_generator == "discriminator":
- model = ElectraForPreTraining(config)
- elif discriminator_or_generator == "generator":
- model = ElectraForMaskedLM(config)
- else:
- raise ValueError("The discriminator_or_generator argument should be either 'discriminator' or 'generator'")
-
- # Load weights from tf checkpoint
- load_tf_weights_in_electra(
- model, config, tf_checkpoint_path, discriminator_or_generator=discriminator_or_generator
- )
-
- # Save pytorch-model
- print(f"Save PyTorch model to {pytorch_dump_path}")
- torch.save(model.state_dict(), pytorch_dump_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
- )
- parser.add_argument(
- "--config_file",
- default=None,
- type=str,
- required=True,
- help="The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.",
- )
- parser.add_argument(
- "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
- parser.add_argument(
- "--discriminator_or_generator",
- default=None,
- type=str,
- required=True,
- help=(
- "Whether to export the generator or the discriminator. Should be a string, either 'discriminator' or "
- "'generator'."
- ),
- )
- args = parser.parse_args()
- convert_tf_checkpoint_to_pytorch(
- args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
- )
diff --git a/transformers/models/electra/modeling_electra.py b/transformers/models/electra/modeling_electra.py
deleted file mode 100644
index 2138aa97c6dca956f57be1cb71fd10f9544f3e45..0000000000000000000000000000000000000000
--- a/transformers/models/electra/modeling_electra.py
+++ /dev/null
@@ -1,1679 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Google AI Language Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch ELECTRA model."""
-
-import math
-import os
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN, get_activation
-from ...modeling_outputs import (
- BaseModelOutputWithCrossAttentions,
- BaseModelOutputWithPastAndCrossAttentions,
- CausalLMOutputWithCrossAttentions,
- MaskedLMOutput,
- MultipleChoiceModelOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel, SequenceSummary
-from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_electra import ElectraConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "google/electra-small-discriminator"
-_CONFIG_FOR_DOC = "ElectraConfig"
-
-
-from ..deprecated._archive_maps import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-def load_tf_weights_in_electra(model, config, tf_checkpoint_path, discriminator_or_generator="discriminator"):
- """Load tf checkpoints in a pytorch model."""
- try:
- import re
-
- import numpy as np
- import tensorflow as tf
- except ImportError:
- logger.error(
- "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
- "https://www.tensorflow.org/install/ for installation instructions."
- )
- raise
- tf_path = os.path.abspath(tf_checkpoint_path)
- logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
- # Load weights from TF model
- init_vars = tf.train.list_variables(tf_path)
- names = []
- arrays = []
- for name, shape in init_vars:
- logger.info(f"Loading TF weight {name} with shape {shape}")
- array = tf.train.load_variable(tf_path, name)
- names.append(name)
- arrays.append(array)
- for name, array in zip(names, arrays):
- original_name: str = name
-
- try:
- if isinstance(model, ElectraForMaskedLM):
- name = name.replace("electra/embeddings/", "generator/embeddings/")
-
- if discriminator_or_generator == "generator":
- name = name.replace("electra/", "discriminator/")
- name = name.replace("generator/", "electra/")
-
- name = name.replace("dense_1", "dense_prediction")
- name = name.replace("generator_predictions/output_bias", "generator_lm_head/bias")
-
- name = name.split("/")
- # print(original_name, name)
- # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
- # which are not required for using pretrained model
- if any(n in ["global_step", "temperature"] for n in name):
- logger.info(f"Skipping {original_name}")
- continue
- pointer = model
- for m_name in name:
- if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
- scope_names = re.split(r"_(\d+)", m_name)
- else:
- scope_names = [m_name]
- if scope_names[0] == "kernel" or scope_names[0] == "gamma":
- pointer = getattr(pointer, "weight")
- elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
- pointer = getattr(pointer, "bias")
- elif scope_names[0] == "output_weights":
- pointer = getattr(pointer, "weight")
- elif scope_names[0] == "squad":
- pointer = getattr(pointer, "classifier")
- else:
- pointer = getattr(pointer, scope_names[0])
- if len(scope_names) >= 2:
- num = int(scope_names[1])
- pointer = pointer[num]
- if m_name.endswith("_embeddings"):
- pointer = getattr(pointer, "weight")
- elif m_name == "kernel":
- array = np.transpose(array)
- try:
- if pointer.shape != array.shape:
- raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
- except ValueError as e:
- e.args += (pointer.shape, array.shape)
- raise
- print(f"Initialize PyTorch weight {name}", original_name)
- pointer.data = torch.from_numpy(array)
- except AttributeError as e:
- print(f"Skipping {original_name}", name, e)
- continue
- return model
-
-
-class ElectraEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.embedding_size)
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.embedding_size)
-
- # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
- # any TensorFlow checkpoint file
- self.LayerNorm = nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
- self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
- self.register_buffer(
- "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
- )
-
- # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.forward
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- past_key_values_length: int = 0,
- ) -> torch.Tensor:
- if input_ids is not None:
- input_shape = input_ids.size()
- else:
- input_shape = inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
-
- # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
- # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
- # issue #5664
- if token_type_ids is None:
- if hasattr(self, "token_type_ids"):
- buffered_token_type_ids = self.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
- embeddings = inputs_embeds + token_type_embeddings
- if self.position_embedding_type == "absolute":
- position_embeddings = self.position_embeddings(position_ids)
- embeddings += position_embeddings
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Electra
-class ElectraSelfAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size)
- self.key = nn.Linear(config.hidden_size, self.all_head_size)
- self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
- self.position_embedding_type = position_embedding_type or getattr(
- config, "position_embedding_type", "absolute"
- )
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- self.max_position_embeddings = config.max_position_embeddings
- self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
-
- self.is_decoder = config.is_decoder
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- mixed_query_layer = self.query(hidden_states)
-
- # If this is instantiated as a cross-attention module, the keys
- # and values come from an encoder; the attention mask needs to be
- # such that the encoder's padding tokens are not attended to.
- is_cross_attention = encoder_hidden_states is not None
-
- if is_cross_attention and past_key_value is not None:
- # reuse k,v, cross_attentions
- key_layer = past_key_value[0]
- value_layer = past_key_value[1]
- attention_mask = encoder_attention_mask
- elif is_cross_attention:
- key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
- value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
- attention_mask = encoder_attention_mask
- elif past_key_value is not None:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
- value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
- else:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- use_cache = past_key_value is not None
- if self.is_decoder:
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_layer, value_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- query_length, key_length = query_layer.shape[2], key_layer.shape[2]
- if use_cache:
- position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(
- -1, 1
- )
- else:
- position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
- position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
- distance = position_ids_l - position_ids_r
-
- positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
- positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
-
- if self.position_embedding_type == "relative_key":
- relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores
- elif self.position_embedding_type == "relative_key_query":
- relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in ElectraModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- if self.is_decoder:
- outputs = outputs + (past_key_value,)
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
-class ElectraSelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Electra
-class ElectraAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- self.self = ElectraSelfAttention(config, position_embedding_type=position_embedding_type)
- self.output = ElectraSelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- self_outputs = self.self(
- hidden_states,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertIntermediate
-class ElectraIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOutput
-class ElectraOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Electra
-class ElectraLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = ElectraAttention(config)
- self.is_decoder = config.is_decoder
- self.add_cross_attention = config.add_cross_attention
- if self.add_cross_attention:
- if not self.is_decoder:
- raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
- self.crossattention = ElectraAttention(config, position_embedding_type="absolute")
- self.intermediate = ElectraIntermediate(config)
- self.output = ElectraOutput(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
- self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
- self_attention_outputs = self.attention(
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions=output_attentions,
- past_key_value=self_attn_past_key_value,
- )
- attention_output = self_attention_outputs[0]
-
- # if decoder, the last output is tuple of self-attn cache
- if self.is_decoder:
- outputs = self_attention_outputs[1:-1]
- present_key_value = self_attention_outputs[-1]
- else:
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- cross_attn_present_key_value = None
- if self.is_decoder and encoder_hidden_states is not None:
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
- " by setting `config.add_cross_attention=True`"
- )
-
- # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
- cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
- cross_attention_outputs = self.crossattention(
- attention_output,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- cross_attn_past_key_value,
- output_attentions,
- )
- attention_output = cross_attention_outputs[0]
- outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
-
- # add cross-attn cache to positions 3,4 of present_key_value tuple
- cross_attn_present_key_value = cross_attention_outputs[-1]
- present_key_value = present_key_value + cross_attn_present_key_value
-
- layer_output = apply_chunking_to_forward(
- self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
- )
- outputs = (layer_output,) + outputs
-
- # if decoder, return the attn key/values as the last output
- if self.is_decoder:
- outputs = outputs + (present_key_value,)
-
- return outputs
-
- def feed_forward_chunk(self, attention_output):
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.output(intermediate_output, attention_output)
- return layer_output
-
-
-# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Electra
-class ElectraEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([ElectraLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = False,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- next_decoder_cache = () if use_cache else None
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
- past_key_value = past_key_values[i] if past_key_values is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
-
- hidden_states = layer_outputs[0]
- if use_cache:
- next_decoder_cache += (layer_outputs[-1],)
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- next_decoder_cache,
- all_hidden_states,
- all_self_attentions,
- all_cross_attentions,
- ]
- if v is not None
- )
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=next_decoder_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-class ElectraDiscriminatorPredictions(nn.Module):
- """Prediction module for the discriminator, made up of two dense layers."""
-
- def __init__(self, config):
- super().__init__()
-
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = get_activation(config.hidden_act)
- self.dense_prediction = nn.Linear(config.hidden_size, 1)
- self.config = config
-
- def forward(self, discriminator_hidden_states):
- hidden_states = self.dense(discriminator_hidden_states)
- hidden_states = self.activation(hidden_states)
- logits = self.dense_prediction(hidden_states).squeeze(-1)
-
- return logits
-
-
-class ElectraGeneratorPredictions(nn.Module):
- """Prediction module for the generator, made up of two dense layers."""
-
- def __init__(self, config):
- super().__init__()
-
- self.activation = get_activation("gelu")
- self.LayerNorm = nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
- self.dense = nn.Linear(config.hidden_size, config.embedding_size)
-
- def forward(self, generator_hidden_states):
- hidden_states = self.dense(generator_hidden_states)
- hidden_states = self.activation(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
-
- return hidden_states
-
-
-class ElectraPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ElectraConfig
- load_tf_weights = load_tf_weights_in_electra
- base_model_prefix = "electra"
- supports_gradient_checkpointing = True
-
- # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-@dataclass
-class ElectraForPreTrainingOutput(ModelOutput):
- """
- Output type of [`ElectraForPreTraining`].
-
- Args:
- loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
- Total loss of the ELECTRA objective.
- logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
- Prediction scores of the head (scores for each token before SoftMax).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-ELECTRA_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`ElectraConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ELECTRA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- encoder_hidden_states (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Electra Model transformer outputting raw hidden-states without any specific head on top. Identical to "
- "the BERT model except that it uses an additional linear layer between the embedding layer and the encoder if the "
- "hidden size and embedding size are different. "
- ""
- "Both the generator and discriminator checkpoints may be loaded into this model.",
- ELECTRA_START_DOCSTRING,
-)
-class ElectraModel(ElectraPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.embeddings = ElectraEmbeddings(config)
-
- if config.embedding_size != config.hidden_size:
- self.embeddings_project = nn.Linear(config.embedding_size, config.hidden_size)
-
- self.encoder = ElectraEncoder(config)
- self.config = config
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithCrossAttentions]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- # past_key_values_length
- past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device)
- if token_type_ids is None:
- if hasattr(self.embeddings, "token_type_ids"):
- buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)
-
- # If a 2D or 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- if self.config.is_decoder and encoder_hidden_states is not None:
- encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
- encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
- if encoder_attention_mask is None:
- encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
- encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
- else:
- encoder_extended_attention_mask = None
-
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- hidden_states = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- past_key_values_length=past_key_values_length,
- )
-
- if hasattr(self, "embeddings_project"):
- hidden_states = self.embeddings_project(hidden_states)
-
- hidden_states = self.encoder(
- hidden_states,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_extended_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- return hidden_states
-
-
-class ElectraClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.activation = get_activation("gelu")
- self.dropout = nn.Dropout(classifier_dropout)
- self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
-
- def forward(self, features, **kwargs):
- x = features[:, 0, :] # take token (equiv. to [CLS])
- x = self.dropout(x)
- x = self.dense(x)
- x = self.activation(x) # although BERT uses tanh here, it seems Electra authors used gelu here
- x = self.dropout(x)
- x = self.out_proj(x)
- return x
-
-
-@add_start_docstrings(
- """
- ELECTRA Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class ElectraForSequenceClassification(ElectraPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
- self.electra = ElectraModel(config)
- self.classifier = ElectraClassificationHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="bhadresh-savani/electra-base-emotion",
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output="'joy'",
- expected_loss=0.06,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- discriminator_hidden_states = self.electra(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = discriminator_hidden_states[0]
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + discriminator_hidden_states[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Electra model with a binary classification head on top as used during pretraining for identifying generated tokens.
-
- It is recommended to load the discriminator checkpoint into that model.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class ElectraForPreTraining(ElectraPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.electra = ElectraModel(config)
- self.discriminator_predictions = ElectraDiscriminatorPredictions(config)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=ElectraForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], ElectraForPreTrainingOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the ELECTRA loss. Input should be a sequence of tokens (see `input_ids` docstring)
- Indices should be in `[0, 1]`:
-
- - 0 indicates the token is an original token,
- - 1 indicates the token was replaced.
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import ElectraForPreTraining, AutoTokenizer
- >>> import torch
-
- >>> discriminator = ElectraForPreTraining.from_pretrained("google/electra-base-discriminator")
- >>> tokenizer = AutoTokenizer.from_pretrained("google/electra-base-discriminator")
-
- >>> sentence = "The quick brown fox jumps over the lazy dog"
- >>> fake_sentence = "The quick brown fox fake over the lazy dog"
-
- >>> fake_tokens = tokenizer.tokenize(fake_sentence, add_special_tokens=True)
- >>> fake_inputs = tokenizer.encode(fake_sentence, return_tensors="pt")
- >>> discriminator_outputs = discriminator(fake_inputs)
- >>> predictions = torch.round((torch.sign(discriminator_outputs[0]) + 1) / 2)
-
- >>> fake_tokens
- ['[CLS]', 'the', 'quick', 'brown', 'fox', 'fake', 'over', 'the', 'lazy', 'dog', '[SEP]']
-
- >>> predictions.squeeze().tolist()
- [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- discriminator_hidden_states = self.electra(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- discriminator_sequence_output = discriminator_hidden_states[0]
-
- logits = self.discriminator_predictions(discriminator_sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = nn.BCEWithLogitsLoss()
- if attention_mask is not None:
- active_loss = attention_mask.view(-1, discriminator_sequence_output.shape[1]) == 1
- active_logits = logits.view(-1, discriminator_sequence_output.shape[1])[active_loss]
- active_labels = labels[active_loss]
- loss = loss_fct(active_logits, active_labels.float())
- else:
- loss = loss_fct(logits.view(-1, discriminator_sequence_output.shape[1]), labels.float())
-
- if not return_dict:
- output = (logits,) + discriminator_hidden_states[1:]
- return ((loss,) + output) if loss is not None else output
-
- return ElectraForPreTrainingOutput(
- loss=loss,
- logits=logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Electra model with a language modeling head on top.
-
- Even though both the discriminator and generator may be loaded into this model, the generator is the only model of
- the two to have been trained for the masked language modeling task.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class ElectraForMaskedLM(ElectraPreTrainedModel):
- _tied_weights_keys = ["generator_lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.electra = ElectraModel(config)
- self.generator_predictions = ElectraGeneratorPredictions(config)
-
- self.generator_lm_head = nn.Linear(config.embedding_size, config.vocab_size)
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.generator_lm_head
-
- def set_output_embeddings(self, word_embeddings):
- self.generator_lm_head = word_embeddings
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="google/electra-small-generator",
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="[MASK]",
- expected_output="'paris'",
- expected_loss=1.22,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- generator_hidden_states = self.electra(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- generator_sequence_output = generator_hidden_states[0]
-
- prediction_scores = self.generator_predictions(generator_sequence_output)
- prediction_scores = self.generator_lm_head(prediction_scores)
-
- loss = None
- # Masked language modeling softmax layer
- if labels is not None:
- loss_fct = nn.CrossEntropyLoss() # -100 index = padding token
- loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + generator_hidden_states[1:]
- return ((loss,) + output) if loss is not None else output
-
- return MaskedLMOutput(
- loss=loss,
- logits=prediction_scores,
- hidden_states=generator_hidden_states.hidden_states,
- attentions=generator_hidden_states.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Electra model with a token classification head on top.
-
- Both the discriminator and generator may be loaded into this model.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class ElectraForTokenClassification(ElectraPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.electra = ElectraModel(config)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="bhadresh-savani/electra-base-discriminator-finetuned-conll03-english",
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output="['B-LOC', 'B-ORG', 'O', 'O', 'O', 'O', 'O', 'B-LOC', 'O', 'B-LOC', 'I-LOC']",
- expected_loss=0.11,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- discriminator_hidden_states = self.electra(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- discriminator_sequence_output = discriminator_hidden_states[0]
-
- discriminator_sequence_output = self.dropout(discriminator_sequence_output)
- logits = self.classifier(discriminator_sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + discriminator_hidden_states[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
-
-@add_start_docstrings(
- """
- ELECTRA Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- ELECTRA_START_DOCSTRING,
-)
-class ElectraForQuestionAnswering(ElectraPreTrainedModel):
- config_class = ElectraConfig
- base_model_prefix = "electra"
-
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.electra = ElectraModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="bhadresh-savani/electra-base-squad2",
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- qa_target_start_index=11,
- qa_target_end_index=12,
- expected_output="'a nice puppet'",
- expected_loss=2.64,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- discriminator_hidden_states = self.electra(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- )
-
- sequence_output = discriminator_hidden_states[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (
- start_logits,
- end_logits,
- ) + discriminator_hidden_states[1:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
-
-@add_start_docstrings(
- """
- ELECTRA Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class ElectraForMultipleChoice(ElectraPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.electra = ElectraModel(config)
- self.sequence_summary = SequenceSummary(config)
- self.classifier = nn.Linear(config.hidden_size, 1)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
- position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
- inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- discriminator_hidden_states = self.electra(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = discriminator_hidden_states[0]
-
- pooled_output = self.sequence_summary(sequence_output)
- logits = self.classifier(pooled_output)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + discriminator_hidden_states[1:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
-
-@add_start_docstrings(
- """ELECTRA Model with a `language modeling` head on top for CLM fine-tuning.""", ELECTRA_START_DOCSTRING
-)
-class ElectraForCausalLM(ElectraPreTrainedModel):
- _tied_weights_keys = ["generator_lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
-
- if not config.is_decoder:
- logger.warning("If you want to use `ElectraForCausalLM` as a standalone, add `is_decoder=True.`")
-
- self.electra = ElectraModel(config)
- self.generator_predictions = ElectraGeneratorPredictions(config)
- self.generator_lm_head = nn.Linear(config.embedding_size, config.vocab_size)
-
- self.init_weights()
-
- def get_output_embeddings(self):
- return self.generator_lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.generator_lm_head = new_embeddings
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.Tensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
- r"""
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
- `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
- ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, ElectraForCausalLM, ElectraConfig
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("google/electra-base-generator")
- >>> config = ElectraConfig.from_pretrained("google/electra-base-generator")
- >>> config.is_decoder = True
- >>> model = ElectraForCausalLM.from_pretrained("google/electra-base-generator", config=config)
-
- >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> prediction_logits = outputs.logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- if labels is not None:
- use_cache = False
-
- outputs = self.electra(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- prediction_scores = self.generator_lm_head(self.generator_predictions(sequence_output))
-
- lm_loss = None
- if labels is not None:
- # we are doing next-token prediction; shift prediction scores and input ids by one
- shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
- labels = labels[:, 1:].contiguous()
- loss_fct = CrossEntropyLoss()
- lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[1:]
- return ((lm_loss,) + output) if lm_loss is not None else output
-
- return CausalLMOutputWithCrossAttentions(
- loss=lm_loss,
- logits=prediction_scores,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- cross_attentions=outputs.cross_attentions,
- )
-
- # Copied from transformers.models.roberta.modeling_roberta.RobertaForCausalLM.prepare_inputs_for_generation
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs):
- input_shape = input_ids.shape
- # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
- if attention_mask is None:
- attention_mask = input_ids.new_ones(input_shape)
-
- # cut decoder_input_ids if past_key_values is used
- if past_key_values is not None:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
-
- return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values}
-
- # Copied from transformers.models.roberta.modeling_roberta.RobertaForCausalLM._reorder_cache
- def _reorder_cache(self, past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
- )
- return reordered_past
diff --git a/transformers/models/electra/modeling_flax_electra.py b/transformers/models/electra/modeling_flax_electra.py
deleted file mode 100644
index 64d49eb17a460ae0a8aca59c54cf0e1557122361..0000000000000000000000000000000000000000
--- a/transformers/models/electra/modeling_flax_electra.py
+++ /dev/null
@@ -1,1601 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The Google Flax Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Callable, Optional, Tuple
-
-import flax
-import flax.linen as nn
-import jax
-import jax.numpy as jnp
-import numpy as np
-from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
-from flax.linen import combine_masks, make_causal_mask
-from flax.linen import partitioning as nn_partitioning
-from flax.linen.attention import dot_product_attention_weights
-from flax.traverse_util import flatten_dict, unflatten_dict
-from jax import lax
-
-from ...modeling_flax_outputs import (
- FlaxBaseModelOutput,
- FlaxBaseModelOutputWithPastAndCrossAttentions,
- FlaxCausalLMOutputWithCrossAttentions,
- FlaxMaskedLMOutput,
- FlaxMultipleChoiceModelOutput,
- FlaxQuestionAnsweringModelOutput,
- FlaxSequenceClassifierOutput,
- FlaxTokenClassifierOutput,
-)
-from ...modeling_flax_utils import (
- ACT2FN,
- FlaxPreTrainedModel,
- append_call_sample_docstring,
- append_replace_return_docstrings,
- overwrite_call_docstring,
-)
-from ...utils import ModelOutput, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_electra import ElectraConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "google/electra-small-discriminator"
-_CONFIG_FOR_DOC = "ElectraConfig"
-
-remat = nn_partitioning.remat
-
-
-@flax.struct.dataclass
-class FlaxElectraForPreTrainingOutput(ModelOutput):
- """
- Output type of [`ElectraForPreTraining`].
-
- Args:
- logits (`jnp.ndarray` of shape `(batch_size, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
- hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- logits: jnp.ndarray = None
- hidden_states: Optional[Tuple[jnp.ndarray]] = None
- attentions: Optional[Tuple[jnp.ndarray]] = None
-
-
-ELECTRA_START_DOCSTRING = r"""
-
- This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading, saving and converting weights from PyTorch models)
-
- This model is also a Flax Linen
- [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a
- regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.
-
- Finally, this model supports inherent JAX features such as:
-
- - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
-
- Parameters:
- config ([`ElectraConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ELECTRA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`numpy.ndarray` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`numpy.ndarray` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`numpy.ndarray` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`numpy.ndarray` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
- head_mask (`numpy.ndarray` of shape `({0})`, `optional):
- Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-
-"""
-
-
-class FlaxElectraEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.word_embeddings = nn.Embed(
- self.config.vocab_size,
- self.config.embedding_size,
- embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.position_embeddings = nn.Embed(
- self.config.max_position_embeddings,
- self.config.embedding_size,
- embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.token_type_embeddings = nn.Embed(
- self.config.type_vocab_size,
- self.config.embedding_size,
- embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.LayerNorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)
- self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob)
-
- # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertEmbeddings.__call__
- def __call__(self, input_ids, token_type_ids, position_ids, attention_mask, deterministic: bool = True):
- # Embed
- inputs_embeds = self.word_embeddings(input_ids.astype("i4"))
- position_embeds = self.position_embeddings(position_ids.astype("i4"))
- token_type_embeddings = self.token_type_embeddings(token_type_ids.astype("i4"))
-
- # Sum all embeddings
- hidden_states = inputs_embeds + token_type_embeddings + position_embeds
-
- # Layer Norm
- hidden_states = self.LayerNorm(hidden_states)
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertSelfAttention with Bert->Electra
-class FlaxElectraSelfAttention(nn.Module):
- config: ElectraConfig
- causal: bool = False
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.head_dim = self.config.hidden_size // self.config.num_attention_heads
- if self.config.hidden_size % self.config.num_attention_heads != 0:
- raise ValueError(
- "`config.hidden_size`: {self.config.hidden_size} has to be a multiple of `config.num_attention_heads` "
- " : {self.config.num_attention_heads}"
- )
-
- self.query = nn.Dense(
- self.config.hidden_size,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
- )
- self.key = nn.Dense(
- self.config.hidden_size,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
- )
- self.value = nn.Dense(
- self.config.hidden_size,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
- )
-
- if self.causal:
- self.causal_mask = make_causal_mask(
- jnp.ones((1, self.config.max_position_embeddings), dtype="bool"), dtype="bool"
- )
-
- def _split_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.config.num_attention_heads, self.head_dim))
-
- def _merge_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.config.hidden_size,))
-
- @nn.compact
- # Copied from transformers.models.bart.modeling_flax_bart.FlaxBartAttention._concatenate_to_cache
- def _concatenate_to_cache(self, key, value, query, attention_mask):
- """
- This function takes projected key, value states from a single input token and concatenates the states to cached
- states from previous steps. This function is slighly adapted from the official Flax repository:
- https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252
- """
- # detect if we're initializing by absence of existing cache data.
- is_initialized = self.has_variable("cache", "cached_key")
- cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype)
- cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype)
- cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
-
- if is_initialized:
- *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape
- # update key, value caches with our new 1d spatial slices
- cur_index = cache_index.value
- indices = (0,) * len(batch_dims) + (cur_index, 0, 0)
- key = lax.dynamic_update_slice(cached_key.value, key, indices)
- value = lax.dynamic_update_slice(cached_value.value, value, indices)
- cached_key.value = key
- cached_value.value = value
- num_updated_cache_vectors = query.shape[1]
- cache_index.value = cache_index.value + num_updated_cache_vectors
- # causal mask for cached decoder self-attention: our single query position should only attend to those key positions that have already been generated and cached, not the remaining zero elements.
- pad_mask = jnp.broadcast_to(
- jnp.arange(max_length) < cur_index + num_updated_cache_vectors,
- tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),
- )
- attention_mask = combine_masks(pad_mask, attention_mask)
- return key, value, attention_mask
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- layer_head_mask,
- key_value_states: Optional[jnp.ndarray] = None,
- init_cache: bool = False,
- deterministic=True,
- output_attentions: bool = False,
- ):
- # if key_value_states are provided this layer is used as a cross-attention layer
- # for the decoder
- is_cross_attention = key_value_states is not None
- batch_size = hidden_states.shape[0]
-
- # get query proj
- query_states = self.query(hidden_states)
- # get key, value proj
- if is_cross_attention:
- # cross_attentions
- key_states = self.key(key_value_states)
- value_states = self.value(key_value_states)
- else:
- # self_attention
- key_states = self.key(hidden_states)
- value_states = self.value(hidden_states)
-
- query_states = self._split_heads(query_states)
- key_states = self._split_heads(key_states)
- value_states = self._split_heads(value_states)
-
- # handle cache prepare causal attention mask
- if self.causal:
- query_length, key_length = query_states.shape[1], key_states.shape[1]
- if self.has_variable("cache", "cached_key"):
- mask_shift = self.variables["cache"]["cache_index"]
- max_decoder_length = self.variables["cache"]["cached_key"].shape[1]
- causal_mask = lax.dynamic_slice(
- self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)
- )
- else:
- causal_mask = self.causal_mask[:, :, :query_length, :key_length]
- causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])
-
- # combine masks if needed
- if attention_mask is not None and self.causal:
- attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
- attention_mask = combine_masks(attention_mask, causal_mask)
- elif self.causal:
- attention_mask = causal_mask
- elif attention_mask is not None:
- attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))
-
- # During fast autoregressive decoding, we feed one position at a time,
- # and cache the keys and values step by step.
- if self.causal and (self.has_variable("cache", "cached_key") or init_cache):
- key_states, value_states, attention_mask = self._concatenate_to_cache(
- key_states, value_states, query_states, attention_mask
- )
-
- # Convert the boolean attention mask to an attention bias.
- if attention_mask is not None:
- # attention mask in the form of attention bias
- attention_bias = lax.select(
- attention_mask > 0,
- jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
- jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
- )
- else:
- attention_bias = None
-
- dropout_rng = None
- if not deterministic and self.config.attention_probs_dropout_prob > 0.0:
- dropout_rng = self.make_rng("dropout")
-
- attn_weights = dot_product_attention_weights(
- query_states,
- key_states,
- bias=attention_bias,
- dropout_rng=dropout_rng,
- dropout_rate=self.config.attention_probs_dropout_prob,
- broadcast_dropout=True,
- deterministic=deterministic,
- dtype=self.dtype,
- precision=None,
- )
-
- # Mask heads if we want to
- if layer_head_mask is not None:
- attn_weights = jnp.einsum("...hqk,h->...hqk", attn_weights, layer_head_mask)
-
- attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states)
- attn_output = attn_output.reshape(attn_output.shape[:2] + (-1,))
-
- outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertSelfOutput with Bert->Electra
-class FlaxElectraSelfOutput(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.dense = nn.Dense(
- self.config.hidden_size,
- kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
- dtype=self.dtype,
- )
- self.LayerNorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)
- self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob)
-
- def __call__(self, hidden_states, input_tensor, deterministic: bool = True):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertAttention with Bert->Electra
-class FlaxElectraAttention(nn.Module):
- config: ElectraConfig
- causal: bool = False
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.self = FlaxElectraSelfAttention(self.config, causal=self.causal, dtype=self.dtype)
- self.output = FlaxElectraSelfOutput(self.config, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- layer_head_mask,
- key_value_states=None,
- init_cache=False,
- deterministic=True,
- output_attentions: bool = False,
- ):
- # Attention mask comes in as attention_mask.shape == (*batch_sizes, kv_length)
- # FLAX expects: attention_mask.shape == (*batch_sizes, 1, 1, kv_length) such that it is broadcastable
- # with attn_weights.shape == (*batch_sizes, num_heads, q_length, kv_length)
- attn_outputs = self.self(
- hidden_states,
- attention_mask,
- layer_head_mask=layer_head_mask,
- key_value_states=key_value_states,
- init_cache=init_cache,
- deterministic=deterministic,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0]
- hidden_states = self.output(attn_output, hidden_states, deterministic=deterministic)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attn_outputs[1],)
-
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertIntermediate with Bert->Electra
-class FlaxElectraIntermediate(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.dense = nn.Dense(
- self.config.intermediate_size,
- kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
- dtype=self.dtype,
- )
- self.activation = ACT2FN[self.config.hidden_act]
-
- def __call__(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.activation(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertOutput with Bert->Electra
-class FlaxElectraOutput(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.dense = nn.Dense(
- self.config.hidden_size,
- kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
- dtype=self.dtype,
- )
- self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob)
- self.LayerNorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)
-
- def __call__(self, hidden_states, attention_output, deterministic: bool = True):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- hidden_states = self.LayerNorm(hidden_states + attention_output)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertLayer with Bert->Electra
-class FlaxElectraLayer(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
-
- def setup(self):
- self.attention = FlaxElectraAttention(self.config, causal=self.config.is_decoder, dtype=self.dtype)
- self.intermediate = FlaxElectraIntermediate(self.config, dtype=self.dtype)
- self.output = FlaxElectraOutput(self.config, dtype=self.dtype)
- if self.config.add_cross_attention:
- self.crossattention = FlaxElectraAttention(self.config, causal=False, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- init_cache: bool = False,
- deterministic: bool = True,
- output_attentions: bool = False,
- ):
- # Self Attention
- attention_outputs = self.attention(
- hidden_states,
- attention_mask,
- layer_head_mask=layer_head_mask,
- init_cache=init_cache,
- deterministic=deterministic,
- output_attentions=output_attentions,
- )
- attention_output = attention_outputs[0]
-
- # Cross-Attention Block
- if encoder_hidden_states is not None:
- cross_attention_outputs = self.crossattention(
- attention_output,
- attention_mask=encoder_attention_mask,
- layer_head_mask=layer_head_mask,
- key_value_states=encoder_hidden_states,
- deterministic=deterministic,
- output_attentions=output_attentions,
- )
- attention_output = cross_attention_outputs[0]
-
- hidden_states = self.intermediate(attention_output)
- hidden_states = self.output(hidden_states, attention_output, deterministic=deterministic)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attention_outputs[1],)
- if encoder_hidden_states is not None:
- outputs += (cross_attention_outputs[1],)
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertLayerCollection with Bert->Electra
-class FlaxElectraLayerCollection(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
- gradient_checkpointing: bool = False
-
- def setup(self):
- if self.gradient_checkpointing:
- FlaxElectraCheckpointLayer = remat(FlaxElectraLayer, static_argnums=(5, 6, 7))
- self.layers = [
- FlaxElectraCheckpointLayer(self.config, name=str(i), dtype=self.dtype)
- for i in range(self.config.num_hidden_layers)
- ]
- else:
- self.layers = [
- FlaxElectraLayer(self.config, name=str(i), dtype=self.dtype)
- for i in range(self.config.num_hidden_layers)
- ]
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- head_mask,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- init_cache: bool = False,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- all_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
- all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
-
- # Check if head_mask has a correct number of layers specified if desired
- if head_mask is not None:
- if head_mask.shape[0] != (len(self.layers)):
- raise ValueError(
- f"The head_mask should be specified for {len(self.layers)} layers, but it is for "
- f" {head_mask.shape[0]}."
- )
-
- for i, layer in enumerate(self.layers):
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- layer_outputs = layer(
- hidden_states,
- attention_mask,
- head_mask[i] if head_mask is not None else None,
- encoder_hidden_states,
- encoder_attention_mask,
- init_cache,
- deterministic,
- output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions += (layer_outputs[1],)
-
- if encoder_hidden_states is not None:
- all_cross_attentions += (layer_outputs[2],)
-
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
-
- if not return_dict:
- return tuple(v for v in outputs if v is not None)
-
- return FlaxBaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertEncoder with Bert->Electra
-class FlaxElectraEncoder(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.layer = FlaxElectraLayerCollection(
- self.config,
- dtype=self.dtype,
- gradient_checkpointing=self.gradient_checkpointing,
- )
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- head_mask,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- init_cache: bool = False,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- return self.layer(
- hidden_states,
- attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- init_cache=init_cache,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
-
-class FlaxElectraGeneratorPredictions(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.LayerNorm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)
- self.dense = nn.Dense(self.config.embedding_size, dtype=self.dtype)
-
- def __call__(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = ACT2FN[self.config.hidden_act](hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
- return hidden_states
-
-
-class FlaxElectraDiscriminatorPredictions(nn.Module):
- """Prediction module for the discriminator, made up of two dense layers."""
-
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.dense = nn.Dense(self.config.hidden_size, dtype=self.dtype)
- self.dense_prediction = nn.Dense(1, dtype=self.dtype)
-
- def __call__(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = ACT2FN[self.config.hidden_act](hidden_states)
- hidden_states = self.dense_prediction(hidden_states).squeeze(-1)
- return hidden_states
-
-
-class FlaxElectraPreTrainedModel(FlaxPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ElectraConfig
- base_model_prefix = "electra"
- module_class: nn.Module = None
-
- def __init__(
- self,
- config: ElectraConfig,
- input_shape: Tuple = (1, 1),
- seed: int = 0,
- dtype: jnp.dtype = jnp.float32,
- _do_init: bool = True,
- gradient_checkpointing: bool = False,
- **kwargs,
- ):
- module = self.module_class(config=config, dtype=dtype, gradient_checkpointing=gradient_checkpointing, **kwargs)
- super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
-
- # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertPreTrainedModel.enable_gradient_checkpointing
- def enable_gradient_checkpointing(self):
- self._module = self.module_class(
- config=self.config,
- dtype=self.dtype,
- gradient_checkpointing=True,
- )
-
- # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertPreTrainedModel.init_weights
- def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
- # init input tensors
- input_ids = jnp.zeros(input_shape, dtype="i4")
- token_type_ids = jnp.zeros_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape)
- attention_mask = jnp.ones_like(input_ids)
- head_mask = jnp.ones((self.config.num_hidden_layers, self.config.num_attention_heads))
-
- params_rng, dropout_rng = jax.random.split(rng)
- rngs = {"params": params_rng, "dropout": dropout_rng}
-
- if self.config.add_cross_attention:
- encoder_hidden_states = jnp.zeros(input_shape + (self.config.hidden_size,))
- encoder_attention_mask = attention_mask
- module_init_outputs = self.module.init(
- rngs,
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- return_dict=False,
- )
- else:
- module_init_outputs = self.module.init(
- rngs, input_ids, attention_mask, token_type_ids, position_ids, head_mask, return_dict=False
- )
-
- random_params = module_init_outputs["params"]
-
- if params is not None:
- random_params = flatten_dict(unfreeze(random_params))
- params = flatten_dict(unfreeze(params))
- for missing_key in self._missing_keys:
- params[missing_key] = random_params[missing_key]
- self._missing_keys = set()
- return freeze(unflatten_dict(params))
- else:
- return random_params
-
- # Copied from transformers.models.bart.modeling_flax_bart.FlaxBartDecoderPreTrainedModel.init_cache
- def init_cache(self, batch_size, max_length):
- r"""
- Args:
- batch_size (`int`):
- batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
- max_length (`int`):
- maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
- cache.
- """
- # init input variables to retrieve cache
- input_ids = jnp.ones((batch_size, max_length), dtype="i4")
- attention_mask = jnp.ones_like(input_ids, dtype="i4")
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
-
- init_variables = self.module.init(
- jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True
- )
- return unfreeze(init_variables["cache"])
-
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- params: dict = None,
- dropout_rng: jax.random.PRNGKey = None,
- train: bool = False,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- past_key_values: dict = None,
- ):
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- # init input tensors if not passed
- if token_type_ids is None:
- token_type_ids = jnp.ones_like(input_ids)
-
- if position_ids is None:
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
-
- if attention_mask is None:
- attention_mask = jnp.ones_like(input_ids)
-
- if head_mask is None:
- head_mask = jnp.ones((self.config.num_hidden_layers, self.config.num_attention_heads))
-
- # Handle any PRNG if needed
- rngs = {}
- if dropout_rng is not None:
- rngs["dropout"] = dropout_rng
-
- inputs = {"params": params or self.params}
-
- if self.config.add_cross_attention:
- # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed
- # down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be
- # changed by FlaxElectraAttention module
- if past_key_values:
- inputs["cache"] = past_key_values
- mutable = ["cache"]
- else:
- mutable = False
-
- outputs = self.module.apply(
- inputs,
- jnp.array(input_ids, dtype="i4"),
- jnp.array(attention_mask, dtype="i4"),
- token_type_ids=jnp.array(token_type_ids, dtype="i4"),
- position_ids=jnp.array(position_ids, dtype="i4"),
- head_mask=jnp.array(head_mask, dtype="i4"),
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- deterministic=not train,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- rngs=rngs,
- mutable=mutable,
- )
-
- # add updated cache to model output
- if past_key_values is not None and return_dict:
- outputs, past_key_values = outputs
- outputs["past_key_values"] = unfreeze(past_key_values["cache"])
- return outputs
- elif past_key_values is not None and not return_dict:
- outputs, past_key_values = outputs
- outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:]
-
- else:
- outputs = self.module.apply(
- inputs,
- jnp.array(input_ids, dtype="i4"),
- jnp.array(attention_mask, dtype="i4"),
- token_type_ids=jnp.array(token_type_ids, dtype="i4"),
- position_ids=jnp.array(position_ids, dtype="i4"),
- head_mask=jnp.array(head_mask, dtype="i4"),
- deterministic=not train,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- rngs=rngs,
- )
-
- return outputs
-
-
-class FlaxElectraModule(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32 # the dtype of the computation
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.embeddings = FlaxElectraEmbeddings(self.config, dtype=self.dtype)
- if self.config.embedding_size != self.config.hidden_size:
- self.embeddings_project = nn.Dense(self.config.hidden_size, dtype=self.dtype)
- self.encoder = FlaxElectraEncoder(
- self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
- )
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask: Optional[np.ndarray] = None,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- init_cache: bool = False,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- embeddings = self.embeddings(
- input_ids, token_type_ids, position_ids, attention_mask, deterministic=deterministic
- )
- if hasattr(self, "embeddings_project"):
- embeddings = self.embeddings_project(embeddings)
-
- return self.encoder(
- embeddings,
- attention_mask,
- head_mask=head_mask,
- deterministic=deterministic,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
-
-@add_start_docstrings(
- "The bare Electra Model transformer outputting raw hidden-states without any specific head on top.",
- ELECTRA_START_DOCSTRING,
-)
-class FlaxElectraModel(FlaxElectraPreTrainedModel):
- module_class = FlaxElectraModule
-
-
-append_call_sample_docstring(FlaxElectraModel, _CHECKPOINT_FOR_DOC, FlaxBaseModelOutput, _CONFIG_FOR_DOC)
-
-
-class FlaxElectraTiedDense(nn.Module):
- embedding_size: int
- dtype: jnp.dtype = jnp.float32
- precision = None
- bias_init: Callable[..., np.ndarray] = jax.nn.initializers.zeros
-
- def setup(self):
- self.bias = self.param("bias", self.bias_init, (self.embedding_size,))
-
- def __call__(self, x, kernel):
- x = jnp.asarray(x, self.dtype)
- kernel = jnp.asarray(kernel, self.dtype)
- y = lax.dot_general(
- x,
- kernel,
- (((x.ndim - 1,), (0,)), ((), ())),
- precision=self.precision,
- )
- bias = jnp.asarray(self.bias, self.dtype)
- return y + bias
-
-
-class FlaxElectraForMaskedLMModule(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.electra = FlaxElectraModule(
- config=self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
- )
- self.generator_predictions = FlaxElectraGeneratorPredictions(config=self.config, dtype=self.dtype)
- if self.config.tie_word_embeddings:
- self.generator_lm_head = FlaxElectraTiedDense(self.config.vocab_size, dtype=self.dtype)
- else:
- self.generator_lm_head = nn.Dense(self.config.vocab_size, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- outputs = self.electra(
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = outputs[0]
- prediction_scores = self.generator_predictions(hidden_states)
-
- if self.config.tie_word_embeddings:
- shared_embedding = self.electra.variables["params"]["embeddings"]["word_embeddings"]["embedding"]
- prediction_scores = self.generator_lm_head(prediction_scores, shared_embedding.T)
- else:
- prediction_scores = self.generator_lm_head(prediction_scores)
-
- if not return_dict:
- return (prediction_scores,) + outputs[1:]
-
- return FlaxMaskedLMOutput(
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings("""Electra Model with a `language modeling` head on top.""", ELECTRA_START_DOCSTRING)
-class FlaxElectraForMaskedLM(FlaxElectraPreTrainedModel):
- module_class = FlaxElectraForMaskedLMModule
-
-
-append_call_sample_docstring(FlaxElectraForMaskedLM, _CHECKPOINT_FOR_DOC, FlaxMaskedLMOutput, _CONFIG_FOR_DOC)
-
-
-class FlaxElectraForPreTrainingModule(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.electra = FlaxElectraModule(
- config=self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
- )
- self.discriminator_predictions = FlaxElectraDiscriminatorPredictions(config=self.config, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- # Model
- outputs = self.electra(
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = outputs[0]
-
- logits = self.discriminator_predictions(hidden_states)
-
- if not return_dict:
- return (logits,) + outputs[1:]
-
- return FlaxElectraForPreTrainingOutput(
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Electra model with a binary classification head on top as used during pretraining for identifying generated tokens.
-
- It is recommended to load the discriminator checkpoint into that model.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class FlaxElectraForPreTraining(FlaxElectraPreTrainedModel):
- module_class = FlaxElectraForPreTrainingModule
-
-
-FLAX_ELECTRA_FOR_PRETRAINING_DOCSTRING = """
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, FlaxElectraForPreTraining
-
- >>> tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")
- >>> model = FlaxElectraForPreTraining.from_pretrained("google/electra-small-discriminator")
-
- >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np")
- >>> outputs = model(**inputs)
-
- >>> prediction_logits = outputs.logits
- ```
-"""
-
-overwrite_call_docstring(
- FlaxElectraForPreTraining,
- ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length") + FLAX_ELECTRA_FOR_PRETRAINING_DOCSTRING,
-)
-append_replace_return_docstrings(
- FlaxElectraForPreTraining, output_type=FlaxElectraForPreTrainingOutput, config_class=_CONFIG_FOR_DOC
-)
-
-
-class FlaxElectraForTokenClassificationModule(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.electra = FlaxElectraModule(
- config=self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
- )
- classifier_dropout = (
- self.config.classifier_dropout
- if self.config.classifier_dropout is not None
- else self.config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Dense(self.config.num_labels, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- # Model
- outputs = self.electra(
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = outputs[0]
-
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- logits = self.classifier(hidden_states)
-
- if not return_dict:
- return (logits,) + outputs[1:]
-
- return FlaxTokenClassifierOutput(
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Electra model with a token classification head on top.
-
- Both the discriminator and generator may be loaded into this model.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class FlaxElectraForTokenClassification(FlaxElectraPreTrainedModel):
- module_class = FlaxElectraForTokenClassificationModule
-
-
-append_call_sample_docstring(
- FlaxElectraForTokenClassification,
- _CHECKPOINT_FOR_DOC,
- FlaxTokenClassifierOutput,
- _CONFIG_FOR_DOC,
-)
-
-
-def identity(x, **kwargs):
- return x
-
-
-class FlaxElectraSequenceSummary(nn.Module):
- r"""
- Compute a single vector summary of a sequence hidden states.
-
- Args:
- config ([`PretrainedConfig`]):
- The config used by the model. Relevant arguments in the config class of the model are (refer to the actual
- config class of your model for the default values it uses):
-
- - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.
- - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes
- (otherwise to `config.hidden_size`).
- - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,
- another string or `None` will add no activation.
- - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.
- - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.
- """
-
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.summary = identity
- if hasattr(self.config, "summary_use_proj") and self.config.summary_use_proj:
- if (
- hasattr(self.config, "summary_proj_to_labels")
- and self.config.summary_proj_to_labels
- and self.config.num_labels > 0
- ):
- num_classes = self.config.num_labels
- else:
- num_classes = self.config.hidden_size
- self.summary = nn.Dense(num_classes, dtype=self.dtype)
-
- activation_string = getattr(self.config, "summary_activation", None)
- self.activation = ACT2FN[activation_string] if activation_string else lambda x: x # noqa F407
-
- self.first_dropout = identity
- if hasattr(self.config, "summary_first_dropout") and self.config.summary_first_dropout > 0:
- self.first_dropout = nn.Dropout(self.config.summary_first_dropout)
-
- self.last_dropout = identity
- if hasattr(self.config, "summary_last_dropout") and self.config.summary_last_dropout > 0:
- self.last_dropout = nn.Dropout(self.config.summary_last_dropout)
-
- def __call__(self, hidden_states, cls_index=None, deterministic: bool = True):
- """
- Compute a single vector summary of a sequence hidden states.
-
- Args:
- hidden_states (`jnp.ndarray` of shape `[batch_size, seq_len, hidden_size]`):
- The hidden states of the last layer.
- cls_index (`jnp.ndarray` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*):
- Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token.
-
- Returns:
- `jnp.ndarray`: The summary of the sequence hidden states.
- """
- # NOTE: this doest "first" type summary always
- output = hidden_states[:, 0]
- output = self.first_dropout(output, deterministic=deterministic)
- output = self.summary(output)
- output = self.activation(output)
- output = self.last_dropout(output, deterministic=deterministic)
- return output
-
-
-class FlaxElectraForMultipleChoiceModule(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.electra = FlaxElectraModule(
- config=self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
- )
- self.sequence_summary = FlaxElectraSequenceSummary(config=self.config, dtype=self.dtype)
- self.classifier = nn.Dense(1, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- num_choices = input_ids.shape[1]
- input_ids = input_ids.reshape(-1, input_ids.shape[-1]) if input_ids is not None else None
- attention_mask = attention_mask.reshape(-1, attention_mask.shape[-1]) if attention_mask is not None else None
- token_type_ids = token_type_ids.reshape(-1, token_type_ids.shape[-1]) if token_type_ids is not None else None
- position_ids = position_ids.reshape(-1, position_ids.shape[-1]) if position_ids is not None else None
-
- # Model
- outputs = self.electra(
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = outputs[0]
- pooled_output = self.sequence_summary(hidden_states, deterministic=deterministic)
- logits = self.classifier(pooled_output)
-
- reshaped_logits = logits.reshape(-1, num_choices)
-
- if not return_dict:
- return (reshaped_logits,) + outputs[1:]
-
- return FlaxMultipleChoiceModelOutput(
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- ELECTRA Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class FlaxElectraForMultipleChoice(FlaxElectraPreTrainedModel):
- module_class = FlaxElectraForMultipleChoiceModule
-
-
-# adapt docstring slightly for FlaxElectraForMultipleChoice
-overwrite_call_docstring(
- FlaxElectraForMultipleChoice, ELECTRA_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
-)
-append_call_sample_docstring(
- FlaxElectraForMultipleChoice,
- _CHECKPOINT_FOR_DOC,
- FlaxMultipleChoiceModelOutput,
- _CONFIG_FOR_DOC,
-)
-
-
-class FlaxElectraForQuestionAnsweringModule(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.electra = FlaxElectraModule(
- config=self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
- )
- self.qa_outputs = nn.Dense(self.config.num_labels, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- # Model
- outputs = self.electra(
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = outputs[0]
- logits = self.qa_outputs(hidden_states)
- start_logits, end_logits = logits.split(self.config.num_labels, axis=-1)
- start_logits = start_logits.squeeze(-1)
- end_logits = end_logits.squeeze(-1)
-
- if not return_dict:
- return (start_logits, end_logits) + outputs[1:]
-
- return FlaxQuestionAnsweringModelOutput(
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- ELECTRA Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- ELECTRA_START_DOCSTRING,
-)
-class FlaxElectraForQuestionAnswering(FlaxElectraPreTrainedModel):
- module_class = FlaxElectraForQuestionAnsweringModule
-
-
-append_call_sample_docstring(
- FlaxElectraForQuestionAnswering,
- _CHECKPOINT_FOR_DOC,
- FlaxQuestionAnsweringModelOutput,
- _CONFIG_FOR_DOC,
-)
-
-
-class FlaxElectraClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.dense = nn.Dense(self.config.hidden_size, dtype=self.dtype)
- classifier_dropout = (
- self.config.classifier_dropout
- if self.config.classifier_dropout is not None
- else self.config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.out_proj = nn.Dense(self.config.num_labels, dtype=self.dtype)
-
- def __call__(self, hidden_states, deterministic: bool = True):
- x = hidden_states[:, 0, :] # take token (equiv. to [CLS])
- x = self.dropout(x, deterministic=deterministic)
- x = self.dense(x)
- x = ACT2FN["gelu"](x) # although BERT uses tanh here, it seems Electra authors used gelu
- x = self.dropout(x, deterministic=deterministic)
- x = self.out_proj(x)
- return x
-
-
-class FlaxElectraForSequenceClassificationModule(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.electra = FlaxElectraModule(
- config=self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
- )
- self.classifier = FlaxElectraClassificationHead(config=self.config, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- # Model
- outputs = self.electra(
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = outputs[0]
- logits = self.classifier(hidden_states, deterministic=deterministic)
-
- if not return_dict:
- return (logits,) + outputs[1:]
-
- return FlaxSequenceClassifierOutput(
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Electra Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class FlaxElectraForSequenceClassification(FlaxElectraPreTrainedModel):
- module_class = FlaxElectraForSequenceClassificationModule
-
-
-append_call_sample_docstring(
- FlaxElectraForSequenceClassification,
- _CHECKPOINT_FOR_DOC,
- FlaxSequenceClassifierOutput,
- _CONFIG_FOR_DOC,
-)
-
-
-class FlaxElectraForCausalLMModule(nn.Module):
- config: ElectraConfig
- dtype: jnp.dtype = jnp.float32
- gradient_checkpointing: bool = False
-
- def setup(self):
- self.electra = FlaxElectraModule(
- config=self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
- )
- self.generator_predictions = FlaxElectraGeneratorPredictions(config=self.config, dtype=self.dtype)
- if self.config.tie_word_embeddings:
- self.generator_lm_head = FlaxElectraTiedDense(self.config.vocab_size, dtype=self.dtype)
- else:
- self.generator_lm_head = nn.Dense(self.config.vocab_size, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask: Optional[jnp.ndarray] = None,
- token_type_ids: Optional[jnp.ndarray] = None,
- position_ids: Optional[jnp.ndarray] = None,
- head_mask: Optional[jnp.ndarray] = None,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- init_cache: bool = False,
- deterministic: bool = True,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- outputs = self.electra(
- input_ids,
- attention_mask,
- token_type_ids,
- position_ids,
- head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- init_cache=init_cache,
- deterministic=deterministic,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = outputs[0]
- prediction_scores = self.generator_predictions(hidden_states)
-
- if self.config.tie_word_embeddings:
- shared_embedding = self.electra.variables["params"]["embeddings"]["word_embeddings"]["embedding"]
- prediction_scores = self.generator_lm_head(prediction_scores, shared_embedding.T)
- else:
- prediction_scores = self.generator_lm_head(prediction_scores)
-
- if not return_dict:
- return (prediction_scores,) + outputs[1:]
-
- return FlaxCausalLMOutputWithCrossAttentions(
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- cross_attentions=outputs.cross_attentions,
- )
-
-
-@add_start_docstrings(
- """
- Electra Model with a language modeling head on top (a linear layer on top of the hidden-states output) e.g for
- autoregressive tasks.
- """,
- ELECTRA_START_DOCSTRING,
-)
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertForCausalLM with Bert->Electra
-class FlaxElectraForCausalLM(FlaxElectraPreTrainedModel):
- module_class = FlaxElectraForCausalLMModule
-
- def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jax.Array] = None):
- # initializing the cache
- batch_size, seq_length = input_ids.shape
-
- past_key_values = self.init_cache(batch_size, max_length)
- # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
- # But since the decoder uses a causal mask, those positions are masked anyway.
- # Thus, we can create a single static attention_mask here, which is more efficient for compilation
- extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
- if attention_mask is not None:
- position_ids = attention_mask.cumsum(axis=-1) - 1
- extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, attention_mask, (0, 0))
- else:
- position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))
-
- return {
- "past_key_values": past_key_values,
- "attention_mask": extended_attention_mask,
- "position_ids": position_ids,
- }
-
- def update_inputs_for_generation(self, model_outputs, model_kwargs):
- model_kwargs["past_key_values"] = model_outputs.past_key_values
- model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1
- return model_kwargs
-
-
-append_call_sample_docstring(
- FlaxElectraForCausalLM,
- _CHECKPOINT_FOR_DOC,
- FlaxCausalLMOutputWithCrossAttentions,
- _CONFIG_FOR_DOC,
-)
diff --git a/transformers/models/electra/modeling_tf_electra.py b/transformers/models/electra/modeling_tf_electra.py
deleted file mode 100644
index ba60cd8f5d575434f62312a573724cbdc34974d8..0000000000000000000000000000000000000000
--- a/transformers/models/electra/modeling_tf_electra.py
+++ /dev/null
@@ -1,1768 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Google AI Language Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF Electra model."""
-
-
-from __future__ import annotations
-
-import math
-import warnings
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutputWithPastAndCrossAttentions,
- TFMaskedLMOutput,
- TFMultipleChoiceModelOutput,
- TFQuestionAnsweringModelOutput,
- TFSequenceClassifierOutput,
- TFTokenClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFMaskedLanguageModelingLoss,
- TFModelInputType,
- TFMultipleChoiceLoss,
- TFPreTrainedModel,
- TFQuestionAnsweringLoss,
- TFSequenceClassificationLoss,
- TFSequenceSummary,
- TFTokenClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_electra import ElectraConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "google/electra-small-discriminator"
-_CONFIG_FOR_DOC = "ElectraConfig"
-
-
-from ..deprecated._archive_maps import TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfAttention with Bert->Electra
-class TFElectraSelfAttention(keras.layers.Layer):
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
-
- if config.hidden_size % config.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number "
- f"of attention heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.sqrt_att_head_size = math.sqrt(self.attention_head_size)
-
- self.query = keras.layers.Dense(
- units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"
- )
- self.key = keras.layers.Dense(
- units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key"
- )
- self.value = keras.layers.Dense(
- units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"
- )
- self.dropout = keras.layers.Dropout(rate=config.attention_probs_dropout_prob)
-
- self.is_decoder = config.is_decoder
- self.config = config
-
- def transpose_for_scores(self, tensor: tf.Tensor, batch_size: int) -> tf.Tensor:
- # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]
- tensor = tf.reshape(tensor=tensor, shape=(batch_size, -1, self.num_attention_heads, self.attention_head_size))
-
- # Transpose the tensor from [batch_size, seq_length, num_attention_heads, attention_head_size] to [batch_size, num_attention_heads, seq_length, attention_head_size]
- return tf.transpose(tensor, perm=[0, 2, 1, 3])
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- head_mask: tf.Tensor,
- encoder_hidden_states: tf.Tensor,
- encoder_attention_mask: tf.Tensor,
- past_key_value: Tuple[tf.Tensor],
- output_attentions: bool,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- batch_size = shape_list(hidden_states)[0]
- mixed_query_layer = self.query(inputs=hidden_states)
-
- # If this is instantiated as a cross-attention module, the keys
- # and values come from an encoder; the attention mask needs to be
- # such that the encoder's padding tokens are not attended to.
- is_cross_attention = encoder_hidden_states is not None
-
- if is_cross_attention and past_key_value is not None:
- # reuse k,v, cross_attentions
- key_layer = past_key_value[0]
- value_layer = past_key_value[1]
- attention_mask = encoder_attention_mask
- elif is_cross_attention:
- key_layer = self.transpose_for_scores(self.key(inputs=encoder_hidden_states), batch_size)
- value_layer = self.transpose_for_scores(self.value(inputs=encoder_hidden_states), batch_size)
- attention_mask = encoder_attention_mask
- elif past_key_value is not None:
- key_layer = self.transpose_for_scores(self.key(inputs=hidden_states), batch_size)
- value_layer = self.transpose_for_scores(self.value(inputs=hidden_states), batch_size)
- key_layer = tf.concat([past_key_value[0], key_layer], axis=2)
- value_layer = tf.concat([past_key_value[1], value_layer], axis=2)
- else:
- key_layer = self.transpose_for_scores(self.key(inputs=hidden_states), batch_size)
- value_layer = self.transpose_for_scores(self.value(inputs=hidden_states), batch_size)
-
- query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
-
- if self.is_decoder:
- # if cross_attention save Tuple(tf.Tensor, tf.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(tf.Tensor, tf.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_layer, value_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- # (batch size, num_heads, seq_len_q, seq_len_k)
- attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
- dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
- attention_scores = tf.divide(attention_scores, dk)
-
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in TFElectraModel call() function)
- attention_scores = tf.add(attention_scores, attention_mask)
-
- # Normalize the attention scores to probabilities.
- attention_probs = stable_softmax(logits=attention_scores, axis=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(inputs=attention_probs, training=training)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = tf.multiply(attention_probs, head_mask)
-
- attention_output = tf.matmul(attention_probs, value_layer)
- attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])
-
- # (batch_size, seq_len_q, all_head_size)
- attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
- outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)
-
- if self.is_decoder:
- outputs = outputs + (past_key_value,)
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "query", None) is not None:
- with tf.name_scope(self.query.name):
- self.query.build([None, None, self.config.hidden_size])
- if getattr(self, "key", None) is not None:
- with tf.name_scope(self.key.name):
- self.key.build([None, None, self.config.hidden_size])
- if getattr(self, "value", None) is not None:
- with tf.name_scope(self.value.name):
- self.value.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfOutput with Bert->Electra
-class TFElectraSelfOutput(keras.layers.Layer):
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.dropout(inputs=hidden_states, training=training)
- hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertAttention with Bert->Electra
-class TFElectraAttention(keras.layers.Layer):
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.self_attention = TFElectraSelfAttention(config, name="self")
- self.dense_output = TFElectraSelfOutput(config, name="output")
-
- def prune_heads(self, heads):
- raise NotImplementedError
-
- def call(
- self,
- input_tensor: tf.Tensor,
- attention_mask: tf.Tensor,
- head_mask: tf.Tensor,
- encoder_hidden_states: tf.Tensor,
- encoder_attention_mask: tf.Tensor,
- past_key_value: Tuple[tf.Tensor],
- output_attentions: bool,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- self_outputs = self.self_attention(
- hidden_states=input_tensor,
- attention_mask=attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- training=training,
- )
- attention_output = self.dense_output(
- hidden_states=self_outputs[0], input_tensor=input_tensor, training=training
- )
- # add attentions (possibly with past_key_value) if we output them
- outputs = (attention_output,) + self_outputs[1:]
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "self_attention", None) is not None:
- with tf.name_scope(self.self_attention.name):
- self.self_attention.build(None)
- if getattr(self, "dense_output", None) is not None:
- with tf.name_scope(self.dense_output.name):
- self.dense_output.build(None)
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertIntermediate with Bert->Electra
-class TFElectraIntermediate(keras.layers.Layer):
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
-
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = get_tf_activation(config.hidden_act)
- else:
- self.intermediate_act_fn = config.hidden_act
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertOutput with Bert->Electra
-class TFElectraOutput(keras.layers.Layer):
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = self.dropout(inputs=hidden_states, training=training)
- hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.intermediate_size])
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertLayer with Bert->Electra
-class TFElectraLayer(keras.layers.Layer):
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.attention = TFElectraAttention(config, name="attention")
- self.is_decoder = config.is_decoder
- self.add_cross_attention = config.add_cross_attention
- if self.add_cross_attention:
- if not self.is_decoder:
- raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
- self.crossattention = TFElectraAttention(config, name="crossattention")
- self.intermediate = TFElectraIntermediate(config, name="intermediate")
- self.bert_output = TFElectraOutput(config, name="output")
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- head_mask: tf.Tensor,
- encoder_hidden_states: tf.Tensor | None,
- encoder_attention_mask: tf.Tensor | None,
- past_key_value: Tuple[tf.Tensor] | None,
- output_attentions: bool,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
- self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
- self_attention_outputs = self.attention(
- input_tensor=hidden_states,
- attention_mask=attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- past_key_value=self_attn_past_key_value,
- output_attentions=output_attentions,
- training=training,
- )
- attention_output = self_attention_outputs[0]
-
- # if decoder, the last output is tuple of self-attn cache
- if self.is_decoder:
- outputs = self_attention_outputs[1:-1]
- present_key_value = self_attention_outputs[-1]
- else:
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- cross_attn_present_key_value = None
- if self.is_decoder and encoder_hidden_states is not None:
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
- " by setting `config.add_cross_attention=True`"
- )
-
- # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
- cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
- cross_attention_outputs = self.crossattention(
- input_tensor=attention_output,
- attention_mask=attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- past_key_value=cross_attn_past_key_value,
- output_attentions=output_attentions,
- training=training,
- )
- attention_output = cross_attention_outputs[0]
- outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
-
- # add cross-attn cache to positions 3,4 of present_key_value tuple
- cross_attn_present_key_value = cross_attention_outputs[-1]
- present_key_value = present_key_value + cross_attn_present_key_value
-
- intermediate_output = self.intermediate(hidden_states=attention_output)
- layer_output = self.bert_output(
- hidden_states=intermediate_output, input_tensor=attention_output, training=training
- )
- outputs = (layer_output,) + outputs # add attentions if we output them
-
- # if decoder, return the attn key/values as the last output
- if self.is_decoder:
- outputs = outputs + (present_key_value,)
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "intermediate", None) is not None:
- with tf.name_scope(self.intermediate.name):
- self.intermediate.build(None)
- if getattr(self, "bert_output", None) is not None:
- with tf.name_scope(self.bert_output.name):
- self.bert_output.build(None)
- if getattr(self, "crossattention", None) is not None:
- with tf.name_scope(self.crossattention.name):
- self.crossattention.build(None)
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertEncoder with Bert->Electra
-class TFElectraEncoder(keras.layers.Layer):
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
- self.config = config
- self.layer = [TFElectraLayer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor,
- head_mask: tf.Tensor,
- encoder_hidden_states: tf.Tensor | None,
- encoder_attention_mask: tf.Tensor | None,
- past_key_values: Tuple[Tuple[tf.Tensor]] | None,
- use_cache: Optional[bool],
- output_attentions: bool,
- output_hidden_states: bool,
- return_dict: bool,
- training: bool = False,
- ) -> Union[TFBaseModelOutputWithPastAndCrossAttentions, Tuple[tf.Tensor]]:
- all_hidden_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
-
- next_decoder_cache = () if use_cache else None
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- past_key_value = past_key_values[i] if past_key_values is not None else None
-
- layer_outputs = layer_module(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- head_mask=head_mask[i],
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- training=training,
- )
- hidden_states = layer_outputs[0]
-
- if use_cache:
- next_decoder_cache += (layer_outputs[-1],)
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
- if self.config.add_cross_attention and encoder_hidden_states is not None:
- all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
-
- # Add last layer
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v for v in [hidden_states, all_hidden_states, all_attentions, all_cross_attentions] if v is not None
- )
-
- return TFBaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=next_decoder_cache,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- cross_attentions=all_cross_attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "layer", None) is not None:
- for layer in self.layer:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertPooler with Bert->Electra
-class TFElectraPooler(keras.layers.Layer):
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size,
- kernel_initializer=get_initializer(config.initializer_range),
- activation="tanh",
- name="dense",
- )
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(inputs=first_token_tensor)
-
- return pooled_output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-# Copied from transformers.models.albert.modeling_tf_albert.TFAlbertEmbeddings with Albert->Electra
-class TFElectraEmbeddings(keras.layers.Layer):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config: ElectraConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embedding_size = config.embedding_size
- self.max_position_embeddings = config.max_position_embeddings
- self.initializer_range = config.initializer_range
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)
-
- def build(self, input_shape=None):
- with tf.name_scope("word_embeddings"):
- self.weight = self.add_weight(
- name="weight",
- shape=[self.config.vocab_size, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
-
- with tf.name_scope("token_type_embeddings"):
- self.token_type_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.config.type_vocab_size, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
-
- with tf.name_scope("position_embeddings"):
- self.position_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.max_position_embeddings, self.embedding_size],
- initializer=get_initializer(self.initializer_range),
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.embedding_size])
-
- # Copied from transformers.models.bert.modeling_tf_bert.TFBertEmbeddings.call
- def call(
- self,
- input_ids: tf.Tensor = None,
- position_ids: tf.Tensor = None,
- token_type_ids: tf.Tensor = None,
- inputs_embeds: tf.Tensor = None,
- past_key_values_length=0,
- training: bool = False,
- ) -> tf.Tensor:
- """
- Applies embedding based on inputs tensor.
-
- Returns:
- final_embeddings (`tf.Tensor`): output embedding tensor.
- """
- if input_ids is None and inputs_embeds is None:
- raise ValueError("Need to provide either `input_ids` or `input_embeds`.")
-
- if input_ids is not None:
- check_embeddings_within_bounds(input_ids, self.config.vocab_size)
- inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
-
- input_shape = shape_list(inputs_embeds)[:-1]
-
- if token_type_ids is None:
- token_type_ids = tf.fill(dims=input_shape, value=0)
-
- if position_ids is None:
- position_ids = tf.expand_dims(
- tf.range(start=past_key_values_length, limit=input_shape[1] + past_key_values_length), axis=0
- )
-
- position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
- token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)
- final_embeddings = inputs_embeds + position_embeds + token_type_embeds
- final_embeddings = self.LayerNorm(inputs=final_embeddings)
- final_embeddings = self.dropout(inputs=final_embeddings, training=training)
-
- return final_embeddings
-
-
-class TFElectraDiscriminatorPredictions(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(config.hidden_size, name="dense")
- self.dense_prediction = keras.layers.Dense(1, name="dense_prediction")
- self.config = config
-
- def call(self, discriminator_hidden_states, training=False):
- hidden_states = self.dense(discriminator_hidden_states)
- hidden_states = get_tf_activation(self.config.hidden_act)(hidden_states)
- logits = tf.squeeze(self.dense_prediction(hidden_states), -1)
-
- return logits
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "dense_prediction", None) is not None:
- with tf.name_scope(self.dense_prediction.name):
- self.dense_prediction.build([None, None, self.config.hidden_size])
-
-
-class TFElectraGeneratorPredictions(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.dense = keras.layers.Dense(config.embedding_size, name="dense")
- self.config = config
-
- def call(self, generator_hidden_states, training=False):
- hidden_states = self.dense(generator_hidden_states)
- hidden_states = get_tf_activation("gelu")(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.embedding_size])
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFElectraPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ElectraConfig
- base_model_prefix = "electra"
- # When the model is loaded from a PT model
- _keys_to_ignore_on_load_unexpected = [r"generator_lm_head.weight"]
- _keys_to_ignore_on_load_missing = [r"dropout"]
-
-
-@keras_serializable
-class TFElectraMainLayer(keras.layers.Layer):
- config_class = ElectraConfig
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.is_decoder = config.is_decoder
-
- self.embeddings = TFElectraEmbeddings(config, name="embeddings")
-
- if config.embedding_size != config.hidden_size:
- self.embeddings_project = keras.layers.Dense(config.hidden_size, name="embeddings_project")
-
- self.encoder = TFElectraEncoder(config, name="encoder")
-
- def get_input_embeddings(self):
- return self.embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.weight = value
- self.embeddings.vocab_size = shape_list(value)[0]
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- raise NotImplementedError
-
- def get_extended_attention_mask(self, attention_mask, input_shape, dtype, past_key_values_length=0):
- batch_size, seq_length = input_shape
-
- if attention_mask is None:
- attention_mask = tf.fill(dims=(batch_size, seq_length + past_key_values_length), value=1)
-
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask_shape = shape_list(attention_mask)
-
- mask_seq_length = seq_length + past_key_values_length
- # Copied from `modeling_tf_t5.py`
- # Provided a padding mask of dimensions [batch_size, mask_seq_length]
- # - if the model is a decoder, apply a causal mask in addition to the padding mask
- # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]
- if self.is_decoder:
- seq_ids = tf.range(mask_seq_length)
- causal_mask = tf.less_equal(
- tf.tile(seq_ids[None, None, :], (batch_size, mask_seq_length, 1)),
- seq_ids[None, :, None],
- )
- causal_mask = tf.cast(causal_mask, dtype=attention_mask.dtype)
- extended_attention_mask = causal_mask * attention_mask[:, None, :]
- attention_mask_shape = shape_list(extended_attention_mask)
- extended_attention_mask = tf.reshape(
- extended_attention_mask, (attention_mask_shape[0], 1, attention_mask_shape[1], attention_mask_shape[2])
- )
- if past_key_values_length > 0:
- extended_attention_mask = extended_attention_mask[:, :, -seq_length:, :]
- else:
- extended_attention_mask = tf.reshape(
- attention_mask, (attention_mask_shape[0], 1, 1, attention_mask_shape[1])
- )
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and -10000.0 for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- extended_attention_mask = tf.cast(extended_attention_mask, dtype=dtype)
- one_cst = tf.constant(1.0, dtype=dtype)
- ten_thousand_cst = tf.constant(-10000.0, dtype=dtype)
- extended_attention_mask = tf.multiply(tf.subtract(one_cst, extended_attention_mask), ten_thousand_cst)
-
- return extended_attention_mask
-
- def get_head_mask(self, head_mask):
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.config.num_hidden_layers
-
- return head_mask
-
- @unpack_inputs
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- encoder_hidden_states: np.ndarray | tf.Tensor | None = None,
- encoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithPastAndCrossAttentions, Tuple[tf.Tensor]]:
- if not self.config.is_decoder:
- use_cache = False
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
-
- if past_key_values is None:
- past_key_values_length = 0
- past_key_values = [None] * len(self.encoder.layer)
- else:
- past_key_values_length = shape_list(past_key_values[0][0])[-2]
-
- if attention_mask is None:
- attention_mask = tf.fill(dims=(batch_size, seq_length + past_key_values_length), value=1)
-
- if token_type_ids is None:
- token_type_ids = tf.fill(dims=input_shape, value=0)
-
- hidden_states = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- past_key_values_length=past_key_values_length,
- training=training,
- )
- extended_attention_mask = self.get_extended_attention_mask(
- attention_mask, input_shape, hidden_states.dtype, past_key_values_length
- )
-
- # Copied from `modeling_tf_t5.py` with -1e9 -> -10000
- if self.is_decoder and encoder_attention_mask is not None:
- # If a 2D ou 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- encoder_attention_mask = tf.cast(encoder_attention_mask, dtype=extended_attention_mask.dtype)
- num_dims_encoder_attention_mask = len(shape_list(encoder_attention_mask))
- if num_dims_encoder_attention_mask == 3:
- encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
- if num_dims_encoder_attention_mask == 2:
- encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
-
- # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
- # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow/transformer/transformer_layers.py#L270
- # encoder_extended_attention_mask = tf.math.equal(encoder_extended_attention_mask,
- # tf.transpose(encoder_extended_attention_mask, perm=(-1, -2)))
-
- encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -10000.0
- else:
- encoder_extended_attention_mask = None
-
- head_mask = self.get_head_mask(head_mask)
-
- if hasattr(self, "embeddings_project"):
- hidden_states = self.embeddings_project(hidden_states, training=training)
-
- hidden_states = self.encoder(
- hidden_states=hidden_states,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_extended_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "embeddings_project", None) is not None:
- with tf.name_scope(self.embeddings_project.name):
- self.embeddings_project.build([None, None, self.config.embedding_size])
-
-
-@dataclass
-class TFElectraForPreTrainingOutput(ModelOutput):
- """
- Output type of [`TFElectraForPreTraining`].
-
- Args:
- loss (*optional*, returned when `labels` is provided, `tf.Tensor` of shape `(1,)`):
- Total loss of the ELECTRA objective.
- logits (`tf.Tensor` of shape `(batch_size, sequence_length)`):
- Prediction scores of the head (scores for each token before SoftMax).
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- logits: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-ELECTRA_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`ElectraConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ELECTRA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`Numpy array` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare Electra Model transformer outputting raw hidden-states without any specific head on top. Identical to "
- "the BERT model except that it uses an additional linear layer between the embedding layer and the encoder if the "
- "hidden size and embedding size are different. "
- ""
- "Both the generator and discriminator checkpoints may be loaded into this model.",
- ELECTRA_START_DOCSTRING,
-)
-class TFElectraModel(TFElectraPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.electra = TFElectraMainLayer(config, name="electra")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutputWithPastAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- encoder_hidden_states: np.ndarray | tf.Tensor | None = None,
- encoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithPastAndCrossAttentions, Tuple[tf.Tensor]]:
- r"""
- encoder_hidden_states (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- past_key_values (`Tuple[Tuple[tf.Tensor]]` of length `config.n_layers`)
- contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*, defaults to `True`):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`). Set to `False` during training, `True` during generation
- """
- outputs = self.electra(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "electra", None) is not None:
- with tf.name_scope(self.electra.name):
- self.electra.build(None)
-
-
-@add_start_docstrings(
- """
- Electra model with a binary classification head on top as used during pretraining for identifying generated tokens.
-
- Even though both the discriminator and generator may be loaded into this model, the discriminator is the only model
- of the two to have the correct classification head to be used for this model.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class TFElectraForPreTraining(TFElectraPreTrainedModel):
- def __init__(self, config, **kwargs):
- super().__init__(config, **kwargs)
-
- self.electra = TFElectraMainLayer(config, name="electra")
- self.discriminator_predictions = TFElectraDiscriminatorPredictions(config, name="discriminator_predictions")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=TFElectraForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFElectraForPreTrainingOutput, Tuple[tf.Tensor]]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> import tensorflow as tf
- >>> from transformers import AutoTokenizer, TFElectraForPreTraining
-
- >>> tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")
- >>> model = TFElectraForPreTraining.from_pretrained("google/electra-small-discriminator")
- >>> input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :] # Batch size 1
- >>> outputs = model(input_ids)
- >>> scores = outputs[0]
- ```"""
- discriminator_hidden_states = self.electra(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- discriminator_sequence_output = discriminator_hidden_states[0]
- logits = self.discriminator_predictions(discriminator_sequence_output)
-
- if not return_dict:
- return (logits,) + discriminator_hidden_states[1:]
-
- return TFElectraForPreTrainingOutput(
- logits=logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "electra", None) is not None:
- with tf.name_scope(self.electra.name):
- self.electra.build(None)
- if getattr(self, "discriminator_predictions", None) is not None:
- with tf.name_scope(self.discriminator_predictions.name):
- self.discriminator_predictions.build(None)
-
-
-class TFElectraMaskedLMHead(keras.layers.Layer):
- def __init__(self, config, input_embeddings, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.embedding_size = config.embedding_size
- self.input_embeddings = input_embeddings
-
- def build(self, input_shape):
- self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")
-
- super().build(input_shape)
-
- def get_output_embeddings(self):
- return self.input_embeddings
-
- def set_output_embeddings(self, value):
- self.input_embeddings.weight = value
- self.input_embeddings.vocab_size = shape_list(value)[0]
-
- def get_bias(self):
- return {"bias": self.bias}
-
- def set_bias(self, value):
- self.bias = value["bias"]
- self.config.vocab_size = shape_list(value["bias"])[0]
-
- def call(self, hidden_states):
- seq_length = shape_list(tensor=hidden_states)[1]
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.embedding_size])
- hidden_states = tf.matmul(a=hidden_states, b=self.input_embeddings.weight, transpose_b=True)
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])
- hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)
-
- return hidden_states
-
-
-@add_start_docstrings(
- """
- Electra model with a language modeling head on top.
-
- Even though both the discriminator and generator may be loaded into this model, the generator is the only model of
- the two to have been trained for the masked language modeling task.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class TFElectraForMaskedLM(TFElectraPreTrainedModel, TFMaskedLanguageModelingLoss):
- def __init__(self, config, **kwargs):
- super().__init__(config, **kwargs)
-
- self.config = config
- self.electra = TFElectraMainLayer(config, name="electra")
- self.generator_predictions = TFElectraGeneratorPredictions(config, name="generator_predictions")
-
- if isinstance(config.hidden_act, str):
- self.activation = get_tf_activation(config.hidden_act)
- else:
- self.activation = config.hidden_act
-
- self.generator_lm_head = TFElectraMaskedLMHead(config, self.electra.embeddings, name="generator_lm_head")
-
- def get_lm_head(self):
- return self.generator_lm_head
-
- def get_prefix_bias_name(self):
- warnings.warn("The method get_prefix_bias_name is deprecated. Please use `get_bias` instead.", FutureWarning)
- return self.name + "/" + self.generator_lm_head.name
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="google/electra-small-generator",
- output_type=TFMaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="[MASK]",
- expected_output="'paris'",
- expected_loss=1.22,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFMaskedLMOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- generator_hidden_states = self.electra(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- generator_sequence_output = generator_hidden_states[0]
- prediction_scores = self.generator_predictions(generator_sequence_output, training=training)
- prediction_scores = self.generator_lm_head(prediction_scores, training=training)
- loss = None if labels is None else self.hf_compute_loss(labels, prediction_scores)
-
- if not return_dict:
- output = (prediction_scores,) + generator_hidden_states[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFMaskedLMOutput(
- loss=loss,
- logits=prediction_scores,
- hidden_states=generator_hidden_states.hidden_states,
- attentions=generator_hidden_states.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "electra", None) is not None:
- with tf.name_scope(self.electra.name):
- self.electra.build(None)
- if getattr(self, "generator_predictions", None) is not None:
- with tf.name_scope(self.generator_predictions.name):
- self.generator_predictions.build(None)
- if getattr(self, "generator_lm_head", None) is not None:
- with tf.name_scope(self.generator_lm_head.name):
- self.generator_lm_head.build(None)
-
-
-class TFElectraClassificationHead(keras.layers.Layer):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- classifier_dropout = (
- config.classifhidden_dropout_probier_dropout
- if config.classifier_dropout is not None
- else config.hidden_dropout_prob
- )
- self.dropout = keras.layers.Dropout(classifier_dropout)
- self.out_proj = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="out_proj"
- )
- self.config = config
-
- def call(self, inputs, **kwargs):
- x = inputs[:, 0, :] # take token (equiv. to [CLS])
- x = self.dropout(x)
- x = self.dense(x)
- x = get_tf_activation("gelu")(x) # although BERT uses tanh here, it seems Electra authors used gelu here
- x = self.dropout(x)
- x = self.out_proj(x)
-
- return x
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "out_proj", None) is not None:
- with tf.name_scope(self.out_proj.name):
- self.out_proj.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- ELECTRA Model transformer with a sequence classification/regression head on top (a linear layer on top of the
- pooled output) e.g. for GLUE tasks.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class TFElectraForSequenceClassification(TFElectraPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
- self.electra = TFElectraMainLayer(config, name="electra")
- self.classifier = TFElectraClassificationHead(config, name="classifier")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="bhadresh-savani/electra-base-emotion",
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output="'joy'",
- expected_loss=0.06,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- outputs = self.electra(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- logits = self.classifier(outputs[0])
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "electra", None) is not None:
- with tf.name_scope(self.electra.name):
- self.electra.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build(None)
-
-
-@add_start_docstrings(
- """
- ELECTRA Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class TFElectraForMultipleChoice(TFElectraPreTrainedModel, TFMultipleChoiceLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.electra = TFElectraMainLayer(config, name="electra")
- self.sequence_summary = TFSequenceSummary(
- config, initializer_range=config.initializer_range, name="sequence_summary"
- )
- self.classifier = keras.layers.Dense(
- 1, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFMultipleChoiceModelOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]`
- where `num_choices` is the size of the second dimension of the input tensors. (See `input_ids` above)
- """
-
- if input_ids is not None:
- num_choices = shape_list(input_ids)[1]
- seq_length = shape_list(input_ids)[2]
- else:
- num_choices = shape_list(inputs_embeds)[1]
- seq_length = shape_list(inputs_embeds)[2]
-
- flat_input_ids = tf.reshape(input_ids, (-1, seq_length)) if input_ids is not None else None
- flat_attention_mask = tf.reshape(attention_mask, (-1, seq_length)) if attention_mask is not None else None
- flat_token_type_ids = tf.reshape(token_type_ids, (-1, seq_length)) if token_type_ids is not None else None
- flat_position_ids = tf.reshape(position_ids, (-1, seq_length)) if position_ids is not None else None
- flat_inputs_embeds = (
- tf.reshape(inputs_embeds, (-1, seq_length, shape_list(inputs_embeds)[3]))
- if inputs_embeds is not None
- else None
- )
- outputs = self.electra(
- input_ids=flat_input_ids,
- attention_mask=flat_attention_mask,
- token_type_ids=flat_token_type_ids,
- position_ids=flat_position_ids,
- head_mask=head_mask,
- inputs_embeds=flat_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- logits = self.sequence_summary(outputs[0])
- logits = self.classifier(logits)
- reshaped_logits = tf.reshape(logits, (-1, num_choices))
- loss = None if labels is None else self.hf_compute_loss(labels, reshaped_logits)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFMultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "electra", None) is not None:
- with tf.name_scope(self.electra.name):
- self.electra.build(None)
- if getattr(self, "sequence_summary", None) is not None:
- with tf.name_scope(self.sequence_summary.name):
- self.sequence_summary.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- Electra model with a token classification head on top.
-
- Both the discriminator and generator may be loaded into this model.
- """,
- ELECTRA_START_DOCSTRING,
-)
-class TFElectraForTokenClassification(TFElectraPreTrainedModel, TFTokenClassificationLoss):
- def __init__(self, config, **kwargs):
- super().__init__(config, **kwargs)
-
- self.electra = TFElectraMainLayer(config, name="electra")
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = keras.layers.Dropout(classifier_dropout)
- self.classifier = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="bhadresh-savani/electra-base-discriminator-finetuned-conll03-english",
- output_type=TFTokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output="['B-LOC', 'B-ORG', 'O', 'O', 'O', 'O', 'O', 'B-LOC', 'O', 'B-LOC', 'I-LOC']",
- expected_loss=0.11,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFTokenClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- discriminator_hidden_states = self.electra(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- discriminator_sequence_output = discriminator_hidden_states[0]
- discriminator_sequence_output = self.dropout(discriminator_sequence_output)
- logits = self.classifier(discriminator_sequence_output)
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + discriminator_hidden_states[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFTokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "electra", None) is not None:
- with tf.name_scope(self.electra.name):
- self.electra.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- Electra Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- ELECTRA_START_DOCSTRING,
-)
-class TFElectraForQuestionAnswering(TFElectraPreTrainedModel, TFQuestionAnsweringLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.num_labels = config.num_labels
- self.electra = TFElectraMainLayer(config, name="electra")
- self.qa_outputs = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="bhadresh-savani/electra-base-squad2",
- output_type=TFQuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- qa_target_start_index=11,
- qa_target_end_index=12,
- expected_output="'a nice puppet'",
- expected_loss=2.64,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- start_positions: np.ndarray | tf.Tensor | None = None,
- end_positions: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFQuestionAnsweringModelOutput, Tuple[tf.Tensor]]:
- r"""
- start_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- discriminator_hidden_states = self.electra(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- discriminator_sequence_output = discriminator_hidden_states[0]
- logits = self.qa_outputs(discriminator_sequence_output)
- start_logits, end_logits = tf.split(logits, 2, axis=-1)
- start_logits = tf.squeeze(start_logits, axis=-1)
- end_logits = tf.squeeze(end_logits, axis=-1)
- loss = None
-
- if start_positions is not None and end_positions is not None:
- labels = {"start_position": start_positions}
- labels["end_position"] = end_positions
- loss = self.hf_compute_loss(labels, (start_logits, end_logits))
-
- if not return_dict:
- output = (
- start_logits,
- end_logits,
- ) + discriminator_hidden_states[1:]
-
- return ((loss,) + output) if loss is not None else output
-
- return TFQuestionAnsweringModelOutput(
- loss=loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "electra", None) is not None:
- with tf.name_scope(self.electra.name):
- self.electra.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.config.hidden_size])
diff --git a/transformers/models/electra/tokenization_electra.py b/transformers/models/electra/tokenization_electra.py
deleted file mode 100644
index ceb3e7560215c2fb79bea37b515a5808c924fb7a..0000000000000000000000000000000000000000
--- a/transformers/models/electra/tokenization_electra.py
+++ /dev/null
@@ -1,503 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The Google AI Team, Stanford University and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import collections
-import os
-import unicodedata
-from typing import List, Optional, Tuple
-
-from ...tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
-
-
-# Copied from transformers.models.bert.tokenization_bert.load_vocab
-def load_vocab(vocab_file):
- """Loads a vocabulary file into a dictionary."""
- vocab = collections.OrderedDict()
- with open(vocab_file, "r", encoding="utf-8") as reader:
- tokens = reader.readlines()
- for index, token in enumerate(tokens):
- token = token.rstrip("\n")
- vocab[token] = index
- return vocab
-
-
-# Copied from transformers.models.bert.tokenization_bert.whitespace_tokenize
-def whitespace_tokenize(text):
- """Runs basic whitespace cleaning and splitting on a piece of text."""
- text = text.strip()
- if not text:
- return []
- tokens = text.split()
- return tokens
-
-
-# Copied from transformers.models.bert.tokenization_bert.BertTokenizer with Bert->Electra,BERT->Electra
-class ElectraTokenizer(PreTrainedTokenizer):
- r"""
- Construct a Electra tokenizer. Based on WordPiece.
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- do_basic_tokenize (`bool`, *optional*, defaults to `True`):
- Whether or not to do basic tokenization before WordPiece.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original Electra).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
-
- def __init__(
- self,
- vocab_file,
- do_lower_case=True,
- do_basic_tokenize=True,
- never_split=None,
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- if not os.path.isfile(vocab_file):
- raise ValueError(
- f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
- " model use `tokenizer = ElectraTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- self.vocab = load_vocab(vocab_file)
- self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
- self.do_basic_tokenize = do_basic_tokenize
- if do_basic_tokenize:
- self.basic_tokenizer = BasicTokenizer(
- do_lower_case=do_lower_case,
- never_split=never_split,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- )
-
- self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token))
-
- super().__init__(
- do_lower_case=do_lower_case,
- do_basic_tokenize=do_basic_tokenize,
- never_split=never_split,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- @property
- def do_lower_case(self):
- return self.basic_tokenizer.do_lower_case
-
- @property
- def vocab_size(self):
- return len(self.vocab)
-
- def get_vocab(self):
- return dict(self.vocab, **self.added_tokens_encoder)
-
- def _tokenize(self, text, split_special_tokens=False):
- split_tokens = []
- if self.do_basic_tokenize:
- for token in self.basic_tokenizer.tokenize(
- text, never_split=self.all_special_tokens if not split_special_tokens else None
- ):
- # If the token is part of the never_split set
- if token in self.basic_tokenizer.never_split:
- split_tokens.append(token)
- else:
- split_tokens += self.wordpiece_tokenizer.tokenize(token)
- else:
- split_tokens = self.wordpiece_tokenizer.tokenize(text)
- return split_tokens
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.vocab.get(token, self.vocab.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.ids_to_tokens.get(index, self.unk_token)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- out_string = " ".join(tokens).replace(" ##", "").strip()
- return out_string
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A Electra sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A Electra sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- index = 0
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- else:
- vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(token + "\n")
- index += 1
- return (vocab_file,)
-
-
-# Copied from transformers.models.bert.tokenization_bert.BasicTokenizer
-class BasicTokenizer(object):
- """
- Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
-
- Args:
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- do_split_on_punc (`bool`, *optional*, defaults to `True`):
- In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
- the full context of the words, such as contractions.
- """
-
- def __init__(
- self,
- do_lower_case=True,
- never_split=None,
- tokenize_chinese_chars=True,
- strip_accents=None,
- do_split_on_punc=True,
- ):
- if never_split is None:
- never_split = []
- self.do_lower_case = do_lower_case
- self.never_split = set(never_split)
- self.tokenize_chinese_chars = tokenize_chinese_chars
- self.strip_accents = strip_accents
- self.do_split_on_punc = do_split_on_punc
-
- def tokenize(self, text, never_split=None):
- """
- Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
-
- Args:
- never_split (`List[str]`, *optional*)
- Kept for backward compatibility purposes. Now implemented directly at the base class level (see
- [`PreTrainedTokenizer.tokenize`]) List of token not to split.
- """
- # union() returns a new set by concatenating the two sets.
- never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
- text = self._clean_text(text)
-
- # This was added on November 1st, 2018 for the multilingual and Chinese
- # models. This is also applied to the English models now, but it doesn't
- # matter since the English models were not trained on any Chinese data
- # and generally don't have any Chinese data in them (there are Chinese
- # characters in the vocabulary because Wikipedia does have some Chinese
- # words in the English Wikipedia.).
- if self.tokenize_chinese_chars:
- text = self._tokenize_chinese_chars(text)
- # prevents treating the same character with different unicode codepoints as different characters
- unicode_normalized_text = unicodedata.normalize("NFC", text)
- orig_tokens = whitespace_tokenize(unicode_normalized_text)
- split_tokens = []
- for token in orig_tokens:
- if token not in never_split:
- if self.do_lower_case:
- token = token.lower()
- if self.strip_accents is not False:
- token = self._run_strip_accents(token)
- elif self.strip_accents:
- token = self._run_strip_accents(token)
- split_tokens.extend(self._run_split_on_punc(token, never_split))
-
- output_tokens = whitespace_tokenize(" ".join(split_tokens))
- return output_tokens
-
- def _run_strip_accents(self, text):
- """Strips accents from a piece of text."""
- text = unicodedata.normalize("NFD", text)
- output = []
- for char in text:
- cat = unicodedata.category(char)
- if cat == "Mn":
- continue
- output.append(char)
- return "".join(output)
-
- def _run_split_on_punc(self, text, never_split=None):
- """Splits punctuation on a piece of text."""
- if not self.do_split_on_punc or (never_split is not None and text in never_split):
- return [text]
- chars = list(text)
- i = 0
- start_new_word = True
- output = []
- while i < len(chars):
- char = chars[i]
- if _is_punctuation(char):
- output.append([char])
- start_new_word = True
- else:
- if start_new_word:
- output.append([])
- start_new_word = False
- output[-1].append(char)
- i += 1
-
- return ["".join(x) for x in output]
-
- def _tokenize_chinese_chars(self, text):
- """Adds whitespace around any CJK character."""
- output = []
- for char in text:
- cp = ord(char)
- if self._is_chinese_char(cp):
- output.append(" ")
- output.append(char)
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
- def _is_chinese_char(self, cp):
- """Checks whether CP is the codepoint of a CJK character."""
- # This defines a "chinese character" as anything in the CJK Unicode block:
- # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
- #
- # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
- # despite its name. The modern Korean Hangul alphabet is a different block,
- # as is Japanese Hiragana and Katakana. Those alphabets are used to write
- # space-separated words, so they are not treated specially and handled
- # like the all of the other languages.
- if (
- (cp >= 0x4E00 and cp <= 0x9FFF)
- or (cp >= 0x3400 and cp <= 0x4DBF) #
- or (cp >= 0x20000 and cp <= 0x2A6DF) #
- or (cp >= 0x2A700 and cp <= 0x2B73F) #
- or (cp >= 0x2B740 and cp <= 0x2B81F) #
- or (cp >= 0x2B820 and cp <= 0x2CEAF) #
- or (cp >= 0xF900 and cp <= 0xFAFF)
- or (cp >= 0x2F800 and cp <= 0x2FA1F) #
- ): #
- return True
-
- return False
-
- def _clean_text(self, text):
- """Performs invalid character removal and whitespace cleanup on text."""
- output = []
- for char in text:
- cp = ord(char)
- if cp == 0 or cp == 0xFFFD or _is_control(char):
- continue
- if _is_whitespace(char):
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
-
-# Copied from transformers.models.bert.tokenization_bert.WordpieceTokenizer
-class WordpieceTokenizer(object):
- """Runs WordPiece tokenization."""
-
- def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
- self.vocab = vocab
- self.unk_token = unk_token
- self.max_input_chars_per_word = max_input_chars_per_word
-
- def tokenize(self, text):
- """
- Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
- tokenization using the given vocabulary.
-
- For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
-
- Args:
- text: A single token or whitespace separated tokens. This should have
- already been passed through *BasicTokenizer*.
-
- Returns:
- A list of wordpiece tokens.
- """
-
- output_tokens = []
- for token in whitespace_tokenize(text):
- chars = list(token)
- if len(chars) > self.max_input_chars_per_word:
- output_tokens.append(self.unk_token)
- continue
-
- is_bad = False
- start = 0
- sub_tokens = []
- while start < len(chars):
- end = len(chars)
- cur_substr = None
- while start < end:
- substr = "".join(chars[start:end])
- if start > 0:
- substr = "##" + substr
- if substr in self.vocab:
- cur_substr = substr
- break
- end -= 1
- if cur_substr is None:
- is_bad = True
- break
- sub_tokens.append(cur_substr)
- start = end
-
- if is_bad:
- output_tokens.append(self.unk_token)
- else:
- output_tokens.extend(sub_tokens)
- return output_tokens
diff --git a/transformers/models/electra/tokenization_electra_fast.py b/transformers/models/electra/tokenization_electra_fast.py
deleted file mode 100644
index 7b9d6a36cb92108d9e8796b5972e50f71d498af5..0000000000000000000000000000000000000000
--- a/transformers/models/electra/tokenization_electra_fast.py
+++ /dev/null
@@ -1,169 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The Google AI Team, Stanford University and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-from typing import List, Optional, Tuple
-
-from tokenizers import normalizers
-
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from .tokenization_electra import ElectraTokenizer
-
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
-
-
-# Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast with Bert->Electra , BERT->ELECTRA
-class ElectraTokenizerFast(PreTrainedTokenizerFast):
- r"""
- Construct a "fast" ELECTRA tokenizer (backed by HuggingFace's *tokenizers* library). Based on WordPiece.
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- clean_text (`bool`, *optional*, defaults to `True`):
- Whether or not to clean the text before tokenization by removing any control characters and replacing all
- whitespaces by the classic one.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see [this
- issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original ELECTRA).
- wordpieces_prefix (`str`, *optional*, defaults to `"##"`):
- The prefix for subwords.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = ElectraTokenizer
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- do_lower_case=True,
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- super().__init__(
- vocab_file,
- tokenizer_file=tokenizer_file,
- do_lower_case=do_lower_case,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- normalizer_state = json.loads(self.backend_tokenizer.normalizer.__getstate__())
- if (
- normalizer_state.get("lowercase", do_lower_case) != do_lower_case
- or normalizer_state.get("strip_accents", strip_accents) != strip_accents
- or normalizer_state.get("handle_chinese_chars", tokenize_chinese_chars) != tokenize_chinese_chars
- ):
- normalizer_class = getattr(normalizers, normalizer_state.pop("type"))
- normalizer_state["lowercase"] = do_lower_case
- normalizer_state["strip_accents"] = strip_accents
- normalizer_state["handle_chinese_chars"] = tokenize_chinese_chars
- self.backend_tokenizer.normalizer = normalizer_class(**normalizer_state)
-
- self.do_lower_case = do_lower_case
-
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A ELECTRA sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- output = [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
-
- if token_ids_1 is not None:
- output += token_ids_1 + [self.sep_token_id]
-
- return output
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A ELECTRA sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
diff --git a/transformers/models/encodec/__init__.py b/transformers/models/encodec/__init__.py
deleted file mode 100644
index d3d9488968bf2cc6316ba5eb4601e3dc3e5878b8..0000000000000000000000000000000000000000
--- a/transformers/models/encodec/__init__.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_encodec": [
- "ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "EncodecConfig",
- ],
- "feature_extraction_encodec": ["EncodecFeatureExtractor"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_encodec"] = [
- "ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST",
- "EncodecModel",
- "EncodecPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_encodec import (
- ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP,
- EncodecConfig,
- )
- from .feature_extraction_encodec import EncodecFeatureExtractor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_encodec import (
- ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST,
- EncodecModel,
- EncodecPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/encodec/__pycache__/__init__.cpython-310.pyc b/transformers/models/encodec/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 4c70acb4b97dc384b0c1a36a22ea41e93b8889ca..0000000000000000000000000000000000000000
Binary files a/transformers/models/encodec/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encodec/__pycache__/configuration_encodec.cpython-310.pyc b/transformers/models/encodec/__pycache__/configuration_encodec.cpython-310.pyc
deleted file mode 100644
index a3a6b0e189e7d4d4b4e915af6af677fc18ac2c39..0000000000000000000000000000000000000000
Binary files a/transformers/models/encodec/__pycache__/configuration_encodec.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encodec/__pycache__/convert_encodec_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/encodec/__pycache__/convert_encodec_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 3c35bcbae93b608e7fd938f07b41a0923dd43b4c..0000000000000000000000000000000000000000
Binary files a/transformers/models/encodec/__pycache__/convert_encodec_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encodec/__pycache__/feature_extraction_encodec.cpython-310.pyc b/transformers/models/encodec/__pycache__/feature_extraction_encodec.cpython-310.pyc
deleted file mode 100644
index fa5bc6848c6a5fe7cc2f942a0757ec68c05d115e..0000000000000000000000000000000000000000
Binary files a/transformers/models/encodec/__pycache__/feature_extraction_encodec.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encodec/__pycache__/modeling_encodec.cpython-310.pyc b/transformers/models/encodec/__pycache__/modeling_encodec.cpython-310.pyc
deleted file mode 100644
index de46edbb4326d703b2ade8a4de7d869ae307b9ac..0000000000000000000000000000000000000000
Binary files a/transformers/models/encodec/__pycache__/modeling_encodec.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encodec/configuration_encodec.py b/transformers/models/encodec/configuration_encodec.py
deleted file mode 100644
index 4e18bb178adf237b143c75e1406234ca36efadc3..0000000000000000000000000000000000000000
--- a/transformers/models/encodec/configuration_encodec.py
+++ /dev/null
@@ -1,193 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Meta Platforms, Inc. and affiliates, and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" EnCodec model configuration"""
-
-
-import math
-from typing import Optional
-
-import numpy as np
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class EncodecConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of an [`EncodecModel`]. It is used to instantiate a
- Encodec model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the
- [facebook/encodec_24khz](https://huggingface.co/facebook/encodec_24khz) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- target_bandwidths (`List[float]`, *optional*, defaults to `[1.5, 3.0, 6.0, 12.0, 24.0]`):
- The range of diffent bandwiths the model can encode audio with.
- sampling_rate (`int`, *optional*, defaults to 24000):
- The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz).
- audio_channels (`int`, *optional*, defaults to 1):
- Number of channels in the audio data. Either 1 for mono or 2 for stereo.
- normalize (`bool`, *optional*, defaults to `False`):
- Whether the audio shall be normalized when passed.
- chunk_length_s (`float`, *optional*):
- If defined the audio is pre-processed into chunks of lengths `chunk_length_s` and then encoded.
- overlap (`float`, *optional*):
- Defines the overlap between each chunk. It is used to compute the `chunk_stride` using the following
- formulae : `int((1.0 - self.overlap) * self.chunk_length)`.
- hidden_size (`int`, *optional*, defaults to 128):
- Intermediate representation dimension.
- num_filters (`int`, *optional*, defaults to 32):
- Number of convolution kernels of first `EncodecConv1d` down sampling layer.
- num_residual_layers (`int`, *optional*, defaults to 1):
- Number of residual layers.
- upsampling_ratios (`Sequence[int]` , *optional*, defaults to `[8, 5, 4, 2]`):
- Kernel size and stride ratios. The encoder uses downsampling ratios instead of upsampling ratios, hence it
- will use the ratios in the reverse order to the ones specified here that must match the decoder order.
- norm_type (`str`, *optional*, defaults to `"weight_norm"`):
- Normalization method. Should be in `["weight_norm", "time_group_norm"]`
- kernel_size (`int`, *optional*, defaults to 7):
- Kernel size for the initial convolution.
- last_kernel_size (`int`, *optional*, defaults to 7):
- Kernel size for the last convolution layer.
- residual_kernel_size (`int`, *optional*, defaults to 3):
- Kernel size for the residual layers.
- dilation_growth_rate (`int`, *optional*, defaults to 2):
- How much to increase the dilation with each layer.
- use_causal_conv (`bool`, *optional*, defaults to `True`):
- Whether to use fully causal convolution.
- pad_mode (`str`, *optional*, defaults to `"reflect"`):
- Padding mode for the convolutions.
- compress (`int`, *optional*, defaults to 2):
- Reduced dimensionality in residual branches (from Demucs v3).
- num_lstm_layers (`int`, *optional*, defaults to 2):
- Number of LSTM layers at the end of the encoder.
- trim_right_ratio (`float`, *optional*, defaults to 1.0):
- Ratio for trimming at the right of the transposed convolution under the `use_causal_conv = True` setup. If
- equal to 1.0, it means that all the trimming is done at the right.
- codebook_size (`int`, *optional*, defaults to 1024):
- Number of discret codes that make up VQVAE.
- codebook_dim (`int`, *optional*):
- Dimension of the codebook vectors. If not defined, uses `hidden_size`.
- use_conv_shortcut (`bool`, *optional*, defaults to `True`):
- Whether to use a convolutional layer as the 'skip' connection in the `EncodecResnetBlock` block. If False,
- an identity function will be used, giving a generic residual connection.
-
- Example:
-
- ```python
- >>> from transformers import EncodecModel, EncodecConfig
-
- >>> # Initializing a "facebook/encodec_24khz" style configuration
- >>> configuration = EncodecConfig()
-
- >>> # Initializing a model (with random weights) from the "facebook/encodec_24khz" style configuration
- >>> model = EncodecModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "encodec"
-
- def __init__(
- self,
- target_bandwidths=[1.5, 3.0, 6.0, 12.0, 24.0],
- sampling_rate=24_000,
- audio_channels=1,
- normalize=False,
- chunk_length_s=None,
- overlap=None,
- hidden_size=128,
- num_filters=32,
- num_residual_layers=1,
- upsampling_ratios=[8, 5, 4, 2],
- norm_type="weight_norm",
- kernel_size=7,
- last_kernel_size=7,
- residual_kernel_size=3,
- dilation_growth_rate=2,
- use_causal_conv=True,
- pad_mode="reflect",
- compress=2,
- num_lstm_layers=2,
- trim_right_ratio=1.0,
- codebook_size=1024,
- codebook_dim=None,
- use_conv_shortcut=True,
- **kwargs,
- ):
- self.target_bandwidths = target_bandwidths
- self.sampling_rate = sampling_rate
- self.audio_channels = audio_channels
- self.normalize = normalize
- self.chunk_length_s = chunk_length_s
- self.overlap = overlap
- self.hidden_size = hidden_size
- self.num_filters = num_filters
- self.num_residual_layers = num_residual_layers
- self.upsampling_ratios = upsampling_ratios
- self.norm_type = norm_type
- self.kernel_size = kernel_size
- self.last_kernel_size = last_kernel_size
- self.residual_kernel_size = residual_kernel_size
- self.dilation_growth_rate = dilation_growth_rate
- self.use_causal_conv = use_causal_conv
- self.pad_mode = pad_mode
- self.compress = compress
- self.num_lstm_layers = num_lstm_layers
- self.trim_right_ratio = trim_right_ratio
- self.codebook_size = codebook_size
- self.codebook_dim = codebook_dim if codebook_dim is not None else hidden_size
- self.use_conv_shortcut = use_conv_shortcut
-
- if self.norm_type not in ["weight_norm", "time_group_norm"]:
- raise ValueError(
- f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}'
- )
-
- super().__init__(**kwargs)
-
- # This is a property because you might want to change the chunk_length_s on the fly
- @property
- def chunk_length(self) -> Optional[int]:
- if self.chunk_length_s is None:
- return None
- else:
- return int(self.chunk_length_s * self.sampling_rate)
-
- # This is a property because you might want to change the chunk_length_s on the fly
- @property
- def chunk_stride(self) -> Optional[int]:
- if self.chunk_length_s is None or self.overlap is None:
- return None
- else:
- return max(1, int((1.0 - self.overlap) * self.chunk_length))
-
- @property
- def frame_rate(self) -> int:
- hop_length = np.prod(self.upsampling_ratios)
- return math.ceil(self.sampling_rate / hop_length)
-
- @property
- def num_quantizers(self) -> int:
- return int(1000 * self.target_bandwidths[-1] // (self.frame_rate * 10))
diff --git a/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py b/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py
deleted file mode 100644
index 3a16a4b7ba0f3b66412e63591055c3fb2afab9ec..0000000000000000000000000000000000000000
--- a/transformers/models/encodec/convert_encodec_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,365 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert EnCodec checkpoints."""
-
-import argparse
-
-import torch
-
-from transformers import (
- EncodecConfig,
- EncodecFeatureExtractor,
- EncodecModel,
- logging,
-)
-
-
-# checkpoints downloaded from:
-# https://dl.fbaipublicfiles.com/encodec/v0/encodec_24khz-d7cc33bc.th
-# https://huggingface.co/facebook/musicgen-small/resolve/main/compression_state_dict.bin
-# https://dl.fbaipublicfiles.com/encodec/v0/encodec_48khz-7e698e3e.th
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger("transformers.models.encodec")
-
-MAPPING_QUANTIZER = {
- "quantizer.vq.layers.*._codebook.inited": "quantizer.layers.*.codebook.inited",
- "quantizer.vq.layers.*._codebook.cluster_size": "quantizer.layers.*.codebook.cluster_size",
- "quantizer.vq.layers.*._codebook.embed": "quantizer.layers.*.codebook.embed",
- "quantizer.vq.layers.*._codebook.embed_avg": "quantizer.layers.*.codebook.embed_avg",
-}
-MAPPING_ENCODER = {
- "encoder.model.0.conv.conv": "encoder.layers.0.conv",
- "encoder.model.1.block.1.conv.conv": "encoder.layers.1.block.1.conv",
- "encoder.model.1.block.3.conv.conv": "encoder.layers.1.block.3.conv",
- "encoder.model.1.shortcut.conv.conv": "encoder.layers.1.shortcut.conv",
- "encoder.model.3.conv.conv": "encoder.layers.3.conv",
- "encoder.model.4.block.1.conv.conv": "encoder.layers.4.block.1.conv",
- "encoder.model.4.block.3.conv.conv": "encoder.layers.4.block.3.conv",
- "encoder.model.4.shortcut.conv.conv": "encoder.layers.4.shortcut.conv",
- "encoder.model.6.conv.conv": "encoder.layers.6.conv",
- "encoder.model.7.block.1.conv.conv": "encoder.layers.7.block.1.conv",
- "encoder.model.7.block.3.conv.conv": "encoder.layers.7.block.3.conv",
- "encoder.model.7.shortcut.conv.conv": "encoder.layers.7.shortcut.conv",
- "encoder.model.9.conv.conv": "encoder.layers.9.conv",
- "encoder.model.10.block.1.conv.conv": "encoder.layers.10.block.1.conv",
- "encoder.model.10.block.3.conv.conv": "encoder.layers.10.block.3.conv",
- "encoder.model.10.shortcut.conv.conv": "encoder.layers.10.shortcut.conv",
- "encoder.model.12.conv.conv": "encoder.layers.12.conv",
- "encoder.model.13.lstm": "encoder.layers.13.lstm",
- "encoder.model.15.conv.conv": "encoder.layers.15.conv",
-}
-MAPPING_ENCODER_48K = {
- "encoder.model.0.conv.norm": "encoder.layers.0.norm",
- "encoder.model.1.block.1.conv.norm": "encoder.layers.1.block.1.norm",
- "encoder.model.1.block.3.conv.norm": "encoder.layers.1.block.3.norm",
- "encoder.model.1.shortcut.conv.norm": "encoder.layers.1.shortcut.norm",
- "encoder.model.3.conv.norm": "encoder.layers.3.norm",
- "encoder.model.4.block.1.conv.norm": "encoder.layers.4.block.1.norm",
- "encoder.model.4.block.3.conv.norm": "encoder.layers.4.block.3.norm",
- "encoder.model.4.shortcut.conv.norm": "encoder.layers.4.shortcut.norm",
- "encoder.model.6.conv.norm": "encoder.layers.6.norm",
- "encoder.model.7.block.1.conv.norm": "encoder.layers.7.block.1.norm",
- "encoder.model.7.block.3.conv.norm": "encoder.layers.7.block.3.norm",
- "encoder.model.7.shortcut.conv.norm": "encoder.layers.7.shortcut.norm",
- "encoder.model.9.conv.norm": "encoder.layers.9.norm",
- "encoder.model.10.block.1.conv.norm": "encoder.layers.10.block.1.norm",
- "encoder.model.10.block.3.conv.norm": "encoder.layers.10.block.3.norm",
- "encoder.model.10.shortcut.conv.norm": "encoder.layers.10.shortcut.norm",
- "encoder.model.12.conv.norm": "encoder.layers.12.norm",
- "encoder.model.15.conv.norm": "encoder.layers.15.norm",
-}
-MAPPING_DECODER = {
- "decoder.model.0.conv.conv": "decoder.layers.0.conv",
- "decoder.model.1.lstm": "decoder.layers.1.lstm",
- "decoder.model.3.convtr.convtr": "decoder.layers.3.conv",
- "decoder.model.4.block.1.conv.conv": "decoder.layers.4.block.1.conv",
- "decoder.model.4.block.3.conv.conv": "decoder.layers.4.block.3.conv",
- "decoder.model.4.shortcut.conv.conv": "decoder.layers.4.shortcut.conv",
- "decoder.model.6.convtr.convtr": "decoder.layers.6.conv",
- "decoder.model.7.block.1.conv.conv": "decoder.layers.7.block.1.conv",
- "decoder.model.7.block.3.conv.conv": "decoder.layers.7.block.3.conv",
- "decoder.model.7.shortcut.conv.conv": "decoder.layers.7.shortcut.conv",
- "decoder.model.9.convtr.convtr": "decoder.layers.9.conv",
- "decoder.model.10.block.1.conv.conv": "decoder.layers.10.block.1.conv",
- "decoder.model.10.block.3.conv.conv": "decoder.layers.10.block.3.conv",
- "decoder.model.10.shortcut.conv.conv": "decoder.layers.10.shortcut.conv",
- "decoder.model.12.convtr.convtr": "decoder.layers.12.conv",
- "decoder.model.13.block.1.conv.conv": "decoder.layers.13.block.1.conv",
- "decoder.model.13.block.3.conv.conv": "decoder.layers.13.block.3.conv",
- "decoder.model.13.shortcut.conv.conv": "decoder.layers.13.shortcut.conv",
- "decoder.model.15.conv.conv": "decoder.layers.15.conv",
-}
-MAPPING_DECODER_48K = {
- "decoder.model.0.conv.norm": "decoder.layers.0.norm",
- "decoder.model.3.convtr.norm": "decoder.layers.3.norm",
- "decoder.model.4.block.1.conv.norm": "decoder.layers.4.block.1.norm",
- "decoder.model.4.block.3.conv.norm": "decoder.layers.4.block.3.norm",
- "decoder.model.4.shortcut.conv.norm": "decoder.layers.4.shortcut.norm",
- "decoder.model.6.convtr.norm": "decoder.layers.6.norm",
- "decoder.model.7.block.1.conv.norm": "decoder.layers.7.block.1.norm",
- "decoder.model.7.block.3.conv.norm": "decoder.layers.7.block.3.norm",
- "decoder.model.7.shortcut.conv.norm": "decoder.layers.7.shortcut.norm",
- "decoder.model.9.convtr.norm": "decoder.layers.9.norm",
- "decoder.model.10.block.1.conv.norm": "decoder.layers.10.block.1.norm",
- "decoder.model.10.block.3.conv.norm": "decoder.layers.10.block.3.norm",
- "decoder.model.10.shortcut.conv.norm": "decoder.layers.10.shortcut.norm",
- "decoder.model.12.convtr.norm": "decoder.layers.12.norm",
- "decoder.model.13.block.1.conv.norm": "decoder.layers.13.block.1.norm",
- "decoder.model.13.block.3.conv.norm": "decoder.layers.13.block.3.norm",
- "decoder.model.13.shortcut.conv.norm": "decoder.layers.13.shortcut.norm",
- "decoder.model.15.conv.norm": "decoder.layers.15.norm",
-}
-MAPPING_24K = {
- **MAPPING_QUANTIZER,
- **MAPPING_ENCODER,
- **MAPPING_DECODER,
-}
-MAPPING_48K = {
- **MAPPING_QUANTIZER,
- **MAPPING_ENCODER,
- **MAPPING_ENCODER_48K,
- **MAPPING_DECODER,
- **MAPPING_DECODER_48K,
-}
-TOP_LEVEL_KEYS = []
-IGNORE_KEYS = []
-
-
-def set_recursively(hf_pointer, key, value, full_name, weight_type):
- for attribute in key.split("."):
- hf_pointer = getattr(hf_pointer, attribute)
-
- if weight_type is not None:
- hf_shape = getattr(hf_pointer, weight_type).shape
- else:
- hf_shape = hf_pointer.shape
-
- if hf_shape != value.shape:
- raise ValueError(
- f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
- f" {value.shape} for {full_name}"
- )
-
- if weight_type == "weight":
- hf_pointer.weight.data = value
- elif weight_type == "weight_g":
- hf_pointer.weight_g.data = value
- elif weight_type == "weight_v":
- hf_pointer.weight_v.data = value
- elif weight_type == "bias":
- hf_pointer.bias.data = value
- elif weight_type == "running_mean":
- hf_pointer.running_mean.data = value
- elif weight_type == "running_var":
- hf_pointer.running_var.data = value
- elif weight_type == "num_batches_tracked":
- hf_pointer.num_batches_tracked.data = value
- elif weight_type == "weight_ih_l0":
- hf_pointer.weight_ih_l0.data = value
- elif weight_type == "weight_hh_l0":
- hf_pointer.weight_hh_l0.data = value
- elif weight_type == "bias_ih_l0":
- hf_pointer.bias_ih_l0.data = value
- elif weight_type == "bias_hh_l0":
- hf_pointer.bias_hh_l0.data = value
- elif weight_type == "weight_ih_l1":
- hf_pointer.weight_ih_l1.data = value
- elif weight_type == "weight_hh_l1":
- hf_pointer.weight_hh_l1.data = value
- elif weight_type == "bias_ih_l1":
- hf_pointer.bias_ih_l1.data = value
- elif weight_type == "bias_hh_l1":
- hf_pointer.bias_hh_l1.data = value
- else:
- hf_pointer.data = value
-
- logger.info(f"{key + ('.' + weight_type if weight_type is not None else '')} was initialized from {full_name}.")
-
-
-def should_ignore(name, ignore_keys):
- for key in ignore_keys:
- if key.endswith(".*"):
- if name.startswith(key[:-1]):
- return True
- elif ".*." in key:
- prefix, suffix = key.split(".*.")
- if prefix in name and suffix in name:
- return True
- elif key in name:
- return True
- return False
-
-
-def recursively_load_weights(orig_dict, hf_model, model_name):
- unused_weights = []
-
- if model_name == "encodec_24khz" or "encodec_32khz":
- MAPPING = MAPPING_24K
- elif model_name == "encodec_48khz":
- MAPPING = MAPPING_48K
- else:
- raise ValueError(f"Unsupported model: {model_name}")
-
- for name, value in orig_dict.items():
- if should_ignore(name, IGNORE_KEYS):
- logger.info(f"{name} was ignored")
- continue
-
- is_used = False
- for key, mapped_key in MAPPING.items():
- if "*" in key:
- prefix, suffix = key.split(".*.")
- if prefix in name and suffix in name:
- key = suffix
-
- if key in name:
- # HACK otherwise .embed gets initialized with .embed_avg too
- if key.endswith("embed") and name.endswith("embed_avg"):
- continue
-
- is_used = True
- if "*" in mapped_key:
- layer_index = name.split(key)[0].split(".")[-2]
- mapped_key = mapped_key.replace("*", layer_index)
- if "weight_g" in name:
- weight_type = "weight_g"
- elif "weight_v" in name:
- weight_type = "weight_v"
- elif "weight_ih_l0" in name:
- weight_type = "weight_ih_l0"
- elif "weight_hh_l0" in name:
- weight_type = "weight_hh_l0"
- elif "bias_ih_l0" in name:
- weight_type = "bias_ih_l0"
- elif "bias_hh_l0" in name:
- weight_type = "bias_hh_l0"
- elif "weight_ih_l1" in name:
- weight_type = "weight_ih_l1"
- elif "weight_hh_l1" in name:
- weight_type = "weight_hh_l1"
- elif "bias_ih_l1" in name:
- weight_type = "bias_ih_l1"
- elif "bias_hh_l1" in name:
- weight_type = "bias_hh_l1"
- elif "bias" in name:
- weight_type = "bias"
- elif "weight" in name:
- weight_type = "weight"
- elif "running_mean" in name:
- weight_type = "running_mean"
- elif "running_var" in name:
- weight_type = "running_var"
- elif "num_batches_tracked" in name:
- weight_type = "num_batches_tracked"
- else:
- weight_type = None
- set_recursively(hf_model, mapped_key, value, name, weight_type)
- continue
- if not is_used:
- unused_weights.append(name)
-
- logger.warning(f"Unused weights: {unused_weights}")
-
-
-@torch.no_grad()
-def convert_checkpoint(
- model_name,
- checkpoint_path,
- pytorch_dump_folder_path,
- config_path=None,
- repo_id=None,
-):
- """
- Copy/paste/tweak model's weights to transformers design.
- """
- if config_path is not None:
- config = EncodecConfig.from_pretrained(config_path)
- else:
- config = EncodecConfig()
-
- if model_name == "encodec_24khz":
- pass # config is already correct
- elif model_name == "encodec_32khz":
- config.upsampling_ratios = [8, 5, 4, 4]
- config.target_bandwidths = [2.2]
- config.num_filters = 64
- config.sampling_rate = 32_000
- config.codebook_size = 2048
- config.use_causal_conv = False
- config.normalize = False
- config.use_conv_shortcut = False
- elif model_name == "encodec_48khz":
- config.upsampling_ratios = [8, 5, 4, 2]
- config.target_bandwidths = [3.0, 6.0, 12.0, 24.0]
- config.sampling_rate = 48_000
- config.audio_channels = 2
- config.use_causal_conv = False
- config.norm_type = "time_group_norm"
- config.normalize = True
- config.chunk_length_s = 1.0
- config.overlap = 0.01
- else:
- raise ValueError(f"Unknown model name: {model_name}")
-
- model = EncodecModel(config)
-
- feature_extractor = EncodecFeatureExtractor(
- feature_size=config.audio_channels,
- sampling_rate=config.sampling_rate,
- chunk_length_s=config.chunk_length_s,
- overlap=config.overlap,
- )
- feature_extractor.save_pretrained(pytorch_dump_folder_path)
-
- original_checkpoint = torch.load(checkpoint_path)
- if "best_state" in original_checkpoint:
- # we might have a training state saved, in which case discard the yaml results and just retain the weights
- original_checkpoint = original_checkpoint["best_state"]
- recursively_load_weights(original_checkpoint, model, model_name)
- model.save_pretrained(pytorch_dump_folder_path)
-
- if repo_id:
- print("Pushing to the hub...")
- feature_extractor.push_to_hub(repo_id)
- model.push_to_hub(repo_id)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--model",
- default="encodec_24khz",
- type=str,
- help="The model to convert. Should be one of 'encodec_24khz', 'encodec_32khz', 'encodec_48khz'.",
- )
- parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint")
- parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
- parser.add_argument(
- "--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model."
- )
- parser.add_argument(
- "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
- )
-
- args = parser.parse_args()
- convert_checkpoint(
- args.model,
- args.checkpoint_path,
- args.pytorch_dump_folder_path,
- args.config_path,
- args.push_to_hub,
- )
diff --git a/transformers/models/encodec/feature_extraction_encodec.py b/transformers/models/encodec/feature_extraction_encodec.py
deleted file mode 100644
index 6f7536a52e9f99deeb97ffc9ef8accbbbed664d2..0000000000000000000000000000000000000000
--- a/transformers/models/encodec/feature_extraction_encodec.py
+++ /dev/null
@@ -1,206 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for EnCodec."""
-
-from typing import List, Optional, Union
-
-import numpy as np
-
-from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
-from ...feature_extraction_utils import BatchFeature
-from ...utils import PaddingStrategy, TensorType, logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class EncodecFeatureExtractor(SequenceFeatureExtractor):
- r"""
- Constructs an EnCodec feature extractor.
-
- This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains
- most of the main methods. Users should refer to this superclass for more information regarding those methods.
-
- Instantiating a feature extractor with the defaults will yield a similar configuration to that of the
- [facebook/encodec_24khz](https://huggingface.co/facebook/encodec_24khz) architecture.
-
- Args:
- feature_size (`int`, *optional*, defaults to 1):
- The feature dimension of the extracted features. Use 1 for mono, 2 for stereo.
- sampling_rate (`int`, *optional*, defaults to 24000):
- The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz).
- padding_value (`float`, *optional*, defaults to 0.0):
- The value that is used to fill the padding values.
- chunk_length_s (`float`, *optional*):
- If defined the audio is pre-processed into chunks of lengths `chunk_length_s` and then encoded.
- overlap (`float`, *optional*):
- Defines the overlap between each chunk. It is used to compute the `chunk_stride` using the following
- formulae : `int((1.0 - self.overlap) * self.chunk_length)`.
- """
-
- model_input_names = ["input_values", "padding_mask"]
-
- def __init__(
- self,
- feature_size: int = 1,
- sampling_rate: int = 24000,
- padding_value: float = 0.0,
- chunk_length_s: float = None,
- overlap: float = None,
- **kwargs,
- ):
- super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
- self.chunk_length_s = chunk_length_s
- self.overlap = overlap
-
- # This is a property because you might want to change the chunk_length_s on the fly
- @property
- def chunk_length(self) -> Optional[int]:
- if self.chunk_length_s is None:
- return None
- else:
- return int(self.chunk_length_s * self.sampling_rate)
-
- # This is a property because you might want to change the chunk_length_s on the fly
- @property
- def chunk_stride(self) -> Optional[int]:
- if self.chunk_length_s is None or self.overlap is None:
- return None
- else:
- return max(1, int((1.0 - self.overlap) * self.chunk_length))
-
- def __call__(
- self,
- raw_audio: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],
- padding: Optional[Union[bool, str, PaddingStrategy]] = None,
- truncation: Optional[bool] = False,
- max_length: Optional[int] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- sampling_rate: Optional[int] = None,
- ) -> BatchFeature:
- """
- Main method to featurize and prepare for the model one or several sequence(s).
-
- Args:
- raw_audio (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`):
- The sequence or batch of sequences to be processed. Each sequence can be a numpy array, a list of float
- values, a list of numpy arrays or a list of list of float values. The numpy array must be of shape
- `(num_samples,)` for mono audio (`feature_size = 1`), or `(2, num_samples)` for stereo audio
- (`feature_size = 2`).
- padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
- Select a strategy to pad the returned sequences (according to the model's padding side and padding
- index) among:
-
- - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
- sequence if provided).
- - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
- acceptable input length for the model if that argument is not provided.
- - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
- lengths).
- truncation (`bool`, *optional*, defaults to `False`):
- Activates truncation to cut input sequences longer than `max_length` to `max_length`.
- max_length (`int`, *optional*):
- Maximum length of the returned list and optionally padding length (see above).
- return_tensors (`str` or [`~utils.TensorType`], *optional*):
- If set, will return tensors instead of list of python integers. Acceptable values are:
-
- - `'tf'`: Return TensorFlow `tf.constant` objects.
- - `'pt'`: Return PyTorch `torch.Tensor` objects.
- - `'np'`: Return Numpy `np.ndarray` objects.
- sampling_rate (`int`, *optional*):
- The sampling rate at which the `audio` input was sampled. It is strongly recommended to pass
- `sampling_rate` at the forward call to prevent silent errors.
- """
- if sampling_rate is not None:
- if sampling_rate != self.sampling_rate:
- raise ValueError(
- f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of"
- f" {self.sampling_rate}. Please make sure that the provided audio input was sampled with"
- f" {self.sampling_rate} and not {sampling_rate}."
- )
- else:
- logger.warning(
- "It is strongly recommended to pass the `sampling_rate` argument to this function. "
- "Failing to do so can result in silent errors that might be hard to debug."
- )
-
- if padding and truncation:
- raise ValueError("Both padding and truncation were set. Make sure you only set one.")
- elif padding is None:
- # by default let's pad the inputs
- padding = True
-
- is_batched = bool(
- isinstance(raw_audio, (list, tuple)) and (isinstance(raw_audio[0], (np.ndarray, tuple, list)))
- )
-
- if is_batched:
- raw_audio = [np.asarray(audio, dtype=np.float32).T for audio in raw_audio]
- elif not is_batched and not isinstance(raw_audio, np.ndarray):
- raw_audio = np.asarray(raw_audio, dtype=np.float32)
- elif isinstance(raw_audio, np.ndarray) and raw_audio.dtype is np.dtype(np.float64):
- raw_audio = raw_audio.astype(np.float32)
-
- # always return batch
- if not is_batched:
- raw_audio = [np.asarray(raw_audio).T]
-
- # verify inputs are valid
- for idx, example in enumerate(raw_audio):
- if example.ndim > 2:
- raise ValueError(f"Expected input shape (channels, length) but got shape {example.shape}")
- if self.feature_size == 1 and example.ndim != 1:
- raise ValueError(f"Expected mono audio but example has {example.shape[-1]} channels")
- if self.feature_size == 2 and example.shape[-1] != 2:
- raise ValueError(f"Expected stereo audio but example has {example.shape[-1]} channels")
-
- padded_inputs = None
- input_values = BatchFeature({"input_values": raw_audio})
- if self.chunk_stride is not None and self.chunk_length is not None and max_length is None:
- if truncation:
- max_length = min(array.shape[0] for array in raw_audio)
- nb_step = int(np.floor(max_length / self.chunk_stride))
- max_length = (nb_step - 1) * self.chunk_stride + self.chunk_length
- elif padding:
- max_length = max(array.shape[0] for array in raw_audio)
- nb_step = int(np.ceil(max_length / self.chunk_stride))
- max_length = (nb_step - 1) * self.chunk_stride + self.chunk_length
- padding = "max_length"
- else:
- padded_inputs = input_values
-
- # normal padding on batch
- if padded_inputs is None:
- padded_inputs = self.pad(
- input_values,
- max_length=max_length,
- truncation=truncation,
- padding=padding,
- return_attention_mask=padding,
- )
- if padding:
- padded_inputs["padding_mask"] = padded_inputs.pop("attention_mask")
-
- input_values = []
- for example in padded_inputs.pop("input_values"):
- if self.feature_size == 1:
- example = example[..., None]
- input_values.append(example.T)
-
- padded_inputs["input_values"] = input_values
- if return_tensors is not None:
- padded_inputs = padded_inputs.convert_to_tensors(return_tensors)
-
- return padded_inputs
diff --git a/transformers/models/encodec/modeling_encodec.py b/transformers/models/encodec/modeling_encodec.py
deleted file mode 100644
index 48498b741d18cab693df38fc35997b3aac8371e6..0000000000000000000000000000000000000000
--- a/transformers/models/encodec/modeling_encodec.py
+++ /dev/null
@@ -1,810 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Meta Platforms, Inc. and affiliates, and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch EnCodec model."""
-
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_encodec import EncodecConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-# General docstring
-_CONFIG_FOR_DOC = "EncodecConfig"
-
-
-from ..deprecated._archive_maps import ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class EncodecOutput(ModelOutput):
- """
- Args:
- audio_codes (`torch.LongTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*):
- Discret code embeddings computed using `model.encode`.
- audio_values (`torch.FlaotTensor` of shape `(batch_size, sequence_length)`, *optional*)
- Decoded audio values, obtained using the decoder part of Encodec.
- """
-
- audio_codes: torch.LongTensor = None
- audio_values: torch.FloatTensor = None
-
-
-@dataclass
-class EncodecEncoderOutput(ModelOutput):
- """
- Args:
- audio_codes (`torch.LongTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*):
- Discret code embeddings computed using `model.encode`.
- audio_scales (`torch.Tensor` of shape `(batch_size, nb_chunks)`, *optional*):
- Scaling factor for each `audio_codes` input. This is used to unscale each chunk of audio when decoding.
- """
-
- audio_codes: torch.LongTensor = None
- audio_scales: torch.FloatTensor = None
-
-
-@dataclass
-class EncodecDecoderOutput(ModelOutput):
- """
- Args:
- audio_values (`torch.FloatTensor` of shape `(batch_size, segment_length)`, *optional*):
- Decoded audio values, obtained using the decoder part of Encodec.
- """
-
- audio_values: torch.FloatTensor = None
-
-
-class EncodecConv1d(nn.Module):
- """Conv1d with asymmetric or causal padding and normalization."""
-
- def __init__(
- self, config, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, dilation: int = 1
- ):
- super().__init__()
- self.causal = config.use_causal_conv
- self.pad_mode = config.pad_mode
- self.norm_type = config.norm_type
-
- if self.norm_type not in ["weight_norm", "time_group_norm"]:
- raise ValueError(
- f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}'
- )
-
- # warn user on unusual setup between dilation and stride
- if stride > 1 and dilation > 1:
- logger.warning(
- "EncodecConv1d has been initialized with stride > 1 and dilation > 1"
- f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
- )
-
- self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, dilation=dilation)
- if self.norm_type == "weight_norm":
- self.conv = nn.utils.weight_norm(self.conv)
- elif self.norm_type == "time_group_norm":
- self.norm = nn.GroupNorm(1, out_channels)
-
- kernel_size = self.conv.kernel_size[0]
- stride = torch.tensor(self.conv.stride[0], dtype=torch.int64)
- dilation = self.conv.dilation[0]
-
- # Effective kernel size with dilations.
- kernel_size = torch.tensor((kernel_size - 1) * dilation + 1, dtype=torch.int64)
-
- self.register_buffer("stride", stride, persistent=False)
- self.register_buffer("kernel_size", kernel_size, persistent=False)
- self.register_buffer("padding_total", torch.tensor(kernel_size - stride, dtype=torch.int64), persistent=False)
-
- def _get_extra_padding_for_conv1d(
- self,
- hidden_states: torch.Tensor,
- ) -> torch.Tensor:
- """See `pad_for_conv1d`."""
- length = hidden_states.shape[-1]
- n_frames = (length - self.kernel_size + self.padding_total) / self.stride + 1
- n_frames = torch.ceil(n_frames).to(torch.int64) - 1
- ideal_length = n_frames * self.stride + self.kernel_size - self.padding_total
-
- return ideal_length - length
-
- @staticmethod
- def _pad1d(hidden_states: torch.Tensor, paddings: Tuple[int, int], mode: str = "zero", value: float = 0.0):
- """Tiny wrapper around torch.nn.functional.pad, just to allow for reflect padding on small input.
- If this is the case, we insert extra 0 padding to the right before the reflection happens.
- """
- length = hidden_states.shape[-1]
- padding_left, padding_right = paddings
- if not mode == "reflect":
- return nn.functional.pad(hidden_states, paddings, mode, value)
-
- max_pad = max(padding_left, padding_right)
- extra_pad = 0
- if length <= max_pad:
- extra_pad = max_pad - length + 1
- hidden_states = nn.functional.pad(hidden_states, (0, extra_pad))
- padded = nn.functional.pad(hidden_states, paddings, mode, value)
- end = padded.shape[-1] - extra_pad
- return padded[..., :end]
-
- def forward(self, hidden_states):
- extra_padding = self._get_extra_padding_for_conv1d(hidden_states)
-
- if self.causal:
- # Left padding for causal
- hidden_states = self._pad1d(hidden_states, (self.padding_total, extra_padding), mode=self.pad_mode)
- else:
- # Asymmetric padding required for odd strides
- padding_right = self.padding_total // 2
- padding_left = self.padding_total - padding_right
- hidden_states = self._pad1d(
- hidden_states, (padding_left, padding_right + extra_padding), mode=self.pad_mode
- )
-
- hidden_states = self.conv(hidden_states)
-
- if self.norm_type == "time_group_norm":
- hidden_states = self.norm(hidden_states)
-
- return hidden_states
-
-
-class EncodecConvTranspose1d(nn.Module):
- """ConvTranspose1d with asymmetric or causal padding and normalization."""
-
- def __init__(self, config, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1):
- super().__init__()
- self.causal = config.use_causal_conv
- self.trim_right_ratio = config.trim_right_ratio
- self.norm_type = config.norm_type
- if self.norm_type not in ["weight_norm", "time_group_norm"]:
- raise ValueError(
- f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}'
- )
-
- self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride)
- if config.norm_type == "weight_norm":
- self.conv = nn.utils.weight_norm(self.conv)
- elif config.norm_type == "time_group_norm":
- self.norm = nn.GroupNorm(1, out_channels)
-
- if not (self.causal or self.trim_right_ratio == 1.0):
- raise ValueError("`trim_right_ratio` != 1.0 only makes sense for causal convolutions")
-
- def forward(self, hidden_states):
- kernel_size = self.conv.kernel_size[0]
- stride = self.conv.stride[0]
- padding_total = kernel_size - stride
-
- hidden_states = self.conv(hidden_states)
-
- if self.norm_type == "time_group_norm":
- hidden_states = self.norm(hidden_states)
-
- # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
- # removed at the very end, when keeping only the right length for the output,
- # as removing it here would require also passing the length at the matching layer
- # in the encoder.
- if self.causal:
- # Trim the padding on the right according to the specified ratio
- # if trim_right_ratio = 1.0, trim everything from right
- padding_right = math.ceil(padding_total * self.trim_right_ratio)
- else:
- # Asymmetric padding required for odd strides
- padding_right = padding_total // 2
-
- padding_left = padding_total - padding_right
-
- # unpad
- end = hidden_states.shape[-1] - padding_right
- hidden_states = hidden_states[..., padding_left:end]
- return hidden_states
-
-
-class EncodecLSTM(nn.Module):
- """
- LSTM without worrying about the hidden state, nor the layout of the data. Expects input as convolutional layout.
- """
-
- def __init__(self, config, dimension):
- super().__init__()
- self.lstm = nn.LSTM(dimension, dimension, config.num_lstm_layers)
-
- def forward(self, hidden_states):
- hidden_states = hidden_states.permute(2, 0, 1)
- hidden_states = self.lstm(hidden_states)[0] + hidden_states
- hidden_states = hidden_states.permute(1, 2, 0)
- return hidden_states
-
-
-class EncodecResnetBlock(nn.Module):
- """
- Residual block from SEANet model as used by EnCodec.
- """
-
- def __init__(self, config: EncodecConfig, dim: int, dilations: List[int]):
- super().__init__()
- kernel_sizes = (config.residual_kernel_size, 1)
- if len(kernel_sizes) != len(dilations):
- raise ValueError("Number of kernel sizes should match number of dilations")
-
- hidden = dim // config.compress
- block = []
- for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
- in_chs = dim if i == 0 else hidden
- out_chs = dim if i == len(kernel_sizes) - 1 else hidden
- block += [nn.ELU()]
- block += [EncodecConv1d(config, in_chs, out_chs, kernel_size, dilation=dilation)]
- self.block = nn.ModuleList(block)
-
- if config.use_conv_shortcut:
- self.shortcut = EncodecConv1d(config, dim, dim, kernel_size=1)
- else:
- self.shortcut = nn.Identity()
-
- def forward(self, hidden_states):
- residual = hidden_states
- for layer in self.block:
- hidden_states = layer(hidden_states)
-
- return self.shortcut(residual) + hidden_states
-
-
-class EncodecEncoder(nn.Module):
- """SEANet encoder as used by EnCodec."""
-
- def __init__(self, config: EncodecConfig):
- super().__init__()
- model = [EncodecConv1d(config, config.audio_channels, config.num_filters, config.kernel_size)]
- scaling = 1
-
- # Downsample to raw audio scale
- for ratio in reversed(config.upsampling_ratios):
- current_scale = scaling * config.num_filters
- # Add residual layers
- for j in range(config.num_residual_layers):
- model += [EncodecResnetBlock(config, current_scale, [config.dilation_growth_rate**j, 1])]
- # Add downsampling layers
- model += [nn.ELU()]
- model += [EncodecConv1d(config, current_scale, current_scale * 2, kernel_size=ratio * 2, stride=ratio)]
- scaling *= 2
-
- model += [EncodecLSTM(config, scaling * config.num_filters)]
- model += [nn.ELU()]
- model += [EncodecConv1d(config, scaling * config.num_filters, config.hidden_size, config.last_kernel_size)]
-
- self.layers = nn.ModuleList(model)
-
- def forward(self, hidden_states):
- for layer in self.layers:
- hidden_states = layer(hidden_states)
- return hidden_states
-
-
-class EncodecDecoder(nn.Module):
- """SEANet decoder as used by EnCodec."""
-
- def __init__(self, config: EncodecConfig):
- super().__init__()
- scaling = int(2 ** len(config.upsampling_ratios))
- model = [EncodecConv1d(config, config.hidden_size, scaling * config.num_filters, config.kernel_size)]
-
- model += [EncodecLSTM(config, scaling * config.num_filters)]
-
- # Upsample to raw audio scale
- for ratio in config.upsampling_ratios:
- current_scale = scaling * config.num_filters
- # Add upsampling layers
- model += [nn.ELU()]
- model += [
- EncodecConvTranspose1d(config, current_scale, current_scale // 2, kernel_size=ratio * 2, stride=ratio)
- ]
- # Add residual layers
- for j in range(config.num_residual_layers):
- model += [EncodecResnetBlock(config, current_scale // 2, (config.dilation_growth_rate**j, 1))]
- scaling //= 2
-
- # Add final layers
- model += [nn.ELU()]
- model += [EncodecConv1d(config, config.num_filters, config.audio_channels, config.last_kernel_size)]
- self.layers = nn.ModuleList(model)
-
- def forward(self, hidden_states):
- for layer in self.layers:
- hidden_states = layer(hidden_states)
- return hidden_states
-
-
-class EncodecEuclideanCodebook(nn.Module):
- """Codebook with Euclidean distance."""
-
- def __init__(self, config: EncodecConfig):
- super().__init__()
- embed = torch.zeros(config.codebook_size, config.codebook_dim)
-
- self.codebook_size = config.codebook_size
-
- self.register_buffer("inited", torch.Tensor([True]))
- self.register_buffer("cluster_size", torch.zeros(config.codebook_size))
- self.register_buffer("embed", embed)
- self.register_buffer("embed_avg", embed.clone())
-
- def quantize(self, hidden_states):
- embed = self.embed.t()
- scaled_states = hidden_states.pow(2).sum(1, keepdim=True)
- dist = -(scaled_states - 2 * hidden_states @ embed + embed.pow(2).sum(0, keepdim=True))
- embed_ind = dist.max(dim=-1).indices
- return embed_ind
-
- def encode(self, hidden_states):
- shape = hidden_states.shape
- # pre-process
- hidden_states = hidden_states.reshape((-1, shape[-1]))
- # quantize
- embed_ind = self.quantize(hidden_states)
- # post-process
- embed_ind = embed_ind.view(*shape[:-1])
- return embed_ind
-
- def decode(self, embed_ind):
- quantize = nn.functional.embedding(embed_ind, self.embed)
- return quantize
-
-
-class EncodecVectorQuantization(nn.Module):
- """
- Vector quantization implementation. Currently supports only euclidean distance.
- """
-
- def __init__(self, config: EncodecConfig):
- super().__init__()
- self.codebook = EncodecEuclideanCodebook(config)
-
- def encode(self, hidden_states):
- hidden_states = hidden_states.permute(0, 2, 1)
- embed_in = self.codebook.encode(hidden_states)
- return embed_in
-
- def decode(self, embed_ind):
- quantize = self.codebook.decode(embed_ind)
- quantize = quantize.permute(0, 2, 1)
- return quantize
-
-
-class EncodecResidualVectorQuantizer(nn.Module):
- """Residual Vector Quantizer."""
-
- def __init__(self, config: EncodecConfig):
- super().__init__()
- self.codebook_size = config.codebook_size
- self.frame_rate = config.frame_rate
- self.num_quantizers = config.num_quantizers
- self.layers = nn.ModuleList([EncodecVectorQuantization(config) for _ in range(config.num_quantizers)])
-
- def get_num_quantizers_for_bandwidth(self, bandwidth: Optional[float] = None) -> int:
- """Return num_quantizers based on specified target bandwidth."""
- bw_per_q = math.log2(self.codebook_size) * self.frame_rate
- num_quantizers = self.num_quantizers
- if bandwidth is not None and bandwidth > 0.0:
- num_quantizers = int(max(1, math.floor(bandwidth * 1000 / bw_per_q)))
- return num_quantizers
-
- def encode(self, embeddings: torch.Tensor, bandwidth: Optional[float] = None) -> torch.Tensor:
- """
- Encode a given input tensor with the specified frame rate at the given bandwidth. The RVQ encode method sets
- the appropriate number of quantizers to use and returns indices for each quantizer.
- """
- num_quantizers = self.get_num_quantizers_for_bandwidth(bandwidth)
- residual = embeddings
- all_indices = []
- for layer in self.layers[:num_quantizers]:
- indices = layer.encode(residual)
- quantized = layer.decode(indices)
- residual = residual - quantized
- all_indices.append(indices)
- out_indices = torch.stack(all_indices)
- return out_indices
-
- def decode(self, codes: torch.Tensor) -> torch.Tensor:
- """Decode the given codes to the quantized representation."""
- quantized_out = torch.tensor(0.0, device=codes.device)
- for i, indices in enumerate(codes):
- layer = self.layers[i]
- quantized = layer.decode(indices)
- quantized_out = quantized_out + quantized
- return quantized_out
-
-
-class EncodecPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = EncodecConfig
- base_model_prefix = "encodec"
- main_input_name = "input_values"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
- elif isinstance(module, nn.Conv1d):
- nn.init.kaiming_normal_(module.weight)
- if module.bias is not None:
- k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))
- nn.init.uniform_(module.bias, a=-k, b=k)
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LSTM):
- for name, param in module.named_parameters():
- if "weight" in name:
- nn.init.xavier_uniform_(param)
- elif "bias" in name:
- nn.init.constant_(param, 0.0)
-
-
-ENCODEC_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`EncodecConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-ENCODEC_INPUTS_DOCSTRING = r"""
- Args:
- input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*):
- Raw audio input converted to Float and padded to the approriate length in order to be encoded using chunks
- of length self.chunk_length and a stride of `config.chunk_stride`.
- padding_mask (`torch.BoolTensor` of shape `(batch_size, channels, sequence_length)`, *optional*):
- Mask to avoid computing scaling factors on padding token indices (can we avoid computing conv on these+).
- Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
-
-
- `padding_mask` should always be passed, unless the input was truncated or not padded. This is because in
- order to process tensors effectively, the input audio should be padded so that `input_length % stride =
- step` with `step = chunk_length-stride`. This ensures that all chunks are of the same shape
-
-
-
- bandwidth (`float`, *optional*):
- The target bandwidth. Must be one of `config.target_bandwidths`. If `None`, uses the smallest possible
- bandwidth. bandwidth is represented as a thousandth of what it is, e.g. 6kbps bandwidth is represented as
- `bandwidth == 6.0`
- audio_codes (`torch.LongTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*):
- Discret code embeddings computed using `model.encode`.
- audio_scales (`torch.Tensor` of shape `(batch_size, nb_chunks)`, *optional*):
- Scaling factor for each `audio_codes` input.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The EnCodec neural audio codec model.",
- ENCODEC_START_DOCSTRING,
-)
-class EncodecModel(EncodecPreTrainedModel):
- def __init__(self, config: EncodecConfig):
- super().__init__(config)
- self.config = config
-
- self.encoder = EncodecEncoder(config)
- self.decoder = EncodecDecoder(config)
-
- self.quantizer = EncodecResidualVectorQuantizer(config)
-
- self.bits_per_codebook = int(math.log2(self.config.codebook_size))
- if 2**self.bits_per_codebook != self.config.codebook_size:
- raise ValueError("The codebook_size must be a power of 2.")
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_encoder(self):
- return self.encoder
-
- def get_decoder(self):
- return self.decoder
-
- def _encode_frame(
- self, input_values: torch.Tensor, bandwidth: float, padding_mask: int
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- """
- Encodes the given input using the underlying VQVAE. If `config.normalize` is set to `True` the input is first
- normalized. The padding mask is required to compute the correct scale.
- """
- length = input_values.shape[-1]
- duration = length / self.config.sampling_rate
-
- if self.config.chunk_length_s is not None and duration > 1e-5 + self.config.chunk_length_s:
- raise RuntimeError(f"Duration of frame ({duration}) is longer than chunk {self.config.chunk_length_s}")
-
- scale = None
- if self.config.normalize:
- # if the padding is non zero
- input_values = input_values * padding_mask
- mono = torch.sum(input_values, 1, keepdim=True) / input_values.shape[1]
- scale = mono.pow(2).mean(dim=-1, keepdim=True).sqrt() + 1e-8
- input_values = input_values / scale
-
- embeddings = self.encoder(input_values)
- codes = self.quantizer.encode(embeddings, bandwidth)
- codes = codes.transpose(0, 1)
- return codes, scale
-
- def encode(
- self,
- input_values: torch.Tensor,
- padding_mask: torch.Tensor = None,
- bandwidth: Optional[float] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor, Optional[torch.Tensor]], EncodecEncoderOutput]:
- """
- Encodes the input audio waveform into discrete codes.
-
- Args:
- input_values (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
- Float values of the input audio waveform.
- padding_mask (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
- Padding mask used to pad the `input_values`.
- bandwidth (`float`, *optional*):
- The target bandwidth. Must be one of `config.target_bandwidths`. If `None`, uses the smallest possible
- bandwidth. bandwidth is represented as a thousandth of what it is, e.g. 6kbps bandwidth is represented
- as bandwidth == 6.0
-
- Returns:
- A list of frames containing the discrete encoded codes for the input audio waveform, along with rescaling
- factors for each chunk when `normalize` is True. Each frames is a tuple `(codebook, scale)`, with
- `codebook` of shape `[batch_size, num_codebooks, frames]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- if bandwidth is None:
- bandwidth = self.config.target_bandwidths[0]
- if bandwidth not in self.config.target_bandwidths:
- raise ValueError(
- f"This model doesn't support the bandwidth {bandwidth}. "
- f"Select one of {self.config.target_bandwidths}."
- )
-
- _, channels, input_length = input_values.shape
-
- if channels < 1 or channels > 2:
- raise ValueError(f"Number of audio channels must be 1 or 2, but got {channels}")
-
- chunk_length = self.config.chunk_length
- if chunk_length is None:
- chunk_length = input_length
- stride = input_length
- else:
- stride = self.config.chunk_stride
-
- if padding_mask is None:
- padding_mask = torch.ones_like(input_values).bool()
-
- encoded_frames = []
- scales = []
-
- step = chunk_length - stride
- if (input_length % stride) - step != 0:
- raise ValueError(
- "The input length is not properly padded for batched chunked decoding. Make sure to pad the input correctly."
- )
-
- for offset in range(0, input_length - step, stride):
- mask = padding_mask[..., offset : offset + chunk_length].bool()
- frame = input_values[:, :, offset : offset + chunk_length]
- encoded_frame, scale = self._encode_frame(frame, bandwidth, mask)
- encoded_frames.append(encoded_frame)
- scales.append(scale)
-
- encoded_frames = torch.stack(encoded_frames)
-
- if not return_dict:
- return (encoded_frames, scales)
-
- return EncodecEncoderOutput(encoded_frames, scales)
-
- @staticmethod
- def _linear_overlap_add(frames: List[torch.Tensor], stride: int):
- # Generic overlap add, with linear fade-in/fade-out, supporting complex scenario
- # e.g., more than 2 frames per position.
- # The core idea is to use a weight function that is a triangle,
- # with a maximum value at the middle of the chunk.
- # We use this weighting when summing the frames, and divide by the sum of weights
- # for each positions at the end. Thus:
- # - if a frame is the only one to cover a position, the weighting is a no-op.
- # - if 2 frames cover a position:
- # ... ...
- # / \/ \
- # / /\ \
- # S T , i.e. S offset of second frame starts, T end of first frame.
- # Then the weight function for each one is: (t - S), (T - t), with `t` a given offset.
- # After the final normalization, the weight of the second frame at position `t` is
- # (t - S) / (t - S + (T - t)) = (t - S) / (T - S), which is exactly what we want.
- #
- # - if more than 2 frames overlap at a given point, we hope that by induction
- # something sensible happens.
- if len(frames) == 0:
- raise ValueError("`frames` cannot be an empty list.")
-
- device = frames[0].device
- dtype = frames[0].dtype
- shape = frames[0].shape[:-1]
- total_size = stride * (len(frames) - 1) + frames[-1].shape[-1]
-
- frame_length = frames[0].shape[-1]
- time_vec = torch.linspace(0, 1, frame_length + 2, device=device, dtype=dtype)[1:-1]
- weight = 0.5 - (time_vec - 0.5).abs()
-
- sum_weight = torch.zeros(total_size, device=device, dtype=dtype)
- out = torch.zeros(*shape, total_size, device=device, dtype=dtype)
- offset: int = 0
-
- for frame in frames:
- frame_length = frame.shape[-1]
- out[..., offset : offset + frame_length] += weight[:frame_length] * frame
- sum_weight[offset : offset + frame_length] += weight[:frame_length]
- offset += stride
-
- if sum_weight.min() == 0:
- raise ValueError(f"`sum_weight` minimum element must be bigger than zero: {sum_weight}`")
-
- return out / sum_weight
-
- def _decode_frame(self, codes: torch.Tensor, scale: Optional[torch.Tensor] = None) -> torch.Tensor:
- codes = codes.transpose(0, 1)
- embeddings = self.quantizer.decode(codes)
- outputs = self.decoder(embeddings)
- if scale is not None:
- outputs = outputs * scale.view(-1, 1, 1)
- return outputs
-
- def decode(
- self,
- audio_codes: torch.Tensor,
- audio_scales: torch.Tensor,
- padding_mask: Optional[torch.Tensor] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], EncodecDecoderOutput]:
- """
- Decodes the given frames into an output audio waveform.
-
- Note that the output might be a bit bigger than the input. In that case, any extra steps at the end can be
- trimmed.
-
- Args:
- audio_codes (`torch.LongTensor` of shape `(batch_size, nb_chunks, chunk_length)`, *optional*):
- Discret code embeddings computed using `model.encode`.
- audio_scales (`torch.Tensor` of shape `(batch_size, nb_chunks)`, *optional*):
- Scaling factor for each `audio_codes` input.
- padding_mask (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
- Padding mask used to pad the `input_values`.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-
- """
- return_dict = return_dict or self.config.return_dict
-
- chunk_length = self.config.chunk_length
- if chunk_length is None:
- if len(audio_codes) != 1:
- raise ValueError(f"Expected one frame, got {len(audio_codes)}")
- audio_values = self._decode_frame(audio_codes[0], audio_scales[0])
- else:
- decoded_frames = []
-
- for frame, scale in zip(audio_codes, audio_scales):
- frames = self._decode_frame(frame, scale)
- decoded_frames.append(frames)
-
- audio_values = self._linear_overlap_add(decoded_frames, self.config.chunk_stride or 1)
-
- # truncate based on padding mask
- if padding_mask is not None and padding_mask.shape[-1] < audio_values.shape[-1]:
- audio_values = audio_values[..., : padding_mask.shape[-1]]
-
- if not return_dict:
- return (audio_values,)
- return EncodecDecoderOutput(audio_values)
-
- @add_start_docstrings_to_model_forward(ENCODEC_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=EncodecOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_values: torch.Tensor,
- padding_mask: Optional[torch.Tensor] = None,
- bandwidth: Optional[float] = None,
- audio_codes: Optional[torch.Tensor] = None,
- audio_scales: Optional[torch.Tensor] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], EncodecOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from datasets import load_dataset
- >>> from transformers import AutoProcessor, EncodecModel
-
- >>> dataset = load_dataset("hf-internal-testing/ashraq-esc50-1-dog-example")
- >>> audio_sample = dataset["train"]["audio"][0]["array"]
-
- >>> model_id = "facebook/encodec_24khz"
- >>> model = EncodecModel.from_pretrained(model_id)
- >>> processor = AutoProcessor.from_pretrained(model_id)
-
- >>> inputs = processor(raw_audio=audio_sample, return_tensors="pt")
-
- >>> outputs = model(**inputs)
- >>> audio_codes = outputs.audio_codes
- >>> audio_values = outputs.audio_values
- ```"""
- return_dict = return_dict or self.config.return_dict
-
- if padding_mask is None:
- padding_mask = torch.ones_like(input_values).bool()
-
- if audio_codes is not None and audio_scales is None:
- raise ValueError("You specified `audio_codes` but did not specify the `audio_scales`")
-
- if audio_scales is not None and audio_codes is None:
- raise ValueError("You specified `audio_scales` but did not specify the `audio_codes`")
-
- if audio_scales is None and audio_codes is None:
- audio_codes, audio_scales = self.encode(input_values, padding_mask, bandwidth, False)
-
- audio_values = self.decode(audio_codes, audio_scales, padding_mask, return_dict=return_dict)[0]
- if not return_dict:
- return (audio_codes, audio_values)
-
- return EncodecOutput(audio_codes=audio_codes, audio_values=audio_values)
diff --git a/transformers/models/encoder_decoder/__init__.py b/transformers/models/encoder_decoder/__init__.py
deleted file mode 100644
index ba71f1f7c7a9e121cf3bdda9c1604cb5021a8a3b..0000000000000000000000000000000000000000
--- a/transformers/models/encoder_decoder/__init__.py
+++ /dev/null
@@ -1,82 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_flax_available,
- is_tf_available,
- is_torch_available,
-)
-
-
-_import_structure = {"configuration_encoder_decoder": ["EncoderDecoderConfig"]}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_encoder_decoder"] = ["EncoderDecoderModel"]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_encoder_decoder"] = ["TFEncoderDecoderModel"]
-
-try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flax_encoder_decoder"] = ["FlaxEncoderDecoderModel"]
-
-if TYPE_CHECKING:
- from .configuration_encoder_decoder import EncoderDecoderConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_encoder_decoder import EncoderDecoderModel
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_encoder_decoder import TFEncoderDecoderModel
-
- try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flax_encoder_decoder import FlaxEncoderDecoderModel
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/encoder_decoder/__pycache__/__init__.cpython-310.pyc b/transformers/models/encoder_decoder/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 034c8a40aeb6c4599b636b9610547cff5ae147a3..0000000000000000000000000000000000000000
Binary files a/transformers/models/encoder_decoder/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encoder_decoder/__pycache__/configuration_encoder_decoder.cpython-310.pyc b/transformers/models/encoder_decoder/__pycache__/configuration_encoder_decoder.cpython-310.pyc
deleted file mode 100644
index b3ef7ead697cef63ddcbe609bf5797c336110a2d..0000000000000000000000000000000000000000
Binary files a/transformers/models/encoder_decoder/__pycache__/configuration_encoder_decoder.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encoder_decoder/__pycache__/modeling_encoder_decoder.cpython-310.pyc b/transformers/models/encoder_decoder/__pycache__/modeling_encoder_decoder.cpython-310.pyc
deleted file mode 100644
index ff431fb068cc015ac39f7302839dd05ca3872e74..0000000000000000000000000000000000000000
Binary files a/transformers/models/encoder_decoder/__pycache__/modeling_encoder_decoder.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encoder_decoder/__pycache__/modeling_flax_encoder_decoder.cpython-310.pyc b/transformers/models/encoder_decoder/__pycache__/modeling_flax_encoder_decoder.cpython-310.pyc
deleted file mode 100644
index 54e9212f91faa96590b3705be5d68e60b99c3928..0000000000000000000000000000000000000000
Binary files a/transformers/models/encoder_decoder/__pycache__/modeling_flax_encoder_decoder.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encoder_decoder/__pycache__/modeling_tf_encoder_decoder.cpython-310.pyc b/transformers/models/encoder_decoder/__pycache__/modeling_tf_encoder_decoder.cpython-310.pyc
deleted file mode 100644
index 328583073e5db0706c433f6f1e050f478b449158..0000000000000000000000000000000000000000
Binary files a/transformers/models/encoder_decoder/__pycache__/modeling_tf_encoder_decoder.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/encoder_decoder/configuration_encoder_decoder.py b/transformers/models/encoder_decoder/configuration_encoder_decoder.py
deleted file mode 100644
index 8c0ae2771e81f16ab1f7e82a69e91f2fa1ad5407..0000000000000000000000000000000000000000
--- a/transformers/models/encoder_decoder/configuration_encoder_decoder.py
+++ /dev/null
@@ -1,106 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class EncoderDecoderConfig(PretrainedConfig):
- r"""
- [`EncoderDecoderConfig`] is the configuration class to store the configuration of a [`EncoderDecoderModel`]. It is
- used to instantiate an Encoder Decoder model according to the specified arguments, defining the encoder and decoder
- configs.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- kwargs (*optional*):
- Dictionary of keyword arguments. Notably:
-
- - **encoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that defines
- the encoder config.
- - **decoder** ([`PretrainedConfig`], *optional*) -- An instance of a configuration object that defines
- the decoder config.
-
- Examples:
-
- ```python
- >>> from transformers import BertConfig, EncoderDecoderConfig, EncoderDecoderModel
-
- >>> # Initializing a BERT google-bert/bert-base-uncased style configuration
- >>> config_encoder = BertConfig()
- >>> config_decoder = BertConfig()
-
- >>> config = EncoderDecoderConfig.from_encoder_decoder_configs(config_encoder, config_decoder)
-
- >>> # Initializing a Bert2Bert model (with random weights) from the google-bert/bert-base-uncased style configurations
- >>> model = EncoderDecoderModel(config=config)
-
- >>> # Accessing the model configuration
- >>> config_encoder = model.config.encoder
- >>> config_decoder = model.config.decoder
- >>> # set decoder config to causal lm
- >>> config_decoder.is_decoder = True
- >>> config_decoder.add_cross_attention = True
-
- >>> # Saving the model, including its configuration
- >>> model.save_pretrained("my-model")
-
- >>> # loading model and config from pretrained folder
- >>> encoder_decoder_config = EncoderDecoderConfig.from_pretrained("my-model")
- >>> model = EncoderDecoderModel.from_pretrained("my-model", config=encoder_decoder_config)
- ```"""
-
- model_type = "encoder-decoder"
- is_composition = True
-
- def __init__(self, **kwargs):
- super().__init__(**kwargs)
- assert (
- "encoder" in kwargs and "decoder" in kwargs
- ), "Config has to be initialized with encoder and decoder config"
- encoder_config = kwargs.pop("encoder")
- encoder_model_type = encoder_config.pop("model_type")
- decoder_config = kwargs.pop("decoder")
- decoder_model_type = decoder_config.pop("model_type")
-
- from ..auto.configuration_auto import AutoConfig
-
- self.encoder = AutoConfig.for_model(encoder_model_type, **encoder_config)
- self.decoder = AutoConfig.for_model(decoder_model_type, **decoder_config)
- self.is_encoder_decoder = True
-
- @classmethod
- def from_encoder_decoder_configs(
- cls, encoder_config: PretrainedConfig, decoder_config: PretrainedConfig, **kwargs
- ) -> PretrainedConfig:
- r"""
- Instantiate a [`EncoderDecoderConfig`] (or a derived class) from a pre-trained encoder model configuration and
- decoder model configuration.
-
- Returns:
- [`EncoderDecoderConfig`]: An instance of a configuration object
- """
- logger.info("Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config")
- decoder_config.is_decoder = True
- decoder_config.add_cross_attention = True
-
- return cls(encoder=encoder_config.to_dict(), decoder=decoder_config.to_dict(), **kwargs)
diff --git a/transformers/models/encoder_decoder/modeling_encoder_decoder.py b/transformers/models/encoder_decoder/modeling_encoder_decoder.py
deleted file mode 100644
index 16248fee64ce593a6d68d309259115abb083aba2..0000000000000000000000000000000000000000
--- a/transformers/models/encoder_decoder/modeling_encoder_decoder.py
+++ /dev/null
@@ -1,693 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Classes to support Encoder-Decoder architectures"""
-
-
-import gc
-import inspect
-import os
-import tempfile
-import warnings
-from typing import Optional, Tuple, Union
-
-import torch
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from ...configuration_utils import PretrainedConfig
-from ...modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
-from ...modeling_utils import PreTrainedModel
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from ..auto.configuration_auto import AutoConfig
-from ..auto.modeling_auto import AutoModel, AutoModelForCausalLM
-from .configuration_encoder_decoder import EncoderDecoderConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "EncoderDecoderConfig"
-
-DEPRECATION_WARNING = (
- "Version v4.12.0 introduces a better way to train encoder-decoder models by computing the loss inside the"
- " encoder-decoder framework rather than in the decoder itself. You may observe training discrepancies if"
- " fine-tuning a model trained with versions anterior to 4.12.0. The decoder_input_ids are now created based on the"
- " labels, no need to pass them yourself anymore."
-)
-
-ENCODER_DECODER_START_DOCSTRING = r"""
- This class can be used to initialize a sequence-to-sequence model with any pretrained autoencoding model as the
- encoder and any pretrained autoregressive model as the decoder. The encoder is loaded via
- [`~AutoModel.from_pretrained`] function and the decoder is loaded via [`~AutoModelForCausalLM.from_pretrained`]
- function. Cross-attention layers are automatically added to the decoder and should be fine-tuned on a downstream
- generative task, like summarization.
-
- The effectiveness of initializing sequence-to-sequence models with pretrained checkpoints for sequence generation
- tasks was shown in [Leveraging Pre-trained Checkpoints for Sequence Generation
- Tasks](https://arxiv.org/abs/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. Michael Matena, Yanqi
- Zhou, Wei Li, Peter J. Liu.
-
- After such an Encoder Decoder model has been trained/fine-tuned, it can be saved/loaded just like any other models
- (see the examples for more information).
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`EncoderDecoderConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ENCODER_DECODER_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
- Indices of decoder input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
-
- If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
- `past_key_values`).
-
- For training, `decoder_input_ids` are automatically created by the model by shifting the `labels` to the
- right, replacing -100 by the `pad_token_id` and prepending them with the `decoder_start_token_id`.
- decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
- Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
- be used by default.
- encoder_outputs (`tuple(torch.FloatTensor)`, *optional*):
- This tuple must consist of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
- `last_hidden_state` (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`) is a tensor
- of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the
- decoder.
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded
- representation. This is useful if you want more control over how to convert `decoder_input_ids` indices
- into associated vectors than the model's internal embedding lookup matrix.
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss for the decoder. Indices should be in `[-100, 0,
- ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- If set to `True`, the model will return a [`~utils.Seq2SeqLMOutput`] instead of a plain tuple.
- kwargs (*optional*): Remaining dictionary of keyword arguments. Keyword arguments come in two flavors:
-
- - Without a prefix which will be input as `**encoder_kwargs` for the encoder forward function.
- - With a *decoder_* prefix which will be input as `**decoder_kwargs` for the decoder forward function.
-"""
-
-
-def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
- """
- Shift input ids one token to the right.
- """
- shifted_input_ids = input_ids.new_zeros(input_ids.shape)
- shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
- if decoder_start_token_id is None:
- raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
- shifted_input_ids[:, 0] = decoder_start_token_id
-
- if pad_token_id is None:
- raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
- # replace possible -100 values in labels by `pad_token_id`
- shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
-
- return shifted_input_ids
-
-
-@add_start_docstrings(ENCODER_DECODER_START_DOCSTRING)
-class EncoderDecoderModel(PreTrainedModel):
- r"""
- [`EncoderDecoderModel`] is a generic model class that will be instantiated as a transformer architecture with one
- of the base model classes of the library as encoder and another one as decoder when created with the
- :meth*~transformers.AutoModel.from_pretrained* class method for the encoder and
- :meth*~transformers.AutoModelForCausalLM.from_pretrained* class method for the decoder.
- """
-
- config_class = EncoderDecoderConfig
- base_model_prefix = "encoder_decoder"
- main_input_name = "input_ids"
- supports_gradient_checkpointing = True
-
- def __init__(
- self,
- config: Optional[PretrainedConfig] = None,
- encoder: Optional[PreTrainedModel] = None,
- decoder: Optional[PreTrainedModel] = None,
- ):
- if config is None and (encoder is None or decoder is None):
- raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")
- if config is None:
- config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config)
- else:
- if not isinstance(config, self.config_class):
- raise ValueError(f"Config: {config} has to be of type {self.config_class}")
-
- if config.decoder.cross_attention_hidden_size is not None:
- if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:
- raise ValueError(
- "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"
- f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
- f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"
- " `config.encoder.hidden_size`."
- )
-
- # initialize with config
- super().__init__(config)
-
- if encoder is None:
- from ..auto.modeling_auto import AutoModel
-
- encoder = AutoModel.from_config(config.encoder)
-
- if decoder is None:
- from ..auto.modeling_auto import AutoModelForCausalLM
-
- decoder = AutoModelForCausalLM.from_config(config.decoder)
-
- self.encoder = encoder
- self.decoder = decoder
-
- if self.encoder.config.to_dict() != self.config.encoder.to_dict():
- logger.warning(
- f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config:"
- f" {self.config.encoder}"
- )
- if self.decoder.config.to_dict() != self.config.decoder.to_dict():
- logger.warning(
- f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"
- f" {self.config.decoder}"
- )
-
- # make sure that the individual model's config refers to the shared config
- # so that the updates to the config will be synced
- self.encoder.config = self.config.encoder
- self.decoder.config = self.config.decoder
-
- # encoder outputs might need to be projected to different dimension for decoder
- if (
- self.encoder.config.hidden_size != self.decoder.config.hidden_size
- and self.decoder.config.cross_attention_hidden_size is None
- ):
- self.enc_to_dec_proj = nn.Linear(self.encoder.config.hidden_size, self.decoder.config.hidden_size)
-
- if self.encoder.get_output_embeddings() is not None:
- raise ValueError(
- f"The encoder {self.encoder} should not have a LM Head. Please use a model without LM Head"
- )
-
- decoder_signature = set(inspect.signature(self.decoder.forward).parameters.keys())
- if "encoder_hidden_states" not in decoder_signature:
- raise ValueError(
- "The selected decoder is not prepared for the encoder hidden states to be passed. Please see the "
- "following discussion on GitHub: https://github.com/huggingface/transformers/issues/23350"
- )
-
- # tie encoder, decoder weights if config set accordingly
- self.tie_weights()
-
- def tie_weights(self):
- # tie encoder & decoder if needed
- if self.config.tie_encoder_decoder:
- # tie encoder and decoder base model
- decoder_base_model_prefix = self.decoder.base_model_prefix
- tied_weights = self._tie_encoder_decoder_weights(
- self.encoder,
- self.decoder._modules[decoder_base_model_prefix],
- self.decoder.base_model_prefix,
- "encoder",
- )
- # Setting a dynamic variable instead of `_tied_weights_keys` because it's a class
- # attributed not an instance member, therefore modifying it will modify the entire class
- # Leading to issues on subsequent calls by different tests or subsequent calls.
- self._dynamic_tied_weights_keys = tied_weights
-
- def get_encoder(self):
- return self.encoder
-
- def get_decoder(self):
- return self.decoder
-
- def get_input_embeddings(self):
- return self.encoder.get_input_embeddings()
-
- def get_output_embeddings(self):
- return self.decoder.get_output_embeddings()
-
- def set_output_embeddings(self, new_embeddings):
- return self.decoder.set_output_embeddings(new_embeddings)
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
- r"""
- Example:
-
- ```python
- >>> from transformers import EncoderDecoderModel
-
- >>> model = EncoderDecoderModel.from_pretrained("patrickvonplaten/bert2bert-cnn_dailymail-fp16")
- ```"""
-
- from_tf = kwargs.pop("from_tf", False)
- if from_tf:
- from transformers import TFEncoderDecoderModel
-
- # a workaround to load from tensorflow checkpoint
- # Using `_tf_model` won't work, because the weight names in the encoder/decoder of `_tf_model` get
- # extended before saving those components. For example, The name of `_tf_model.encoder.vit` is
- # `[top model name]/encoder/vit`, but the name of `tf_model.encoder.vit` is `[top model name]/vit`. The
- # [top model name] is handled (stripped) by the conversion method, and the former case gets extra `encoder`,
- # which should not occur when we want to save the components alone.
- # There was a (very) ugly potential fix, which wasn't integrated to `transformers`: see
- # https://github.com/huggingface/transformers/pull/13222/commits/dbb3c9de76eee235791d2064094654637c99f36d#r697304245
- # (the change in `src/transformers/modeling_tf_utils.py`)
- _tf_model = TFEncoderDecoderModel.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
- config = _tf_model.config
-
- # Using `tf_model` instead
- encoder = _tf_model.encoder.__class__(_tf_model.config.encoder)
- decoder = _tf_model.decoder.__class__(_tf_model.config.decoder)
- # Make sure models are built
- encoder(encoder.dummy_inputs)
- decoder(decoder.dummy_inputs)
-
- # Get the variable correspondence between `_tf_model` and `encoder` and `decoder`
- encoder_variables = {}
- for v in encoder.trainable_variables + encoder.non_trainable_variables:
- encoder_variables["/".join(v.name.split("/")[1:])] = v
- decoder_variables = {}
- for v in decoder.trainable_variables + decoder.non_trainable_variables:
- decoder_variables["/".join(v.name.split("/")[1:])] = v
-
- _encoder_variables = {}
- for v in _tf_model.encoder.trainable_variables + _tf_model.encoder.non_trainable_variables:
- _encoder_variables["/".join(v.name.split("/")[2:])] = v
- _decoder_variables = {}
- for v in _tf_model.decoder.trainable_variables + _tf_model.decoder.non_trainable_variables:
- _decoder_variables["/".join(v.name.split("/")[2:])] = v
-
- # assign weight values to `encoder` and `decoder` from `_tf_model`
- for name, v in encoder_variables.items():
- v.assign(_encoder_variables[name])
- for name, v in decoder_variables.items():
- v.assign(_decoder_variables[name])
-
- tf_model = TFEncoderDecoderModel(encoder=encoder, decoder=decoder)
-
- # Deal with `enc_to_dec_proj`
- if hasattr(_tf_model, "enc_to_dec_proj"):
- tf_model(tf_model.dummy_inputs)
- tf_model.enc_to_dec_proj.kernel.assign(_tf_model.enc_to_dec_proj.kernel)
- tf_model.enc_to_dec_proj.bias.assign(_tf_model.enc_to_dec_proj.bias)
-
- with tempfile.TemporaryDirectory() as tmpdirname:
- encoder_dir = os.path.join(tmpdirname, "encoder")
- decoder_dir = os.path.join(tmpdirname, "decoder")
- tf_model.encoder.save_pretrained(encoder_dir)
- tf_model.decoder.save_pretrained(decoder_dir)
-
- if hasattr(tf_model, "enc_to_dec_proj"):
- enc_to_dec_proj_weight = torch.transpose(
- torch.from_numpy(tf_model.enc_to_dec_proj.kernel.numpy()), 1, 0
- )
- enc_to_dec_proj_bias = torch.from_numpy(tf_model.enc_to_dec_proj.bias.numpy())
-
- del _tf_model
- del tf_model
- gc.collect()
-
- model = EncoderDecoderModel.from_encoder_decoder_pretrained(
- encoder_dir, decoder_dir, encoder_from_tf=True, decoder_from_tf=True
- )
- # This is only for copying some specific attributes of this particular model.
- model.config = config
-
- if hasattr(model, "enc_to_dec_proj"):
- model.enc_to_dec_proj.weight.data = enc_to_dec_proj_weight.contiguous()
- model.enc_to_dec_proj.bias.data = enc_to_dec_proj_bias.contiguous()
-
- return model
-
- # At the moment fast initialization is not supported for composite models
- if kwargs.get("_fast_init", False):
- logger.warning(
- "Fast initialization is currently not supported for EncoderDecoderModel. "
- "Falling back to slow initialization..."
- )
- kwargs["_fast_init"] = False
-
- return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
-
- @classmethod
- def from_encoder_decoder_pretrained(
- cls,
- encoder_pretrained_model_name_or_path: str = None,
- decoder_pretrained_model_name_or_path: str = None,
- *model_args,
- **kwargs,
- ) -> PreTrainedModel:
- r"""
- Instantiate an encoder and a decoder from one or two base classes of the library from pretrained model
- checkpoints.
-
-
- The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
- the model, you need to first set it back in training mode with `model.train()`.
-
- Params:
- encoder_pretrained_model_name_or_path (`str`, *optional*):
- Information necessary to initiate the encoder. Can be either:
-
- - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
- - A path to a *directory* containing model weights saved using
- [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
- - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
- this case, `from_tf` should be set to `True` and a configuration object should be provided as
- `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
- PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
-
- decoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
- Information necessary to initiate the decoder. Can be either:
-
- - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
- - A path to a *directory* containing model weights saved using
- [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
- - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
- this case, `from_tf` should be set to `True` and a configuration object should be provided as
- `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
- PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
-
- model_args (remaining positional arguments, *optional*):
- All remaining positional arguments will be passed to the underlying model's `__init__` method.
-
- kwargs (remaining dictionary of keyword arguments, *optional*):
- Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
- `output_attentions=True`).
-
- - To update the encoder configuration, use the prefix *encoder_* for each configuration parameter.
- - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.
- - To update the parent model configuration, do not use a prefix for each configuration parameter.
-
- Behaves differently depending on whether a `config` is provided or automatically loaded.
-
- Example:
-
- ```python
- >>> from transformers import EncoderDecoderModel
-
- >>> # initialize a bert2bert from two pretrained BERT models. Note that the cross-attention layers will be randomly initialized
- >>> model = EncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-uncased", "google-bert/bert-base-uncased")
- >>> # saving model after fine-tuning
- >>> model.save_pretrained("./bert2bert")
- >>> # load fine-tuned model
- >>> model = EncoderDecoderModel.from_pretrained("./bert2bert")
- ```"""
-
- kwargs_encoder = {
- argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
- }
-
- kwargs_decoder = {
- argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
- }
-
- # remove encoder, decoder kwargs from kwargs
- for key in kwargs_encoder.keys():
- del kwargs["encoder_" + key]
- for key in kwargs_decoder.keys():
- del kwargs["decoder_" + key]
-
- # Load and initialize the encoder and decoder
- # The distinction between encoder and decoder at the model level is made
- # by the value of the flag `is_decoder` that we need to set correctly.
- encoder = kwargs_encoder.pop("model", None)
- if encoder is None:
- if encoder_pretrained_model_name_or_path is None:
- raise ValueError(
- "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
- "to be defined."
- )
-
- if "config" not in kwargs_encoder:
- encoder_config, kwargs_encoder = AutoConfig.from_pretrained(
- encoder_pretrained_model_name_or_path, **kwargs_encoder, return_unused_kwargs=True
- )
-
- if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
- logger.info(
- f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
- "from a decoder model. Cross-attention and casual mask are disabled."
- )
- encoder_config.is_decoder = False
- encoder_config.add_cross_attention = False
-
- kwargs_encoder["config"] = encoder_config
-
- encoder = AutoModel.from_pretrained(encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder)
-
- decoder = kwargs_decoder.pop("model", None)
- if decoder is None:
- if decoder_pretrained_model_name_or_path is None:
- raise ValueError(
- "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
- "to be defined."
- )
-
- if "config" not in kwargs_decoder:
- decoder_config, kwargs_decoder = AutoConfig.from_pretrained(
- decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True
- )
-
- if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
- logger.info(
- f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
- f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
- f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
- )
- decoder_config.is_decoder = True
- decoder_config.add_cross_attention = True
-
- kwargs_decoder["config"] = decoder_config
-
- if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
- logger.warning(
- f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
- f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
- "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
- "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "
- "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"
- )
-
- decoder = AutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
-
- # instantiate config with corresponding kwargs
- config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
- return cls(encoder=encoder, decoder=decoder, config=config)
-
- @add_start_docstrings_to_model_forward(ENCODER_DECODER_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- decoder_input_ids: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.BoolTensor] = None,
- encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
- past_key_values: Tuple[Tuple[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- **kwargs,
- ) -> Union[Tuple, Seq2SeqLMOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import EncoderDecoderModel, BertTokenizer
- >>> import torch
-
- >>> tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")
- >>> model = EncoderDecoderModel.from_encoder_decoder_pretrained(
- ... "google-bert/bert-base-uncased", "google-bert/bert-base-uncased"
- ... ) # initialize Bert2Bert from pre-trained checkpoints
-
- >>> # training
- >>> model.config.decoder_start_token_id = tokenizer.cls_token_id
- >>> model.config.pad_token_id = tokenizer.pad_token_id
- >>> model.config.vocab_size = model.config.decoder.vocab_size
-
- >>> input_ids = tokenizer("This is a really long text", return_tensors="pt").input_ids
- >>> labels = tokenizer("This is the corresponding summary", return_tensors="pt").input_ids
- >>> outputs = model(input_ids=input_ids, labels=labels)
- >>> loss, logits = outputs.loss, outputs.logits
-
- >>> # save and load from pretrained
- >>> model.save_pretrained("bert2bert")
- >>> model = EncoderDecoderModel.from_pretrained("bert2bert")
-
- >>> # generation
- >>> generated = model.generate(input_ids)
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}
-
- kwargs_decoder = {
- argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
- }
-
- if encoder_outputs is None:
- encoder_outputs = self.encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- **kwargs_encoder,
- )
- elif isinstance(encoder_outputs, tuple):
- encoder_outputs = BaseModelOutput(*encoder_outputs)
-
- encoder_hidden_states = encoder_outputs[0]
-
- # optionally project encoder_hidden_states
- if (
- self.encoder.config.hidden_size != self.decoder.config.hidden_size
- and self.decoder.config.cross_attention_hidden_size is None
- ):
- encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
-
- if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
- decoder_input_ids = shift_tokens_right(
- labels, self.config.pad_token_id, self.config.decoder_start_token_id
- )
- if decoder_attention_mask is None:
- decoder_attention_mask = decoder_input_ids.new_tensor(decoder_input_ids != self.config.pad_token_id)
-
- # Decode
- decoder_outputs = self.decoder(
- input_ids=decoder_input_ids,
- attention_mask=decoder_attention_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=attention_mask,
- inputs_embeds=decoder_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- use_cache=use_cache,
- past_key_values=past_key_values,
- return_dict=return_dict,
- **kwargs_decoder,
- )
-
- # Compute loss independent from decoder (as some shift the logits inside them)
- loss = None
- if labels is not None:
- warnings.warn(DEPRECATION_WARNING, FutureWarning)
- logits = decoder_outputs.logits if return_dict else decoder_outputs[0]
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- if loss is not None:
- return (loss,) + decoder_outputs + encoder_outputs
- else:
- return decoder_outputs + encoder_outputs
-
- return Seq2SeqLMOutput(
- loss=loss,
- logits=decoder_outputs.logits,
- past_key_values=decoder_outputs.past_key_values,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- )
-
- def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
- return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
-
- def prepare_inputs_for_generation(
- self, input_ids, past_key_values=None, attention_mask=None, use_cache=None, encoder_outputs=None, **kwargs
- ):
- decoder_inputs = self.decoder.prepare_inputs_for_generation(input_ids, past_key_values=past_key_values)
- decoder_attention_mask = decoder_inputs["attention_mask"] if "attention_mask" in decoder_inputs else None
- input_dict = {
- "attention_mask": attention_mask,
- "decoder_attention_mask": decoder_attention_mask,
- "decoder_input_ids": decoder_inputs["input_ids"],
- "encoder_outputs": encoder_outputs,
- "past_key_values": decoder_inputs["past_key_values"],
- "use_cache": use_cache,
- }
- return input_dict
-
- def resize_token_embeddings(self, *args, **kwargs):
- raise NotImplementedError(
- "Resizing the embedding layers via the EncoderDecoderModel directly is not supported. Please use the"
- " respective methods of the wrapped objects (model.encoder.resize_token_embeddings(...) or"
- " model.decoder.resize_token_embeddings(...))"
- )
-
- def _reorder_cache(self, past_key_values, beam_idx):
- # apply decoder cache reordering here
- return self.decoder._reorder_cache(past_key_values, beam_idx)
diff --git a/transformers/models/encoder_decoder/modeling_flax_encoder_decoder.py b/transformers/models/encoder_decoder/modeling_flax_encoder_decoder.py
deleted file mode 100644
index beecd080328e167a56e68b943c23822fa78d36b8..0000000000000000000000000000000000000000
--- a/transformers/models/encoder_decoder/modeling_flax_encoder_decoder.py
+++ /dev/null
@@ -1,899 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Classes to support Flax Encoder-Decoder architectures"""
-
-
-import os
-from typing import Optional, Tuple, Union
-
-import flax.linen as nn
-import jax
-import jax.numpy as jnp
-from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
-from flax.traverse_util import flatten_dict, unflatten_dict
-from jax import lax
-from jax.random import PRNGKey
-
-from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutputWithCrossAttentions, FlaxSeq2SeqLMOutput
-from ...modeling_flax_utils import FlaxPreTrainedModel
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from ..auto.configuration_auto import AutoConfig
-from ..auto.modeling_flax_auto import FlaxAutoModel, FlaxAutoModelForCausalLM
-from .configuration_encoder_decoder import EncoderDecoderConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "EncoderDecoderConfig"
-
-ENCODER_DECODER_START_DOCSTRING = r"""
- This class can be used to initialize a sequence-to-sequence model with any pretrained autoencoding model as the
- encoder and any pretrained autoregressive model as the decoder. The encoder is loaded via
- [`~AutoModel.from_pretrained`] function and the decoder is loaded via [`~AutoModelForCausalLM.from_pretrained`]
- function. Cross-attention layers are automatically added to the decoder and should be fine-tuned on a downstream
- generative task, like summarization.
-
- The effectiveness of initializing sequence-to-sequence models with pretrained checkpoints for sequence generation
- tasks was shown in [Leveraging Pre-trained Checkpoints for Sequence Generation
- Tasks](https://arxiv.org/abs/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. Michael Matena, Yanqi
- Zhou, Wei Li, Peter J. Liu.
-
- After such an Encoder Decoder model has been trained/fine-tuned, it can be saved/loaded just like any other models
- (see the examples for more information).
-
- This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a Flax Linen
- [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a
- regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.
-
- Parameters:
- config ([`EncoderDecoderConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.
- dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
- The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
- `jax.numpy.bfloat16` (on TPUs).
-
- This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
- specified all the computation will be performed with the given `dtype`.
-
- **Note that this only specifies the dtype of the computation and does not influence the dtype of model
- parameters.**
-
- If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and
- [`~FlaxPreTrainedModel.to_bf16`].
-"""
-
-ENCODER_DECODER_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
- it.
-
- Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
- Indices of decoder input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are decoder input IDs?](../glossary#decoder-input-ids)
-
- For sequence to sequence training, `decoder_input_ids` should be provided. `decoder_input_ids` should be
- created outside of the model by shifting the `labels` to the right, replacing -100 by the `pad_token_id`
- and prepending them with the `decoder_start_token_id`.
- decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
- Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
- be used by default.
- position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.encoder.max_position_embeddings - 1]`.
- decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
- range `[0, config.decoder.max_position_embeddings - 1]`.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- If set to `True`, the model will return a [`~utils.FlaxSeq2SeqLMOutput`] instead of a plain tuple.
-"""
-
-ENCODER_DECODER_ENCODE_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
- it.
-
- Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.encoder.max_position_embeddings - 1]`.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- If set to `True`, the model will return a [`~utils.FlaxBaseModelOutput`] instead of a plain tuple.
-"""
-
-ENCODER_DECODER_DECODE_INPUTS_DOCSTRING = r"""
- Args:
- decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
- Indices of decoder input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are decoder input IDs?](../glossary#decoder-input-ids)
-
- If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
- `past_key_values`).
-
- For sequence to sequence training, `decoder_input_ids` should be provided. `decoder_input_ids` should be
- created outside of the model by shifting the `labels` to the right, replacing -100 by the `pad_token_id`
- and prepending them with the `decoder_start_token_id`.
- encoder_outputs (`tuple(tuple(jnp.ndarray)`):
- Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
- `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of
- hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
- encoder_attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
- Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
- be used by default.
- decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
- range `[0, config.decoder.max_position_embeddings - 1]`.
- past_key_values (`Dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):
- Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
- auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- If set to `True`, the model will return a [`~utils.FlaxCausalLMOutputWithCrossAttentions`] instead of a
- plain tuple.
-"""
-
-
-class FlaxEncoderDecoderModule(nn.Module):
- config: EncoderDecoderConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- encoder_config = self.config.encoder
- decoder_config = self.config.decoder
-
- # Copied from `modeling_hybrid_clip.py` with modifications.
- from ...models.auto.modeling_flax_auto import FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, FLAX_MODEL_MAPPING
-
- encoder_module = FLAX_MODEL_MAPPING[encoder_config.__class__].module_class
- decoder_module = FLAX_MODEL_FOR_CAUSAL_LM_MAPPING[decoder_config.__class__].module_class
-
- self.encoder = encoder_module(encoder_config, dtype=self.dtype)
- self.decoder = decoder_module(decoder_config, dtype=self.dtype)
-
- # encoder outputs might need to be projected to different dimension for decoder
- if (
- self.encoder.config.hidden_size != self.decoder.config.hidden_size
- and self.decoder.config.cross_attention_hidden_size is None
- ):
- self.enc_to_dec_proj = nn.Dense(
- self.decoder.config.hidden_size,
- kernel_init=jax.nn.initializers.normal(self.decoder.config.initializer_range),
- dtype=self.dtype,
- )
- else:
- self.enc_to_dec_proj = None
-
- def _get_encoder_module(self):
- return self.encoder
-
- def _get_projection_module(self):
- return self.enc_to_dec_proj
-
- def _get_decoder_module(self):
- return self.decoder
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- decoder_input_ids,
- decoder_attention_mask,
- position_ids,
- decoder_position_ids,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- deterministic: bool = True,
- ):
- encoder_outputs = self.encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- deterministic=deterministic,
- )
-
- encoder_hidden_states = encoder_outputs[0]
-
- # optionally project encoder_hidden_states
- if self.enc_to_dec_proj is not None:
- encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
-
- decoder_outputs = self.decoder(
- input_ids=decoder_input_ids,
- attention_mask=decoder_attention_mask,
- position_ids=decoder_position_ids,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- deterministic=deterministic,
- )
-
- if not return_dict:
- return decoder_outputs + encoder_outputs
-
- return FlaxSeq2SeqLMOutput(
- logits=decoder_outputs.logits,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(ENCODER_DECODER_START_DOCSTRING)
-class FlaxEncoderDecoderModel(FlaxPreTrainedModel):
- r"""
- [`FlaxEncoderDecoderModel`] is a generic model class that will be instantiated as a transformer architecture with
- the module (flax.nn.Module) of one of the base model classes of the library as encoder module and another one as
- decoder module when created with the :meth*~transformers.FlaxAutoModel.from_pretrained* class method for the
- encoder and :meth*~transformers.FlaxAutoModelForCausalLM.from_pretrained* class method for the decoder.
- """
-
- config_class = EncoderDecoderConfig
- base_model_prefix = "encoder_decoder"
- module_class = FlaxEncoderDecoderModule
-
- def __init__(
- self,
- config: EncoderDecoderConfig,
- input_shape: Optional[Tuple] = None,
- seed: int = 0,
- dtype: jnp.dtype = jnp.float32,
- _do_init: bool = True,
- **kwargs,
- ):
- if input_shape is None:
- input_shape = ((1, 1), (1, 1))
-
- if not _do_init:
- raise ValueError(
- "`FlaxEncoderDecoderModel` cannot be created without initializing, `_do_init` must be `True`."
- )
-
- if config.decoder.cross_attention_hidden_size is not None:
- if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:
- raise ValueError(
- "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"
- f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
- f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"
- " `config.encoder.hidden_size`."
- )
-
- module = self.module_class(config=config, dtype=dtype, **kwargs)
- super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
-
- def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
- encoder_input_shape, decoder_input_shape = input_shape
-
- # init input tensors
- input_ids = jnp.zeros(encoder_input_shape, dtype="i4")
- attention_mask = jnp.ones_like(input_ids)
- decoder_input_ids = jnp.zeros(decoder_input_shape, dtype="i4")
- decoder_attention_mask = jnp.ones_like(decoder_input_ids)
-
- batch_size, sequence_length = input_ids.shape
- position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
-
- decoder_batch_size, decoder_sequence_length = decoder_input_ids.shape
- if not decoder_batch_size == batch_size:
- raise ValueError(
- f"The inputs of encoder and decoder should have the same batch size, but got {batch_size} for encoder"
- f" and {decoder_batch_size} for decoder."
- )
- decoder_position_ids = jnp.broadcast_to(
- jnp.arange(decoder_sequence_length)[None, :], (decoder_batch_size, decoder_sequence_length)
- )
-
- params_rng, dropout_rng = jax.random.split(rng)
- rngs = {"params": params_rng, "dropout": dropout_rng}
-
- random_params = self.module.init(
- rngs,
- input_ids,
- attention_mask,
- decoder_input_ids,
- decoder_attention_mask,
- position_ids,
- decoder_position_ids,
- )["params"]
-
- if params is not None:
- random_params = flatten_dict(unfreeze(random_params))
- params = flatten_dict(unfreeze(params))
- for missing_key in self._missing_keys:
- params[missing_key] = random_params[missing_key]
- self._missing_keys = set()
- return freeze(unflatten_dict(params))
- else:
- return random_params
-
- def init_cache(self, batch_size, max_length, encoder_outputs):
- r"""
- Args:
- batch_size (`int`):
- batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
- max_length (`int`):
- maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
- cache.
- encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):
- `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:
- `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)
- is a sequence of hidden-states at the output of the last layer of the encoder. Used in the
- cross-attention of the decoder.
- """
- # init input variables to retrieve cache
- decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")
- decoder_attention_mask = jnp.ones_like(decoder_input_ids)
- decoder_position_ids = jnp.broadcast_to(
- jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape
- )
-
- def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
- decoder_module = module._get_decoder_module()
- return decoder_module(
- input_ids=decoder_input_ids,
- attention_mask=decoder_attention_mask,
- position_ids=decoder_position_ids,
- **kwargs,
- )
-
- init_variables = self.module.init(
- jax.random.PRNGKey(0),
- decoder_input_ids=decoder_input_ids,
- decoder_attention_mask=decoder_attention_mask,
- decoder_position_ids=decoder_position_ids,
- encoder_hidden_states=encoder_outputs[0],
- init_cache=True,
- method=_decoder_forward, # we only need to call the decoder to init the cache
- )
- return unfreeze(init_variables["cache"])
-
- @add_start_docstrings(ENCODER_DECODER_ENCODE_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=_CONFIG_FOR_DOC)
- def encode(
- self,
- input_ids: jnp.ndarray,
- attention_mask: Optional[jnp.ndarray] = None,
- position_ids: Optional[jnp.ndarray] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- train: bool = False,
- params: dict = None,
- dropout_rng: PRNGKey = None,
- ):
- r"""
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import FlaxEncoderDecoderModel, BertTokenizer
-
- >>> # initialize a bert2gpt2 from pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized
- >>> model = FlaxEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-cased", "openai-community/gpt2")
-
- >>> tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-cased")
-
- >>> text = "My friends are cool but they eat too many carbs."
- >>> input_ids = tokenizer.encode(text, return_tensors="np")
- >>> encoder_outputs = model.encode(input_ids)
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- if attention_mask is None:
- attention_mask = jnp.ones_like(input_ids)
- if position_ids is None:
- batch_size, sequence_length = input_ids.shape
- position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
-
- # Handle any PRNG if needed
- rngs = {}
- if dropout_rng is not None:
- rngs["dropout"] = dropout_rng
-
- def _encoder_forward(module, input_ids, attention_mask, position_ids, **kwargs):
- encode_module = module._get_encoder_module()
- return encode_module(input_ids, attention_mask, position_ids, **kwargs)
-
- outputs = self.module.apply(
- {"params": params or self.params},
- input_ids=jnp.array(input_ids, dtype="i4"),
- attention_mask=jnp.array(attention_mask, dtype="i4"),
- position_ids=jnp.array(position_ids, dtype="i4"),
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- deterministic=not train,
- rngs=rngs,
- method=_encoder_forward,
- )
-
- if return_dict:
- outputs = FlaxBaseModelOutput(
- last_hidden_state=outputs.last_hidden_state,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- return outputs
-
- @add_start_docstrings(ENCODER_DECODER_DECODE_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=FlaxCausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
- def decode(
- self,
- decoder_input_ids,
- encoder_outputs,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- decoder_attention_mask: Optional[jnp.ndarray] = None,
- decoder_position_ids: Optional[jnp.ndarray] = None,
- past_key_values: dict = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- train: bool = False,
- params: dict = None,
- dropout_rng: PRNGKey = None,
- ):
- r"""
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import FlaxEncoderDecoderModel, BertTokenizer
- >>> import jax.numpy as jnp
-
- >>> # initialize a bert2gpt2 from pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized
- >>> model = FlaxEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-cased", "openai-community/gpt2")
-
- >>> tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-cased")
-
- >>> text = "My friends are cool but they eat too many carbs."
- >>> input_ids = tokenizer.encode(text, max_length=1024, return_tensors="np")
- >>> encoder_outputs = model.encode(input_ids)
-
- >>> decoder_start_token_id = model.config.decoder.bos_token_id
- >>> decoder_input_ids = jnp.ones((input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
-
- >>> outputs = model.decode(decoder_input_ids, encoder_outputs)
- >>> logits = outputs.logits
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- encoder_hidden_states = encoder_outputs[0]
- if encoder_attention_mask is None:
- batch_size, sequence_length = encoder_hidden_states.shape[:2]
- encoder_attention_mask = jnp.ones((batch_size, sequence_length))
-
- batch_size, sequence_length = decoder_input_ids.shape
- if decoder_attention_mask is None:
- decoder_attention_mask = jnp.ones((batch_size, sequence_length))
-
- if decoder_position_ids is None:
- if past_key_values is not None:
- raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")
-
- decoder_position_ids = jnp.broadcast_to(
- jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
- )
-
- # Handle any PRNG if needed
- rngs = {}
- if dropout_rng is not None:
- rngs["dropout"] = dropout_rng
-
- inputs = {"params": params or self.params}
-
- # if past_key_values are passed then cache is already initialized a private flag init_cache has to be
- # passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that
- # it can be changed by FlaxBartAttention module
- if past_key_values:
- inputs["cache"] = past_key_values
- mutable = ["cache"]
- else:
- mutable = False
-
- def _decoder_forward(
- module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, encoder_hidden_states, **kwargs
- ):
- projection_module = module._get_projection_module()
- decoder_module = module._get_decoder_module()
-
- # optionally project encoder_hidden_states
- if projection_module is not None:
- encoder_hidden_states = projection_module(encoder_hidden_states)
-
- return decoder_module(
- decoder_input_ids,
- decoder_attention_mask,
- decoder_position_ids,
- encoder_hidden_states=encoder_hidden_states,
- **kwargs,
- )
-
- outputs = self.module.apply(
- inputs,
- decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
- decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
- decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- deterministic=not train,
- rngs=rngs,
- mutable=mutable,
- method=_decoder_forward,
- )
-
- # add updated cache to model output
- if past_key_values is not None and return_dict:
- outputs, past = outputs
- outputs["past_key_values"] = unfreeze(past["cache"])
- return outputs
- elif past_key_values is not None and not return_dict:
- outputs, past = outputs
- outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]
-
- return outputs
-
- @add_start_docstrings_to_model_forward(ENCODER_DECODER_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=FlaxSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
- def __call__(
- self,
- input_ids: jnp.ndarray,
- attention_mask: Optional[jnp.ndarray] = None,
- decoder_input_ids: Optional[jnp.ndarray] = None,
- decoder_attention_mask: Optional[jnp.ndarray] = None,
- position_ids: Optional[jnp.ndarray] = None,
- decoder_position_ids: Optional[jnp.ndarray] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- train: bool = False,
- params: dict = None,
- dropout_rng: PRNGKey = None,
- ):
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import FlaxEncoderDecoderModel, BertTokenizer, GPT2Tokenizer
-
- >>> # load a fine-tuned bert2gpt2 model
- >>> model = FlaxEncoderDecoderModel.from_pretrained("patrickvonplaten/bert2gpt2-cnn_dailymail-fp16")
- >>> # load input & output tokenizer
- >>> tokenizer_input = BertTokenizer.from_pretrained("google-bert/bert-base-cased")
- >>> tokenizer_output = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
-
- >>> article = '''Sigma Alpha Epsilon is under fire for a video showing party-bound fraternity members
- >>> singing a racist chant. SAE's national chapter suspended the students,
- >>> but University of Oklahoma President David Boren took it a step further,
- >>> saying the university's affiliation with the fraternity is permanently done.'''
-
- >>> input_ids = tokenizer_input(article, add_special_tokens=True, return_tensors="np").input_ids
-
- >>> # use GPT2's eos_token as the pad as well as eos token
- >>> model.config.eos_token_id = model.config.decoder.eos_token_id
- >>> model.config.pad_token_id = model.config.eos_token_id
-
- >>> sequences = model.generate(input_ids, num_beams=4, max_length=12).sequences
-
- >>> summary = tokenizer_output.batch_decode(sequences, skip_special_tokens=True)[0]
- >>> assert summary == "SAS Alpha Epsilon suspended Sigma Alpha Epsilon members"
- ```
- """
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- # prepare encoder inputs
- if attention_mask is None:
- attention_mask = jnp.ones_like(input_ids)
- if position_ids is None:
- batch_size, sequence_length = input_ids.shape
- position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
-
- # prepare decoder inputs
- if decoder_input_ids is None:
- raise ValueError(
- "`decoder_input_ids` cannot be `None`. For sequence to sequence training, `decoder_position_ids` must"
- " be specified as an input argument."
- )
- if decoder_attention_mask is None:
- decoder_attention_mask = jnp.ones_like(decoder_input_ids)
- if decoder_position_ids is None:
- batch_size, sequence_length = decoder_input_ids.shape
- decoder_position_ids = jnp.broadcast_to(
- jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
- )
-
- # Handle any PRNG if needed
- rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}
-
- return self.module.apply(
- {"params": params or self.params},
- input_ids=jnp.array(input_ids, dtype="i4"),
- attention_mask=jnp.array(attention_mask, dtype="i4"),
- decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
- decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
- position_ids=jnp.array(position_ids, dtype="i4"),
- decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- deterministic=not train,
- rngs=rngs,
- )
-
- def prepare_inputs_for_generation(
- self,
- decoder_input_ids,
- max_length,
- attention_mask: Optional[jax.Array] = None,
- decoder_attention_mask: Optional[jax.Array] = None,
- encoder_outputs=None,
- **kwargs,
- ):
- # initializing the cache
- batch_size, seq_length = decoder_input_ids.shape
-
- past_key_values = self.init_cache(batch_size, max_length, encoder_outputs)
- # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
- # But since the decoder uses a causal mask, those positions are masked anyways.
- # Thus we can create a single static attention_mask here, which is more efficient for compilation
- extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
- if decoder_attention_mask is not None:
- decoder_position_ids = decoder_attention_mask.cumsum(axis=-1) - 1
- extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, decoder_attention_mask, (0, 0))
- else:
- decoder_position_ids = jnp.broadcast_to(
- jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length)
- )
-
- return {
- "past_key_values": past_key_values,
- "encoder_outputs": encoder_outputs,
- "encoder_attention_mask": attention_mask,
- "decoder_attention_mask": extended_attention_mask,
- "decoder_position_ids": decoder_position_ids,
- }
-
- def update_inputs_for_generation(self, model_outputs, model_kwargs):
- model_kwargs["past_key_values"] = model_outputs.past_key_values
- model_kwargs["decoder_position_ids"] = model_kwargs["decoder_position_ids"][:, -1:] + 1
- return model_kwargs
-
- @classmethod
- def from_encoder_decoder_pretrained(
- cls,
- encoder_pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
- decoder_pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
- *model_args,
- **kwargs,
- ) -> FlaxPreTrainedModel:
- r"""
- Instantiate an encoder and a decoder from one or two base classes of the library from pretrained model
- checkpoints.
-
- Params:
- encoder_pretrained_model_name_or_path (`Union[str, os.PathLike]`, *optional*):
- Information necessary to initiate the encoder. Can be either:
-
- - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
- - A path to a *directory* containing model weights saved using
- [`~FlaxPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
-
- decoder_pretrained_model_name_or_path (`Union[str, os.PathLike]`, *optional*, defaults to `None`):
- Information necessary to initiate the decoder. Can be either:
-
- - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
- - A path to a *directory* containing model weights saved using
- [`~FlaxPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
-
- model_args (remaining positional arguments, *optional*):
- All remaning positional arguments will be passed to the underlying model's `__init__` method.
-
- kwargs (remaining dictionary of keyword arguments, *optional*):
- Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
- `output_attentions=True`).
-
- - To update the encoder configuration, use the prefix *encoder_* for each configuration parameter.
- - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.
- - To update the parent model configuration, do not use a prefix for each configuration parameter.
-
- Behaves differently depending on whether a `config` is provided or automatically loaded.
-
- Example:
-
- ```python
- >>> from transformers import FlaxEncoderDecoderModel
-
- >>> # initialize a bert2gpt2 from pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized
- >>> model = FlaxEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-cased", "openai-community/gpt2")
- >>> # saving model after fine-tuning
- >>> model.save_pretrained("./bert2gpt2")
- >>> # load fine-tuned model
- >>> model = FlaxEncoderDecoderModel.from_pretrained("./bert2gpt2")
- ```"""
-
- kwargs_encoder = {
- argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
- }
-
- kwargs_decoder = {
- argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
- }
-
- # remove encoder, decoder kwargs from kwargs
- for key in kwargs_encoder.keys():
- del kwargs["encoder_" + key]
- for key in kwargs_decoder.keys():
- del kwargs["decoder_" + key]
-
- # Load and initialize the encoder and decoder
- # The distinction between encoder and decoder at the model level is made
- # by the value of the flag `is_decoder` that we need to set correctly.
- encoder = kwargs_encoder.pop("model", None)
- if encoder is None:
- if encoder_pretrained_model_name_or_path is None:
- raise ValueError(
- "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
- "to be defined."
- )
-
- if "config" not in kwargs_encoder:
- encoder_config, kwargs_encoder = AutoConfig.from_pretrained(
- encoder_pretrained_model_name_or_path, **kwargs_encoder, return_unused_kwargs=True
- )
- if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
- logger.info(
- f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
- "from a decoder model. Cross-attention and casual mask are disabled."
- )
- encoder_config.is_decoder = False
- encoder_config.add_cross_attention = False
-
- kwargs_encoder["config"] = encoder_config
-
- encoder = FlaxAutoModel.from_pretrained(
- encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder
- )
-
- decoder = kwargs_decoder.pop("model", None)
- if decoder is None:
- if decoder_pretrained_model_name_or_path is None:
- raise ValueError(
- "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
- "to be defined."
- )
-
- if "config" not in kwargs_decoder:
- decoder_config, kwargs_decoder = AutoConfig.from_pretrained(
- decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True
- )
- if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
- logger.info(
- f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
- f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
- f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
- )
- decoder_config.is_decoder = True
- decoder_config.add_cross_attention = True
-
- kwargs_decoder["config"] = decoder_config
-
- if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
- logger.warning(
- f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
- f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
- "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
- "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "
- "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"
- )
-
- decoder = FlaxAutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
-
- # instantiate config with corresponding kwargs
- dtype = kwargs.pop("dtype", jnp.float32)
- config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
-
- # init model
- model = cls(config, dtype=dtype)
- model.params["encoder"] = encoder.params
- model.params["decoder"] = decoder.params
-
- return model
diff --git a/transformers/models/encoder_decoder/modeling_tf_encoder_decoder.py b/transformers/models/encoder_decoder/modeling_tf_encoder_decoder.py
deleted file mode 100644
index 855fb767d13d73173365034f485b43174ed583d9..0000000000000000000000000000000000000000
--- a/transformers/models/encoder_decoder/modeling_tf_encoder_decoder.py
+++ /dev/null
@@ -1,663 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Classes to support TF Encoder-Decoder architectures"""
-
-
-from __future__ import annotations
-
-import inspect
-import re
-import warnings
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...configuration_utils import PretrainedConfig
-from ...modeling_tf_outputs import TFBaseModelOutput, TFSeq2SeqLMOutput
-from ...modeling_tf_utils import (
- TFCausalLanguageModelingLoss,
- TFModelInputType,
- TFPreTrainedModel,
- get_initializer,
- keras,
- unpack_inputs,
-)
-from ...tf_utils import shape_list
-from ...utils import (
- ModelOutput,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from ..auto.configuration_auto import AutoConfig
-from ..auto.modeling_tf_auto import TFAutoModel, TFAutoModelForCausalLM
-from .configuration_encoder_decoder import EncoderDecoderConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "EncoderDecoderConfig"
-
-DEPRECATION_WARNING = (
- "Version v4.17.0 introduces a better way to train encoder-decoder models by computing the loss inside the"
- " encoder-decoder framework rather than in the decoder itself. You may observe training discrepancies if"
- " fine-tuning a model trained with versions anterior to 4.17.0. The decoder_input_ids are now created based on the"
- " labels, no need to pass them yourself anymore."
-)
-
-ENCODER_DECODER_START_DOCSTRING = r"""
- This class can be used to initialize a sequence-to-sequence model with any pretrained autoencoding model as the
- encoder and any pretrained autoregressive model as the decoder. The encoder is loaded via
- [`~TFAutoModel.from_pretrained`] function and the decoder is loaded via [`~TFAutoModelForCausalLM.from_pretrained`]
- function. Cross-attention layers are automatically added to the decoder and should be fine-tuned on a downstream
- generative task, like summarization.
-
- The effectiveness of initializing sequence-to-sequence models with pretrained checkpoints for sequence generation
- tasks was shown in [Leveraging Pre-trained Checkpoints for Sequence Generation
- Tasks](https://arxiv.org/abs/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. Michael Matena, Yanqi
- Zhou, Wei Li, Peter J. Liu.
-
- After such an Encoder Decoder model has been trained/fine-tuned, it can be saved/loaded just like any other models
- (see the examples for more information).
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`EncoderDecoderConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ENCODER_DECODER_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`np.ndarray`, `tf.Tensor`, `List[tf.Tensor]` ``Dict[str, tf.Tensor]` or `Dict[str, np.ndarray]` and each example must have the shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- decoder_input_ids (`np.ndarray` or `tf.Tensor` of shape `(batch_size, target_sequence_length)`, *optional*):
- Indices of decoder input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
-
- If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
- `past_key_values`).
-
- Provide for sequence to sequence training to the decoder. Indices can be obtained using
- [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for
- details.
- decoder_attention_mask (`np.ndarray` or `tf.Tensor` of shape `(batch_size, target_sequence_length)`, *optional*):
- Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
- be used by default.
- encoder_outputs (`tuple(tuple(tf.Tensor)`, *optional*):
- This tuple must consist of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
- `last_hidden_state` (`tf.Tensor` of shape `({0}, hidden_size)`) is a tensor of hidden-states at the output
- of the last layer of the encoder. Used in the cross-attention of the decoder.
- past_key_values (`tuple(tuple(tf.Tensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `({0})`.
- inputs_embeds (`np.ndarray` or `tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- decoder_inputs_embeds (`np.ndarray` or `tf.Tensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded
- representation. This is useful if you want more control over how to convert `decoder_input_ids` indices
- into associated vectors than the model's internal embedding lookup matrix.
- labels (`np.ndarray` or `tf.Tensor` of shape `({0})`, *optional*):
- Labels for computing the masked language modeling loss for the decoder. Indices should be in `[-100, 0,
- ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- If set to `True`, the model will return a [`~utils.Seq2SeqLMOutput`] instead of a plain tuple.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
- kwargs (*optional*): Remaining dictionary of keyword arguments. Keyword arguments come in two flavors:
-
- - Without a prefix which will be input as `**encoder_kwargs` for the encoder forward function.
- - With a *decoder_* prefix which will be input as `**decoder_kwargs`` for the decoder forward function.
-"""
-
-
-def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, decoder_start_token_id: int):
- if pad_token_id is None:
- raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
- pad_token_id = tf.cast(pad_token_id, input_ids.dtype)
-
- if decoder_start_token_id is None:
- raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
- decoder_start_token_id = tf.cast(decoder_start_token_id, input_ids.dtype)
-
- start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id)
- shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)
- # replace possible -100 values in labels by `pad_token_id`
- shifted_input_ids = tf.where(
- shifted_input_ids == -100, tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids
- )
-
- # "Verify that `labels` has only positive values and -100"
- assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=input_ids.dtype))
-
- # Make sure the assertion op is called by wrapping the result in an identity no-op
- with tf.control_dependencies([assert_gte0]):
- shifted_input_ids = tf.identity(shifted_input_ids)
-
- return shifted_input_ids
-
-
-@add_start_docstrings(ENCODER_DECODER_START_DOCSTRING)
-class TFEncoderDecoderModel(TFPreTrainedModel, TFCausalLanguageModelingLoss):
- r"""
- [`TFEncoderDecoderModel`] is a generic model class that will be instantiated as a transformer architecture with one
- of the base model classes of the library as encoder and another one as decoder when created with the
- [`~TFAutoModel.from_pretrained`] class method for the encoder and [`~TFAutoModelForCausalLM.from_pretrained`] class
- method for the decoder.
- """
-
- config_class = EncoderDecoderConfig
- base_model_prefix = "encoder_decoder"
- load_weight_prefix = "tf_encoder_decoder_model"
-
- def __init__(
- self,
- config: Optional[PretrainedConfig] = None,
- encoder: Optional[TFPreTrainedModel] = None,
- decoder: Optional[TFPreTrainedModel] = None,
- ):
- if config is None and (encoder is None or decoder is None):
- raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")
- if config is None:
- config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config)
- else:
- if not isinstance(config, self.config_class):
- raise ValueError(f"config: {config} has to be of type {self.config_class}")
-
- if config.decoder.cross_attention_hidden_size is not None:
- if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:
- raise ValueError(
- "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"
- f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
- f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"
- " `config.encoder.hidden_size`."
- )
-
- # initialize with config
- super().__init__(config)
-
- if encoder is None:
- encoder = TFAutoModel.from_config(config.encoder, name="encoder")
-
- if decoder is None:
- decoder = TFAutoModelForCausalLM.from_config(config.decoder, name="decoder")
-
- self.encoder = encoder
- self.decoder = decoder
-
- if self.encoder.config.to_dict() != self.config.encoder.to_dict():
- logger.warning(
- f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config:"
- f" {self.config.encoder}"
- )
- if self.decoder.config.to_dict() != self.config.decoder.to_dict():
- logger.warning(
- f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"
- f" {self.config.decoder}"
- )
-
- # make sure that the individual model's config refers to the shared config
- # so that the updates to the config will be synced
- self.encoder.config = self.config.encoder
- self.decoder.config = self.config.decoder
-
- # encoder outputs might need to be projected to different dimension for decoder
- if (
- self.encoder.config.hidden_size != self.decoder.config.hidden_size
- and self.decoder.config.cross_attention_hidden_size is None
- ):
- self.enc_to_dec_proj = keras.layers.Dense(
- units=self.decoder.config.hidden_size,
- kernel_initializer=get_initializer(config.encoder.initializer_range),
- name="enc_to_dec_proj",
- )
-
- if self.encoder.get_output_embeddings() is not None:
- raise ValueError(
- f"The encoder {self.encoder} should not have a LM Head. Please use a model without LM Head"
- )
-
- decoder_signature = set(inspect.signature(self.decoder.call).parameters.keys())
- if "encoder_hidden_states" not in decoder_signature:
- raise ValueError(
- "The selected decoder is not prepared for the encoder hidden states to be passed. Please see the "
- "following discussion on GitHub: https://github.com/huggingface/transformers/issues/23350"
- )
-
- def get_encoder(self):
- return self.encoder
-
- def get_decoder(self):
- return self.decoder
-
- def get_input_embeddings(self):
- return self.encoder.get_input_embeddings()
-
- def get_output_embeddings(self):
- return self.decoder.get_output_embeddings()
-
- def set_output_embeddings(self, new_embeddings):
- return self.decoder.set_output_embeddings(new_embeddings)
-
- def tf_to_pt_weight_rename(self, tf_weight):
- # Matt: The TF and PT weights don't align because our TF base classes have an extra layer compared to PT models
- # (the main model stem is in the MainLayer class). If we remove that layer, then weight names sync up as normal.
- # However, the name of that extra layer is the name of the MainLayer in the base model. We make the assumption
- # here that the config model_type is the same as the name of the MainLayer. I don't know of anywhere that's
- # not the case, and I wasn't sure how else to go from the config to the correct MainLayer name!
-
- # This override is only needed in the case where we're crossloading weights from PT. However, since weights are
- # often safetensors now, we don't know if we're going to be crossloading until we sniff the weights file.
- # Therefore, we specify tf_to_pt_weight_rename anyway, and let the super method figure out if it needs it
- # or not.
- encoder_model_type = self.config.encoder.model_type
- if "encoder" in tf_weight and "decoder" not in tf_weight:
- return (re.sub(rf"encoder\.{encoder_model_type}\.", "encoder.", tf_weight),)
- else:
- return (tf_weight,)
-
- @classmethod
- def from_encoder_decoder_pretrained(
- cls,
- encoder_pretrained_model_name_or_path: str = None,
- decoder_pretrained_model_name_or_path: str = None,
- *model_args,
- **kwargs,
- ) -> TFPreTrainedModel:
- r"""
- Instantiate an encoder and a decoder from one or two base classes of the library from pretrained model
- checkpoints.
-
-
- Params:
- encoder_pretrained_model_name_or_path (`str`, *optional*):
- Information necessary to initiate the encoder. Can be either:
-
- - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
- - A path to a *directory* containing model weights saved using
- [`~TFPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
- - A path or url to a *pytorch index checkpoint file* (e.g, `./pt_model/`). In this case,
- `encoder_from_pt` should be set to `True`.
-
- decoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
- Information necessary to initiate the decoder. Can be either:
-
- - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
- - A path to a *directory* containing model weights saved using
- [`~TFPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
- - A path or url to a *pytorch checkpoint file* (e.g, `./pt_model/`). In this case,
- `decoder_from_pt` should be set to `True`.
-
- model_args (remaining positional arguments, *optional*):
- All remaning positional arguments will be passed to the underlying model's `__init__` method.
-
- kwargs (remaining dictionary of keyword arguments, *optional*):
- Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
- `output_attentions=True`).
-
- - To update the encoder configuration, use the prefix *encoder_* for each configuration parameter.
- - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.
- - To update the parent model configuration, do not use a prefix for each configuration parameter.
-
- Behaves differently depending on whether a `config` is provided or automatically loaded.
-
- Example:
-
- ```python
- >>> from transformers import TFEncoderDecoderModel
-
- >>> # initialize a bert2gpt2 from two pretrained BERT models. Note that the cross-attention layers will be randomly initialized
- >>> model = TFEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-uncased", "openai-community/gpt2")
- >>> # saving model after fine-tuning
- >>> model.save_pretrained("./bert2gpt2")
- >>> # load fine-tuned model
- >>> model = TFEncoderDecoderModel.from_pretrained("./bert2gpt2")
- ```"""
-
- kwargs_encoder = {
- argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
- }
-
- kwargs_decoder = {
- argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
- }
-
- # remove encoder, decoder kwargs from kwargs
- for key in kwargs_encoder.keys():
- del kwargs["encoder_" + key]
- for key in kwargs_decoder.keys():
- del kwargs["decoder_" + key]
-
- # Load and initialize the encoder and decoder
- # The distinction between encoder and decoder at the model level is made
- # by the value of the flag `is_decoder` that we need to set correctly.
- encoder = kwargs_encoder.pop("model", None)
- if encoder is None:
- if encoder_pretrained_model_name_or_path is None:
- raise ValueError(
- "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
- "to be defined."
- )
-
- if "config" not in kwargs_encoder:
- encoder_config = AutoConfig.from_pretrained(encoder_pretrained_model_name_or_path)
- if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
- logger.info(
- f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
- "from a decoder model. Cross-attention and casual mask are disabled."
- )
- encoder_config.is_decoder = False
- encoder_config.add_cross_attention = False
-
- kwargs_encoder["config"] = encoder_config
-
- kwargs_encoder["name"] = "encoder"
- kwargs_encoder["load_weight_prefix"] = cls.load_weight_prefix
- encoder = TFAutoModel.from_pretrained(encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder)
-
- decoder = kwargs_decoder.pop("model", None)
- if decoder is None:
- if decoder_pretrained_model_name_or_path is None:
- raise ValueError(
- "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
- "to be defined."
- )
-
- if "config" not in kwargs_decoder:
- decoder_config = AutoConfig.from_pretrained(decoder_pretrained_model_name_or_path)
- if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
- logger.info(
- f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
- f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
- f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
- )
- decoder_config.is_decoder = True
- decoder_config.add_cross_attention = True
-
- kwargs_decoder["config"] = decoder_config
-
- if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
- logger.warning(
- f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
- f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
- "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
- "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "
- "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"
- )
-
- kwargs_decoder["name"] = "decoder"
- kwargs_decoder["load_weight_prefix"] = cls.load_weight_prefix
- decoder = TFAutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
-
- # Make sure these 2 `keras.Model` have fixed names so `from_pretrained` could load model weights correctly.
- if encoder.name != "encoder":
- raise ValueError("encoder model must be created with the name `encoder`.")
- if decoder.name != "decoder":
- raise ValueError("decoder model must be created with the name `decoder`.")
-
- # instantiate config with corresponding kwargs
- config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
- return cls(encoder=encoder, decoder=decoder, config=config)
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ENCODER_DECODER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- decoder_input_ids: np.ndarray | tf.Tensor | None = None,
- decoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- encoder_outputs: np.ndarray | tf.Tensor | None = None,
- past_key_values: Tuple[Tuple[tf.Tensor]] | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- decoder_inputs_embeds: np.ndarray | tf.Tensor | None = None,
- labels: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- **kwargs,
- ) -> Union[TFSeq2SeqLMOutput, Tuple[tf.Tensor]]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import TFEncoderDecoderModel, BertTokenizer
-
- >>> # initialize a bert2gpt2 from a pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized
- >>> model = TFEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-cased", "openai-community/gpt2")
-
- >>> tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-cased")
-
- >>> # forward
- >>> input_ids = tokenizer.encode(
- ... "Hello, my dog is cute", add_special_tokens=True, return_tensors="tf"
- ... ) # Batch size 1
- >>> outputs = model(input_ids=input_ids, decoder_input_ids=input_ids)
-
- >>> # training
- >>> outputs = model(input_ids=input_ids, decoder_input_ids=input_ids, labels=input_ids)
- >>> loss, logits = outputs.loss, outputs.logits
-
- >>> # save and load from pretrained
- >>> model.save_pretrained("bert2gpt2")
- >>> model = TFEncoderDecoderModel.from_pretrained("bert2gpt2")
-
- >>> # generation
- >>> generated = model.generate(input_ids, decoder_start_token_id=model.config.decoder.bos_token_id)
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}
-
- kwargs_decoder = {
- argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
- }
-
- # Let the user be responsible for the expected format.
- if encoder_outputs is not None:
- if return_dict and not isinstance(encoder_outputs, ModelOutput):
- raise ValueError(
- "If `return_dict=True` and `encoder_outputs` is provided, it should be an instance of "
- f"`ModelOutput`. Got an instance {type(encoder_outputs)} for `encoder_outputs`."
- )
-
- if encoder_outputs is None:
- encoder_inputs = {
- "input_ids": input_ids,
- "attention_mask": attention_mask,
- "inputs_embeds": inputs_embeds,
- "output_attentions": output_attentions,
- "output_hidden_states": output_hidden_states,
- "return_dict": return_dict,
- "training": training,
- }
-
- # Add arguments to encoder from `kwargs_encoder`
- encoder_inputs.update(kwargs_encoder)
-
- # Handle the case where the inputs are passed as a single dict which contains `labels`.
- # The `labels` shouldn't be passed to `self.encoder` below, because it is a based model without this
- # parameter (otherwise, an error occurs when `input_processing` is called inside `self.encoder.call()`).
- if "labels" in encoder_inputs:
- labels = encoder_inputs.pop("labels")
-
- # handle the init case where `dummy_inputs` returns a dict containing `decoder_input_ids`.
- if "decoder_input_ids" in encoder_inputs:
- decoder_input_ids = encoder_inputs.pop("decoder_input_ids")
- # handle the init case where `dummy_inputs` returns a dict containing `decoder_input_ids`.
- if "decoder_attention_mask" in encoder_inputs:
- decoder_attention_mask = encoder_inputs.pop("decoder_attention_mask")
-
- encoder_outputs = self.encoder(**encoder_inputs)
-
- encoder_hidden_states = encoder_outputs[0]
-
- # optionally project encoder_hidden_states
- if (
- self.encoder.config.hidden_size != self.decoder.config.hidden_size
- and self.decoder.config.cross_attention_hidden_size is None
- ):
- encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
-
- if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
- decoder_input_ids = shift_tokens_right(
- labels, self.config.pad_token_id, self.config.decoder_start_token_id
- )
-
- decoder_inputs = {
- "input_ids": decoder_input_ids,
- "attention_mask": decoder_attention_mask,
- "encoder_hidden_states": encoder_hidden_states,
- "encoder_attention_mask": attention_mask,
- "inputs_embeds": decoder_inputs_embeds,
- "output_attentions": output_attentions,
- "output_hidden_states": output_hidden_states,
- "use_cache": use_cache,
- "past_key_values": past_key_values,
- "return_dict": return_dict,
- "training": training,
- }
-
- # Add arguments to decoder from `kwargs_decoder`
- decoder_inputs.update(kwargs_decoder)
-
- decoder_outputs = self.decoder(**decoder_inputs)
-
- logits = decoder_outputs[0]
-
- # Compute loss independent from decoder (as some shift the logits inside them)
- loss = None
- if labels is not None:
- warnings.warn(DEPRECATION_WARNING, FutureWarning)
- loss = self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- past_key_values = None
- if use_cache:
- past_key_values = decoder_outputs[1]
- # The starting index of the remaining elements in `decoder_outputs`
- start_index = sum([1 if x is not None else 0 for x in (loss, logits, past_key_values)])
-
- if not isinstance(encoder_outputs, tuple):
- encoder_outputs = encoder_outputs.to_tuple()
- output = (loss, logits, past_key_values) + decoder_outputs[start_index:] + encoder_outputs
- output = tuple([x for x in output if x is not None])
- return output
-
- return TFSeq2SeqLMOutput(
- loss=loss,
- logits=decoder_outputs.logits,
- past_key_values=decoder_outputs.past_key_values,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- )
-
- def prepare_inputs_for_generation(
- self, input_ids, past_key_values=None, attention_mask=None, use_cache=None, encoder_outputs=None, **kwargs
- ):
- decoder_inputs = self.decoder.prepare_inputs_for_generation(input_ids, past_key_values=past_key_values)
- decoder_attention_mask = decoder_inputs["attention_mask"] if "attention_mask" in decoder_inputs else None
- past_key_values = decoder_inputs.get("past_key_values")
- if past_key_values is None:
- past_key_values = decoder_inputs.get("past") # e.g. on TF GPT2
- input_dict = {
- "input_ids": None, # needs to be passed to make Keras.layer.__call__ happy
- "attention_mask": attention_mask,
- "decoder_attention_mask": decoder_attention_mask,
- "decoder_input_ids": decoder_inputs["input_ids"],
- # TODO (joao): the `TFBaseModelOutput` wrapper should not be needed after the generate refactor is complete
- "encoder_outputs": TFBaseModelOutput(last_hidden_state=encoder_outputs[0]),
- "past_key_values": past_key_values,
- "use_cache": use_cache,
- }
- return input_dict
-
- def prepare_decoder_input_ids_from_labels(self, labels: tf.Tensor):
- return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
-
- def resize_token_embeddings(self, *args, **kwargs):
- raise NotImplementedError(
- "Resizing the embedding layers via the TFEncoderDecoderModel directly is not supported.Please use the"
- " respective methods of the wrapped objects (model.encoder.resize_token_embeddings(...) or"
- " model.decoder.resize_token_embeddings(...))"
- )
-
- def _reorder_cache(self, past, beam_idx):
- # apply decoder cache reordering here
- return self.decoder._reorder_cache(past, beam_idx)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "enc_to_dec_proj", None) is not None:
- with tf.name_scope(self.enc_to_dec_proj.name):
- self.enc_to_dec_proj.build([None, None, self.encoder.config.hidden_size])
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "decoder", None) is not None:
- with tf.name_scope(self.decoder.name):
- self.decoder.build(None)
diff --git a/transformers/models/ernie/__init__.py b/transformers/models/ernie/__init__.py
deleted file mode 100644
index ea7f077f928d39527ab5cf9ba4f195a62445bb84..0000000000000000000000000000000000000000
--- a/transformers/models/ernie/__init__.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tensorflow_text_available, is_torch_available
-
-
-_import_structure = {
- "configuration_ernie": ["ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP", "ErnieConfig", "ErnieOnnxConfig"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_ernie"] = [
- "ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST",
- "ErnieForCausalLM",
- "ErnieForMaskedLM",
- "ErnieForMultipleChoice",
- "ErnieForNextSentencePrediction",
- "ErnieForPreTraining",
- "ErnieForQuestionAnswering",
- "ErnieForSequenceClassification",
- "ErnieForTokenClassification",
- "ErnieModel",
- "ErniePreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_ernie import ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieConfig, ErnieOnnxConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_ernie import (
- ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST,
- ErnieForCausalLM,
- ErnieForMaskedLM,
- ErnieForMultipleChoice,
- ErnieForNextSentencePrediction,
- ErnieForPreTraining,
- ErnieForQuestionAnswering,
- ErnieForSequenceClassification,
- ErnieForTokenClassification,
- ErnieModel,
- ErniePreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/ernie/__pycache__/__init__.cpython-310.pyc b/transformers/models/ernie/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 6f5821800eb5e5bd65c79f45acc33cb7a2ff39ad..0000000000000000000000000000000000000000
Binary files a/transformers/models/ernie/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ernie/__pycache__/configuration_ernie.cpython-310.pyc b/transformers/models/ernie/__pycache__/configuration_ernie.cpython-310.pyc
deleted file mode 100644
index eda536bceae9cdc3f1829500a9d4bd92fbe27148..0000000000000000000000000000000000000000
Binary files a/transformers/models/ernie/__pycache__/configuration_ernie.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ernie/__pycache__/modeling_ernie.cpython-310.pyc b/transformers/models/ernie/__pycache__/modeling_ernie.cpython-310.pyc
deleted file mode 100644
index 63c30ec1dc59994f72df4aace4ef2ad0b50ac73a..0000000000000000000000000000000000000000
Binary files a/transformers/models/ernie/__pycache__/modeling_ernie.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ernie/configuration_ernie.py b/transformers/models/ernie/configuration_ernie.py
deleted file mode 100644
index 81ed03596303ee4cf02bcfc97a914a54b5decda3..0000000000000000000000000000000000000000
--- a/transformers/models/ernie/configuration_ernie.py
+++ /dev/null
@@ -1,162 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The Google AI Language Team Authors and The HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" ERNIE model configuration"""
-from collections import OrderedDict
-from typing import Mapping
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class ErnieConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`ErnieModel`] or a [`TFErnieModel`]. It is used to
- instantiate a ERNIE model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the ERNIE
- [nghuyong/ernie-3.0-base-zh](https://huggingface.co/nghuyong/ernie-3.0-base-zh) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the ERNIE model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`ErnieModel`] or [`TFErnieModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 2):
- The vocabulary size of the `token_type_ids` passed when calling [`ErnieModel`] or [`TFErnieModel`].
- task_type_vocab_size (`int`, *optional*, defaults to 3):
- The vocabulary size of the `task_type_ids` for ERNIE2.0/ERNIE3.0 model
- use_task_id (`bool`, *optional*, defaults to `False`):
- Whether or not the model support `task_type_ids`
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- pad_token_id (`int`, *optional*, defaults to 0):
- Padding token id.
- position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
- Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
- positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
- [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
- For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
- with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- classifier_dropout (`float`, *optional*):
- The dropout ratio for the classification head.
-
- Examples:
-
- ```python
- >>> from transformers import ErnieConfig, ErnieModel
-
- >>> # Initializing a ERNIE nghuyong/ernie-3.0-base-zh style configuration
- >>> configuration = ErnieConfig()
-
- >>> # Initializing a model (with random weights) from the nghuyong/ernie-3.0-base-zh style configuration
- >>> model = ErnieModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "ernie"
-
- def __init__(
- self,
- vocab_size=30522,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=2,
- task_type_vocab_size=3,
- use_task_id=False,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- pad_token_id=0,
- position_embedding_type="absolute",
- use_cache=True,
- classifier_dropout=None,
- **kwargs,
- ):
- super().__init__(pad_token_id=pad_token_id, **kwargs)
-
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.hidden_act = hidden_act
- self.intermediate_size = intermediate_size
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.type_vocab_size = type_vocab_size
- self.task_type_vocab_size = task_type_vocab_size
- self.use_task_id = use_task_id
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.position_embedding_type = position_embedding_type
- self.use_cache = use_cache
- self.classifier_dropout = classifier_dropout
-
-
-class ErnieOnnxConfig(OnnxConfig):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- if self.task == "multiple-choice":
- dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
- else:
- dynamic_axis = {0: "batch", 1: "sequence"}
- return OrderedDict(
- [
- ("input_ids", dynamic_axis),
- ("attention_mask", dynamic_axis),
- ("token_type_ids", dynamic_axis),
- ("task_type_ids", dynamic_axis),
- ]
- )
diff --git a/transformers/models/ernie/modeling_ernie.py b/transformers/models/ernie/modeling_ernie.py
deleted file mode 100644
index a65f453205d5c599e29f0416cf76b585f26266f6..0000000000000000000000000000000000000000
--- a/transformers/models/ernie/modeling_ernie.py
+++ /dev/null
@@ -1,1820 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch ERNIE model."""
-
-
-import math
-import warnings
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutputWithPastAndCrossAttentions,
- BaseModelOutputWithPoolingAndCrossAttentions,
- CausalLMOutputWithCrossAttentions,
- MaskedLMOutput,
- MultipleChoiceModelOutput,
- NextSentencePredictorOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_ernie import ErnieConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "nghuyong/ernie-1.0-base-zh"
-_CONFIG_FOR_DOC = "ErnieConfig"
-
-
-from ..deprecated._archive_maps import ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class ErnieEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
- self.use_task_id = config.use_task_id
- if config.use_task_id:
- self.task_type_embeddings = nn.Embedding(config.task_type_vocab_size, config.hidden_size)
-
- # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
- # any TensorFlow checkpoint file
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
- self.register_buffer(
- "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
- )
-
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- task_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- past_key_values_length: int = 0,
- ) -> torch.Tensor:
- if input_ids is not None:
- input_shape = input_ids.size()
- else:
- input_shape = inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
-
- # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
- # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
- # issue #5664
- if token_type_ids is None:
- if hasattr(self, "token_type_ids"):
- buffered_token_type_ids = self.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
- embeddings = inputs_embeds + token_type_embeddings
- if self.position_embedding_type == "absolute":
- position_embeddings = self.position_embeddings(position_ids)
- embeddings += position_embeddings
-
- # add `task_type_id` for ERNIE model
- if self.use_task_id:
- if task_type_ids is None:
- task_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
- task_type_embeddings = self.task_type_embeddings(task_type_ids)
- embeddings += task_type_embeddings
-
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Ernie
-class ErnieSelfAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size)
- self.key = nn.Linear(config.hidden_size, self.all_head_size)
- self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
- self.position_embedding_type = position_embedding_type or getattr(
- config, "position_embedding_type", "absolute"
- )
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- self.max_position_embeddings = config.max_position_embeddings
- self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
-
- self.is_decoder = config.is_decoder
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- mixed_query_layer = self.query(hidden_states)
-
- # If this is instantiated as a cross-attention module, the keys
- # and values come from an encoder; the attention mask needs to be
- # such that the encoder's padding tokens are not attended to.
- is_cross_attention = encoder_hidden_states is not None
-
- if is_cross_attention and past_key_value is not None:
- # reuse k,v, cross_attentions
- key_layer = past_key_value[0]
- value_layer = past_key_value[1]
- attention_mask = encoder_attention_mask
- elif is_cross_attention:
- key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
- value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
- attention_mask = encoder_attention_mask
- elif past_key_value is not None:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
- value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
- else:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- use_cache = past_key_value is not None
- if self.is_decoder:
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_layer, value_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- query_length, key_length = query_layer.shape[2], key_layer.shape[2]
- if use_cache:
- position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(
- -1, 1
- )
- else:
- position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
- position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
- distance = position_ids_l - position_ids_r
-
- positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
- positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
-
- if self.position_embedding_type == "relative_key":
- relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores
- elif self.position_embedding_type == "relative_key_query":
- relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in ErnieModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- if self.is_decoder:
- outputs = outputs + (past_key_value,)
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->Ernie
-class ErnieSelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Ernie
-class ErnieAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- self.self = ErnieSelfAttention(config, position_embedding_type=position_embedding_type)
- self.output = ErnieSelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- self_outputs = self.self(
- hidden_states,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->Ernie
-class ErnieIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->Ernie
-class ErnieOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Ernie
-class ErnieLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = ErnieAttention(config)
- self.is_decoder = config.is_decoder
- self.add_cross_attention = config.add_cross_attention
- if self.add_cross_attention:
- if not self.is_decoder:
- raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
- self.crossattention = ErnieAttention(config, position_embedding_type="absolute")
- self.intermediate = ErnieIntermediate(config)
- self.output = ErnieOutput(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
- self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
- self_attention_outputs = self.attention(
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions=output_attentions,
- past_key_value=self_attn_past_key_value,
- )
- attention_output = self_attention_outputs[0]
-
- # if decoder, the last output is tuple of self-attn cache
- if self.is_decoder:
- outputs = self_attention_outputs[1:-1]
- present_key_value = self_attention_outputs[-1]
- else:
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- cross_attn_present_key_value = None
- if self.is_decoder and encoder_hidden_states is not None:
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
- " by setting `config.add_cross_attention=True`"
- )
-
- # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
- cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
- cross_attention_outputs = self.crossattention(
- attention_output,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- cross_attn_past_key_value,
- output_attentions,
- )
- attention_output = cross_attention_outputs[0]
- outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
-
- # add cross-attn cache to positions 3,4 of present_key_value tuple
- cross_attn_present_key_value = cross_attention_outputs[-1]
- present_key_value = present_key_value + cross_attn_present_key_value
-
- layer_output = apply_chunking_to_forward(
- self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
- )
- outputs = (layer_output,) + outputs
-
- # if decoder, return the attn key/values as the last output
- if self.is_decoder:
- outputs = outputs + (present_key_value,)
-
- return outputs
-
- def feed_forward_chunk(self, attention_output):
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.output(intermediate_output, attention_output)
- return layer_output
-
-
-# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Ernie
-class ErnieEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([ErnieLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = False,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- next_decoder_cache = () if use_cache else None
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
- past_key_value = past_key_values[i] if past_key_values is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
-
- hidden_states = layer_outputs[0]
- if use_cache:
- next_decoder_cache += (layer_outputs[-1],)
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- next_decoder_cache,
- all_hidden_states,
- all_self_attentions,
- all_cross_attentions,
- ]
- if v is not None
- )
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=next_decoder_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->Ernie
-class ErniePooler(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = nn.Tanh()
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(first_token_tensor)
- pooled_output = self.activation(pooled_output)
- return pooled_output
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPredictionHeadTransform with Bert->Ernie
-class ErniePredictionHeadTransform(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = ACT2FN[config.hidden_act]
- else:
- self.transform_act_fn = config.hidden_act
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->Ernie
-class ErnieLMPredictionHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.transform = ErniePredictionHeadTransform(config)
-
- # The output weights are the same as the input embeddings, but there is
- # an output-only bias for each token.
- self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
- self.bias = nn.Parameter(torch.zeros(config.vocab_size))
-
- # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
- self.decoder.bias = self.bias
-
- def forward(self, hidden_states):
- hidden_states = self.transform(hidden_states)
- hidden_states = self.decoder(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOnlyMLMHead with Bert->Ernie
-class ErnieOnlyMLMHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.predictions = ErnieLMPredictionHead(config)
-
- def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
- prediction_scores = self.predictions(sequence_output)
- return prediction_scores
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOnlyNSPHead with Bert->Ernie
-class ErnieOnlyNSPHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.seq_relationship = nn.Linear(config.hidden_size, 2)
-
- def forward(self, pooled_output):
- seq_relationship_score = self.seq_relationship(pooled_output)
- return seq_relationship_score
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPreTrainingHeads with Bert->Ernie
-class ErniePreTrainingHeads(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.predictions = ErnieLMPredictionHead(config)
- self.seq_relationship = nn.Linear(config.hidden_size, 2)
-
- def forward(self, sequence_output, pooled_output):
- prediction_scores = self.predictions(sequence_output)
- seq_relationship_score = self.seq_relationship(pooled_output)
- return prediction_scores, seq_relationship_score
-
-
-class ErniePreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ErnieConfig
- base_model_prefix = "ernie"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-@dataclass
-# Copied from transformers.models.bert.modeling_bert.BertForPreTrainingOutput with Bert->Ernie
-class ErnieForPreTrainingOutput(ModelOutput):
- """
- Output type of [`ErnieForPreTraining`].
-
- Args:
- loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
- Total loss as the sum of the masked language modeling loss and the next sequence prediction
- (classification) loss.
- prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
- seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`):
- Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
- before SoftMax).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- prediction_logits: torch.FloatTensor = None
- seq_relationship_logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-ERNIE_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`ErnieConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ERNIE_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- task_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Task type embedding is a special embedding to represent the characteristic of different tasks, such as
- word-aware pre-training task, structure-aware pre-training task and semantic-aware pre-training task. We
- assign a `task_type_id` to each task and the `task_type_id` is in the range `[0,
- config.task_type_vocab_size-1]
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Ernie Model transformer outputting raw hidden-states without any specific head on top.",
- ERNIE_START_DOCSTRING,
-)
-class ErnieModel(ErniePreTrainedModel):
- """
-
- The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
- cross-attention is added between the self-attention layers, following the architecture described in [Attention is
- all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
- Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
-
- To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
- to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
- `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
- """
-
- # Copied from transformers.models.bert.modeling_bert.BertModel.__init__ with Bert->Ernie
- def __init__(self, config, add_pooling_layer=True):
- super().__init__(config)
- self.config = config
-
- self.embeddings = ErnieEmbeddings(config)
- self.encoder = ErnieEncoder(config)
-
- self.pooler = ErniePooler(config) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- # Copied from transformers.models.bert.modeling_bert.BertModel.get_input_embeddings
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- # Copied from transformers.models.bert.modeling_bert.BertModel.set_input_embeddings
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
- # Copied from transformers.models.bert.modeling_bert.BertModel._prune_heads
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPoolingAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
- r"""
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if self.config.is_decoder:
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- else:
- use_cache = False
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- # past_key_values_length
- past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
-
- if attention_mask is None:
- attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
-
- if token_type_ids is None:
- if hasattr(self.embeddings, "token_type_ids"):
- buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
- # ourselves in which case we just need to make it broadcastable to all heads.
- extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
-
- # If a 2D or 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- if self.config.is_decoder and encoder_hidden_states is not None:
- encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
- encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
- if encoder_attention_mask is None:
- encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
- encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
- else:
- encoder_extended_attention_mask = None
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- inputs_embeds=inputs_embeds,
- past_key_values_length=past_key_values_length,
- )
- encoder_outputs = self.encoder(
- embedding_output,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_extended_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- return (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPoolingAndCrossAttentions(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- past_key_values=encoder_outputs.past_key_values,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- cross_attentions=encoder_outputs.cross_attentions,
- )
-
-
-@add_start_docstrings(
- """
- Ernie Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `next
- sentence prediction (classification)` head.
- """,
- ERNIE_START_DOCSTRING,
-)
-class ErnieForPreTraining(ErniePreTrainedModel):
- _tied_weights_keys = ["cls.predictions.decoder.bias", "cls.predictions.decoder.weight"]
-
- # Copied from transformers.models.bert.modeling_bert.BertForPreTraining.__init__ with Bert->Ernie,bert->ernie
- def __init__(self, config):
- super().__init__(config)
-
- self.ernie = ErnieModel(config)
- self.cls = ErniePreTrainingHeads(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- # Copied from transformers.models.bert.modeling_bert.BertForPreTraining.get_output_embeddings
- def get_output_embeddings(self):
- return self.cls.predictions.decoder
-
- # Copied from transformers.models.bert.modeling_bert.BertForPreTraining.set_output_embeddings
- def set_output_embeddings(self, new_embeddings):
- self.cls.predictions.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=ErnieForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- next_sentence_label: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], ErnieForPreTrainingOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked),
- the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the next sequence prediction (classification) loss. Input should be a sequence
- pair (see `input_ids` docstring) Indices should be in `[0, 1]`:
-
- - 0 indicates sequence B is a continuation of sequence A,
- - 1 indicates sequence B is a random sequence.
- kwargs (`Dict[str, any]`, optional, defaults to *{}*):
- Used to hide legacy arguments that have been deprecated.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, ErnieForPreTraining
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("nghuyong/ernie-1.0-base-zh")
- >>> model = ErnieForPreTraining.from_pretrained("nghuyong/ernie-1.0-base-zh")
-
- >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> prediction_logits = outputs.prediction_logits
- >>> seq_relationship_logits = outputs.seq_relationship_logits
- ```
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output, pooled_output = outputs[:2]
- prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
-
- total_loss = None
- if labels is not None and next_sentence_label is not None:
- loss_fct = CrossEntropyLoss()
- masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
- next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
- total_loss = masked_lm_loss + next_sentence_loss
-
- if not return_dict:
- output = (prediction_scores, seq_relationship_score) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return ErnieForPreTrainingOutput(
- loss=total_loss,
- prediction_logits=prediction_scores,
- seq_relationship_logits=seq_relationship_score,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """Ernie Model with a `language modeling` head on top for CLM fine-tuning.""", ERNIE_START_DOCSTRING
-)
-class ErnieForCausalLM(ErniePreTrainedModel):
- _tied_weights_keys = ["cls.predictions.decoder.bias", "cls.predictions.decoder.weight"]
-
- # Copied from transformers.models.bert.modeling_bert.BertLMHeadModel.__init__ with BertLMHeadModel->ErnieForCausalLM,Bert->Ernie,bert->ernie
- def __init__(self, config):
- super().__init__(config)
-
- if not config.is_decoder:
- logger.warning("If you want to use `ErnieForCausalLM` as a standalone, add `is_decoder=True.`")
-
- self.ernie = ErnieModel(config, add_pooling_layer=False)
- self.cls = ErnieOnlyMLMHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- # Copied from transformers.models.bert.modeling_bert.BertLMHeadModel.get_output_embeddings
- def get_output_embeddings(self):
- return self.cls.predictions.decoder
-
- # Copied from transformers.models.bert.modeling_bert.BertLMHeadModel.set_output_embeddings
- def set_output_embeddings(self, new_embeddings):
- self.cls.predictions.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutputWithCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.Tensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
- r"""
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
- `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
- ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- if labels is not None:
- use_cache = False
-
- outputs = self.ernie(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- prediction_scores = self.cls(sequence_output)
-
- lm_loss = None
- if labels is not None:
- # we are doing next-token prediction; shift prediction scores and input ids by one
- shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
- labels = labels[:, 1:].contiguous()
- loss_fct = CrossEntropyLoss()
- lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((lm_loss,) + output) if lm_loss is not None else output
-
- return CausalLMOutputWithCrossAttentions(
- loss=lm_loss,
- logits=prediction_scores,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- cross_attentions=outputs.cross_attentions,
- )
-
- # Copied from transformers.models.bert.modeling_bert.BertLMHeadModel.prepare_inputs_for_generation
- def prepare_inputs_for_generation(
- self, input_ids, past_key_values=None, attention_mask=None, use_cache=True, **model_kwargs
- ):
- input_shape = input_ids.shape
- # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
- if attention_mask is None:
- attention_mask = input_ids.new_ones(input_shape)
-
- # cut decoder_input_ids if past_key_values is used
- if past_key_values is not None:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
-
- return {
- "input_ids": input_ids,
- "attention_mask": attention_mask,
- "past_key_values": past_key_values,
- "use_cache": use_cache,
- }
-
- # Copied from transformers.models.bert.modeling_bert.BertLMHeadModel._reorder_cache
- def _reorder_cache(self, past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
- )
- return reordered_past
-
-
-@add_start_docstrings("""Ernie Model with a `language modeling` head on top.""", ERNIE_START_DOCSTRING)
-class ErnieForMaskedLM(ErniePreTrainedModel):
- _tied_weights_keys = ["cls.predictions.decoder.bias", "cls.predictions.decoder.weight"]
-
- # Copied from transformers.models.bert.modeling_bert.BertForMaskedLM.__init__ with Bert->Ernie,bert->ernie
- def __init__(self, config):
- super().__init__(config)
-
- if config.is_decoder:
- logger.warning(
- "If you want to use `ErnieForMaskedLM` make sure `config.is_decoder=False` for "
- "bi-directional self-attention."
- )
-
- self.ernie = ErnieModel(config, add_pooling_layer=False)
- self.cls = ErnieOnlyMLMHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- # Copied from transformers.models.bert.modeling_bert.BertForMaskedLM.get_output_embeddings
- def get_output_embeddings(self):
- return self.cls.predictions.decoder
-
- # Copied from transformers.models.bert.modeling_bert.BertForMaskedLM.set_output_embeddings
- def set_output_embeddings(self, new_embeddings):
- self.cls.predictions.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output="'paris'",
- expected_loss=0.88,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- prediction_scores = self.cls(sequence_output)
-
- masked_lm_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss() # -100 index = padding token
- masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return MaskedLMOutput(
- loss=masked_lm_loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- # Copied from transformers.models.bert.modeling_bert.BertForMaskedLM.prepare_inputs_for_generation
- def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):
- input_shape = input_ids.shape
- effective_batch_size = input_shape[0]
-
- # add a dummy token
- if self.config.pad_token_id is None:
- raise ValueError("The PAD token should be defined for generation")
-
- attention_mask = torch.cat([attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], dim=-1)
- dummy_token = torch.full(
- (effective_batch_size, 1), self.config.pad_token_id, dtype=torch.long, device=input_ids.device
- )
- input_ids = torch.cat([input_ids, dummy_token], dim=1)
-
- return {"input_ids": input_ids, "attention_mask": attention_mask}
-
-
-@add_start_docstrings(
- """Ernie Model with a `next sentence prediction (classification)` head on top.""",
- ERNIE_START_DOCSTRING,
-)
-class ErnieForNextSentencePrediction(ErniePreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForNextSentencePrediction.__init__ with Bert->Ernie,bert->ernie
- def __init__(self, config):
- super().__init__(config)
-
- self.ernie = ErnieModel(config)
- self.cls = ErnieOnlyNSPHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=NextSentencePredictorOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- **kwargs,
- ) -> Union[Tuple[torch.Tensor], NextSentencePredictorOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
- (see `input_ids` docstring). Indices should be in `[0, 1]`:
-
- - 0 indicates sequence B is a continuation of sequence A,
- - 1 indicates sequence B is a random sequence.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, ErnieForNextSentencePrediction
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("nghuyong/ernie-1.0-base-zh")
- >>> model = ErnieForNextSentencePrediction.from_pretrained("nghuyong/ernie-1.0-base-zh")
-
- >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
- >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
- >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
-
- >>> outputs = model(**encoding, labels=torch.LongTensor([1]))
- >>> logits = outputs.logits
- >>> assert logits[0, 0] < logits[0, 1] # next sentence was random
- ```
- """
-
- if "next_sentence_label" in kwargs:
- warnings.warn(
- "The `next_sentence_label` argument is deprecated and will be removed in a future version, use"
- " `labels` instead.",
- FutureWarning,
- )
- labels = kwargs.pop("next_sentence_label")
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- seq_relationship_scores = self.cls(pooled_output)
-
- next_sentence_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1))
-
- if not return_dict:
- output = (seq_relationship_scores,) + outputs[2:]
- return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output
-
- return NextSentencePredictorOutput(
- loss=next_sentence_loss,
- logits=seq_relationship_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Ernie Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
- output) e.g. for GLUE tasks.
- """,
- ERNIE_START_DOCSTRING,
-)
-class ErnieForSequenceClassification(ErniePreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForSequenceClassification.__init__ with Bert->Ernie,bert->ernie
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
-
- self.ernie = ErnieModel(config)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Ernie Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- ERNIE_START_DOCSTRING,
-)
-class ErnieForMultipleChoice(ErniePreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForMultipleChoice.__init__ with Bert->Ernie,bert->ernie
- def __init__(self, config):
- super().__init__(config)
-
- self.ernie = ErnieModel(config)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, 1)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
- position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
- inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.ernie(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Ernie Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- ERNIE_START_DOCSTRING,
-)
-class ErnieForTokenClassification(ErniePreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForTokenClassification.__init__ with Bert->Ernie,bert->ernie
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.ernie = ErnieModel(config, add_pooling_layer=False)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Ernie Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- ERNIE_START_DOCSTRING,
-)
-class ErnieForQuestionAnswering(ErniePreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForQuestionAnswering.__init__ with Bert->Ernie,bert->ernie
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.ernie = ErnieModel(config, add_pooling_layer=False)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- task_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- task_type_ids=task_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/ernie_m/__init__.py b/transformers/models/ernie_m/__init__.py
deleted file mode 100644
index b7cd3bdd0681c130f2d81b70faa6321e5cce9df6..0000000000000000000000000000000000000000
--- a/transformers/models/ernie_m/__init__.py
+++ /dev/null
@@ -1,82 +0,0 @@
-# Copyright 2023 The HuggingFace and Baidu Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-# rely on isort to merge the imports
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available
-
-
-_import_structure = {
- "configuration_ernie_m": ["ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP", "ErnieMConfig"],
-}
-
-try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_ernie_m"] = ["ErnieMTokenizer"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_ernie_m"] = [
- "ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST",
- "ErnieMForMultipleChoice",
- "ErnieMForQuestionAnswering",
- "ErnieMForSequenceClassification",
- "ErnieMForTokenClassification",
- "ErnieMModel",
- "ErnieMPreTrainedModel",
- "ErnieMForInformationExtraction",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_ernie_m import ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieMConfig
-
- try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_ernie_m import ErnieMTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_ernie_m import (
- ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST,
- ErnieMForInformationExtraction,
- ErnieMForMultipleChoice,
- ErnieMForQuestionAnswering,
- ErnieMForSequenceClassification,
- ErnieMForTokenClassification,
- ErnieMModel,
- ErnieMPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/ernie_m/__pycache__/__init__.cpython-310.pyc b/transformers/models/ernie_m/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 48be2ba6f8c33884367facffeeddd53a1ff1fc60..0000000000000000000000000000000000000000
Binary files a/transformers/models/ernie_m/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ernie_m/__pycache__/configuration_ernie_m.cpython-310.pyc b/transformers/models/ernie_m/__pycache__/configuration_ernie_m.cpython-310.pyc
deleted file mode 100644
index 48646911924b2e91199d9b0e0004376dc9e21b3c..0000000000000000000000000000000000000000
Binary files a/transformers/models/ernie_m/__pycache__/configuration_ernie_m.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ernie_m/__pycache__/modeling_ernie_m.cpython-310.pyc b/transformers/models/ernie_m/__pycache__/modeling_ernie_m.cpython-310.pyc
deleted file mode 100644
index 08be3bfe6fa4bb2cf24789ca59cab07a2d49e63c..0000000000000000000000000000000000000000
Binary files a/transformers/models/ernie_m/__pycache__/modeling_ernie_m.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ernie_m/__pycache__/tokenization_ernie_m.cpython-310.pyc b/transformers/models/ernie_m/__pycache__/tokenization_ernie_m.cpython-310.pyc
deleted file mode 100644
index 97f3f9ac5b6eed7d439316966062ef6217c07e2f..0000000000000000000000000000000000000000
Binary files a/transformers/models/ernie_m/__pycache__/tokenization_ernie_m.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/ernie_m/configuration_ernie_m.py b/transformers/models/ernie_m/configuration_ernie_m.py
deleted file mode 100644
index 96451c9d9c999cd744efa41c90510822757561ed..0000000000000000000000000000000000000000
--- a/transformers/models/ernie_m/configuration_ernie_m.py
+++ /dev/null
@@ -1,112 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" ErnieM model configuration"""
-# Adapted from original paddlenlp repository.(https://github.com/PaddlePaddle/PaddleNLP/blob/develop/paddlenlp/transformers/ernie_m/configuration.py)
-
-from __future__ import annotations
-
-from typing import Dict
-
-from ...configuration_utils import PretrainedConfig
-from ..deprecated._archive_maps import ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class ErnieMConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`ErnieMModel`]. It is used to instantiate a
- Ernie-M model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the `Ernie-M`
- [susnato/ernie-m-base_pytorch](https://huggingface.co/susnato/ernie-m-base_pytorch) architecture.
-
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 250002):
- Vocabulary size of `inputs_ids` in [`ErnieMModel`]. Also is the vocab size of token embedding matrix.
- Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling
- [`ErnieMModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the embedding layer, encoder layers and pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the feed-forward (ff) layer in the encoder. Input tensors to feed-forward layers are
- firstly projected from hidden_size to intermediate_size, and then projected back to hidden_size. Typically
- intermediate_size is larger than hidden_size.
- hidden_act (`str`, *optional*, defaults to `"gelu"`):
- The non-linear activation function in the feed-forward layer. `"gelu"`, `"relu"` and any other torch
- supported activation functions are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings and encoder.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability used in `MultiHeadAttention` in all encoder layers to drop some attention target.
- max_position_embeddings (`int`, *optional*, defaults to 514):
- The maximum value of the dimensionality of position encoding, which dictates the maximum supported length
- of an input sequence.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the normal initializer for initializing all weight matrices. The index of padding
- token in the token vocabulary.
- pad_token_id (`int`, *optional*, defaults to 1):
- Padding token id.
- layer_norm_eps (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the layer normalization layers.
- classifier_dropout (`float`, *optional*):
- The dropout ratio for the classification head.
- act_dropout (`float`, *optional*, defaults to 0.0):
- This dropout probability is used in `ErnieMEncoderLayer` after activation.
-
- A normal_initializer initializes weight matrices as normal distributions. See
- `ErnieMPretrainedModel._init_weights()` for how weights are initialized in `ErnieMModel`.
- """
-
- model_type = "ernie_m"
- attribute_map: Dict[str, str] = {"dropout": "classifier_dropout", "num_classes": "num_labels"}
-
- def __init__(
- self,
- vocab_size: int = 250002,
- hidden_size: int = 768,
- num_hidden_layers: int = 12,
- num_attention_heads: int = 12,
- intermediate_size: int = 3072,
- hidden_act: str = "gelu",
- hidden_dropout_prob: float = 0.1,
- attention_probs_dropout_prob: float = 0.1,
- max_position_embeddings: int = 514,
- initializer_range: float = 0.02,
- pad_token_id: int = 1,
- layer_norm_eps: float = 1e-05,
- classifier_dropout=None,
- act_dropout=0.0,
- **kwargs,
- ):
- super().__init__(pad_token_id=pad_token_id, **kwargs)
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.classifier_dropout = classifier_dropout
- self.act_dropout = act_dropout
diff --git a/transformers/models/ernie_m/modeling_ernie_m.py b/transformers/models/ernie_m/modeling_ernie_m.py
deleted file mode 100644
index ac56e120a0c3d40769d109c85910cb62142c8268..0000000000000000000000000000000000000000
--- a/transformers/models/ernie_m/modeling_ernie_m.py
+++ /dev/null
@@ -1,1058 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch ErnieM model."""
-
-
-import math
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn, tensor
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutputWithPastAndCrossAttentions,
- BaseModelOutputWithPoolingAndCrossAttentions,
- MultipleChoiceModelOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_ernie_m import ErnieMConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "susnato/ernie-m-base_pytorch"
-_CONFIG_FOR_DOC = "ErnieMConfig"
-_TOKENIZER_FOR_DOC = "ErnieMTokenizer"
-
-
-from ..deprecated._archive_maps import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Adapted from paddlenlp.transformers.ernie_m.modeling.ErnieEmbeddings
-class ErnieMEmbeddings(nn.Module):
- """Construct the embeddings from word and position embeddings."""
-
- def __init__(self, config):
- super().__init__()
- self.hidden_size = config.hidden_size
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(
- config.max_position_embeddings, config.hidden_size, padding_idx=config.pad_token_id
- )
- self.layer_norm = nn.LayerNorm(normalized_shape=config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(p=config.hidden_dropout_prob)
- self.padding_idx = config.pad_token_id
-
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.LongTensor] = None,
- past_key_values_length: int = 0,
- ) -> torch.Tensor:
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
- if position_ids is None:
- input_shape = inputs_embeds.size()[:-1]
- ones = torch.ones(input_shape, dtype=torch.int64, device=inputs_embeds.device)
- seq_length = torch.cumsum(ones, dim=1)
- position_ids = seq_length - ones
-
- if past_key_values_length > 0:
- position_ids = position_ids + past_key_values_length
- # to mimic paddlenlp implementation
- position_ids += 2
- position_embeddings = self.position_embeddings(position_ids)
- embeddings = inputs_embeds + position_embeddings
- embeddings = self.layer_norm(embeddings)
- embeddings = self.dropout(embeddings)
-
- return embeddings
-
-
-# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->ErnieM,self.value->self.v_proj,self.key->self.k_proj,self.query->self.q_proj
-class ErnieMSelfAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.q_proj = nn.Linear(config.hidden_size, self.all_head_size)
- self.k_proj = nn.Linear(config.hidden_size, self.all_head_size)
- self.v_proj = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
- self.position_embedding_type = position_embedding_type or getattr(
- config, "position_embedding_type", "absolute"
- )
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- self.max_position_embeddings = config.max_position_embeddings
- self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
-
- self.is_decoder = config.is_decoder
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- mixed_query_layer = self.q_proj(hidden_states)
-
- # If this is instantiated as a cross-attention module, the keys
- # and values come from an encoder; the attention mask needs to be
- # such that the encoder's padding tokens are not attended to.
- is_cross_attention = encoder_hidden_states is not None
-
- if is_cross_attention and past_key_value is not None:
- # reuse k,v, cross_attentions
- key_layer = past_key_value[0]
- value_layer = past_key_value[1]
- attention_mask = encoder_attention_mask
- elif is_cross_attention:
- key_layer = self.transpose_for_scores(self.k_proj(encoder_hidden_states))
- value_layer = self.transpose_for_scores(self.v_proj(encoder_hidden_states))
- attention_mask = encoder_attention_mask
- elif past_key_value is not None:
- key_layer = self.transpose_for_scores(self.k_proj(hidden_states))
- value_layer = self.transpose_for_scores(self.v_proj(hidden_states))
- key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
- value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
- else:
- key_layer = self.transpose_for_scores(self.k_proj(hidden_states))
- value_layer = self.transpose_for_scores(self.v_proj(hidden_states))
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- use_cache = past_key_value is not None
- if self.is_decoder:
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_layer, value_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- query_length, key_length = query_layer.shape[2], key_layer.shape[2]
- if use_cache:
- position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(
- -1, 1
- )
- else:
- position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
- position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
- distance = position_ids_l - position_ids_r
-
- positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
- positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
-
- if self.position_embedding_type == "relative_key":
- relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores
- elif self.position_embedding_type == "relative_key_query":
- relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in ErnieMModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- if self.is_decoder:
- outputs = outputs + (past_key_value,)
- return outputs
-
-
-class ErnieMAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- self.self_attn = ErnieMSelfAttention(config, position_embedding_type=position_embedding_type)
- self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self_attn.num_attention_heads, self.self_attn.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self_attn.q_proj = prune_linear_layer(self.self_attn.q_proj, index)
- self.self_attn.k_proj = prune_linear_layer(self.self_attn.k_proj, index)
- self.self_attn.v_proj = prune_linear_layer(self.self_attn.v_proj, index)
- self.out_proj = prune_linear_layer(self.out_proj, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self_attn.num_attention_heads = self.self_attn.num_attention_heads - len(heads)
- self.self_attn.all_head_size = self.self_attn.attention_head_size * self.self_attn.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- self_outputs = self.self_attn(
- hidden_states,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- attention_output = self.out_proj(self_outputs[0])
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-class ErnieMEncoderLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- # to mimic paddlenlp implementation
- dropout = 0.1 if config.hidden_dropout_prob is None else config.hidden_dropout_prob
- act_dropout = config.hidden_dropout_prob if config.act_dropout is None else config.act_dropout
-
- self.self_attn = ErnieMAttention(config)
- self.linear1 = nn.Linear(config.hidden_size, config.intermediate_size)
- self.dropout = nn.Dropout(act_dropout)
- self.linear2 = nn.Linear(config.intermediate_size, config.hidden_size)
- self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout1 = nn.Dropout(dropout)
- self.dropout2 = nn.Dropout(dropout)
- if isinstance(config.hidden_act, str):
- self.activation = ACT2FN[config.hidden_act]
- else:
- self.activation = config.hidden_act
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = True,
- ):
- residual = hidden_states
- if output_attentions:
- hidden_states, attention_opt_weights = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- head_mask=head_mask,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- )
-
- else:
- hidden_states = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- head_mask=head_mask,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- )
- hidden_states = residual + self.dropout1(hidden_states)
- hidden_states = self.norm1(hidden_states)
- residual = hidden_states
-
- hidden_states = self.linear1(hidden_states)
- hidden_states = self.activation(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.linear2(hidden_states)
- hidden_states = residual + self.dropout2(hidden_states)
- hidden_states = self.norm2(hidden_states)
-
- if output_attentions:
- return hidden_states, attention_opt_weights
- else:
- return hidden_states
-
-
-class ErnieMEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.layers = nn.ModuleList([ErnieMEncoderLayer(config) for _ in range(config.num_hidden_layers)])
-
- def forward(
- self,
- input_embeds: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- output_hidden_states: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
- hidden_states = () if output_hidden_states else None
- attentions = () if output_attentions else None
-
- output = input_embeds
- if output_hidden_states:
- hidden_states = hidden_states + (output,)
- for i, layer in enumerate(self.layers):
- layer_head_mask = head_mask[i] if head_mask is not None else None
- past_key_value = past_key_values[i] if past_key_values is not None else None
-
- output, opt_attn_weights = layer(
- hidden_states=output,
- attention_mask=attention_mask,
- head_mask=layer_head_mask,
- past_key_value=past_key_value,
- )
-
- if output_hidden_states:
- hidden_states = hidden_states + (output,)
- if output_attentions:
- attentions = attentions + (opt_attn_weights,)
-
- last_hidden_state = output
- if not return_dict:
- return tuple(v for v in [last_hidden_state, hidden_states, attentions] if v is not None)
-
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=last_hidden_state, hidden_states=hidden_states, attentions=attentions
- )
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->ErnieM
-class ErnieMPooler(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = nn.Tanh()
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(first_token_tensor)
- pooled_output = self.activation(pooled_output)
- return pooled_output
-
-
-class ErnieMPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = ErnieMConfig
- base_model_prefix = "ernie_m"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-ERNIE_M_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`ErnieMConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ERNIE_M_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`ErnieMTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare ErnieM Model transformer outputting raw hidden-states without any specific head on top.",
- ERNIE_M_START_DOCSTRING,
-)
-class ErnieMModel(ErnieMPreTrainedModel):
- def __init__(self, config, add_pooling_layer=True):
- super(ErnieMModel, self).__init__(config)
- self.initializer_range = config.initializer_range
- self.embeddings = ErnieMEmbeddings(config)
- self.encoder = ErnieMEncoder(config)
- self.pooler = ErnieMPooler(config) if add_pooling_layer else None
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layers[layer].self_attn.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(ERNIE_M_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- processor_class=_TOKENIZER_FOR_DOC,
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPastAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[tensor] = None,
- position_ids: Optional[tensor] = None,
- attention_mask: Optional[tensor] = None,
- head_mask: Optional[tensor] = None,
- inputs_embeds: Optional[tensor] = None,
- past_key_values: Optional[Tuple[Tuple[tensor]]] = None,
- use_cache: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.FloatTensor], BaseModelOutputWithPoolingAndCrossAttentions]:
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time.")
-
- # init the default bool value
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- past_key_values_length = 0
- if past_key_values is not None:
- past_key_values_length = past_key_values[0][0].shape[2]
-
- # Adapted from paddlenlp.transformers.ernie_m.ErnieMModel
- if attention_mask is None:
- attention_mask = (input_ids == self.config.pad_token_id).to(torch.float32)
- attention_mask *= torch.finfo(attention_mask.dtype).min
- if past_key_values is not None:
- batch_size = past_key_values[0][0].shape[0]
- past_mask = torch.zeros([batch_size, 1, 1, past_key_values_length], dtype=attention_mask.dtype)
- attention_mask = torch.concat([past_mask, attention_mask], dim=-1)
- # For 2D attention_mask from tokenizer
- elif attention_mask.ndim == 2:
- attention_mask = attention_mask.to(torch.float32)
- attention_mask = 1.0 - attention_mask
- attention_mask *= torch.finfo(attention_mask.dtype).min
-
- extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- past_key_values_length=past_key_values_length,
- )
- encoder_outputs = self.encoder(
- embedding_output,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- past_key_values=past_key_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- sequence_output = encoder_outputs[0]
- pooler_output = self.pooler(sequence_output) if self.pooler is not None else None
- return (sequence_output, pooler_output) + encoder_outputs[1:]
-
- sequence_output = encoder_outputs["last_hidden_state"]
- pooler_output = self.pooler(sequence_output) if self.pooler is not None else None
- hidden_states = None if not output_hidden_states else encoder_outputs["hidden_states"]
- attentions = None if not output_attentions else encoder_outputs["attentions"]
-
- return BaseModelOutputWithPoolingAndCrossAttentions(
- last_hidden_state=sequence_output,
- pooler_output=pooler_output,
- hidden_states=hidden_states,
- attentions=attentions,
- )
-
-
-@add_start_docstrings(
- """ErnieM Model transformer with a sequence classification/regression head on top (a linear layer on top of
- the pooled output) e.g. for GLUE tasks.""",
- ERNIE_M_START_DOCSTRING,
-)
-class ErnieMForSequenceClassification(ErnieMPreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForSequenceClassification.__init__ with Bert->ErnieM,bert->ernie_m
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
-
- self.ernie_m = ErnieMModel(config)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_M_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- processor_class=_TOKENIZER_FOR_DOC,
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.Tensor]] = None,
- use_cache: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- return_dict: Optional[bool] = True,
- labels: Optional[torch.Tensor] = None,
- ) -> Union[Tuple[torch.FloatTensor], SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie_m(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- past_key_values=past_key_values,
- output_hidden_states=output_hidden_states,
- output_attentions=output_attentions,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """ErnieM Model with a multiple choice classification head on top (a linear layer on top of
- the pooled output and a softmax) e.g. for RocStories/SWAG tasks.""",
- ERNIE_M_START_DOCSTRING,
-)
-class ErnieMForMultipleChoice(ErnieMPreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForMultipleChoice.__init__ with Bert->ErnieM,bert->ernie_m
- def __init__(self, config):
- super().__init__(config)
-
- self.ernie_m = ErnieMModel(config)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, 1)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_M_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple[torch.FloatTensor], MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
- inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.ernie_m(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """ErnieM Model with a token classification head on top (a linear layer on top of
- the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks.""",
- ERNIE_M_START_DOCSTRING,
-)
-class ErnieMForTokenClassification(ErnieMPreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForTokenClassification.__init__ with Bert->ErnieM,bert->ernie_m
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.ernie_m = ErnieMModel(config, add_pooling_layer=False)
- classifier_dropout = (
- config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
- )
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_M_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- processor_class=_TOKENIZER_FOR_DOC,
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.Tensor]] = None,
- output_hidden_states: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- return_dict: Optional[bool] = True,
- labels: Optional[torch.Tensor] = None,
- ) -> Union[Tuple[torch.FloatTensor], TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie_m(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- past_key_values=past_key_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """ErnieM Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).""",
- ERNIE_M_START_DOCSTRING,
-)
-class ErnieMForQuestionAnswering(ErnieMPreTrainedModel):
- # Copied from transformers.models.bert.modeling_bert.BertForQuestionAnswering.__init__ with Bert->ErnieM,bert->ernie_m
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.ernie_m = ErnieMModel(config, add_pooling_layer=False)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_M_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- processor_class=_TOKENIZER_FOR_DOC,
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple[torch.FloatTensor], QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.ernie_m(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """ErnieMForInformationExtraction is a Ernie-M Model with two linear layer on top of the hidden-states output to
- compute `start_prob` and `end_prob`, designed for Universal Information Extraction.""",
- ERNIE_M_START_DOCSTRING,
-)
-# Copied from paddlenlp.transformers.ernie_m.modeling.UIEM
-class ErnieMForInformationExtraction(ErnieMPreTrainedModel):
- def __init__(self, config):
- super(ErnieMForInformationExtraction, self).__init__(config)
- self.ernie_m = ErnieMModel(config)
- self.linear_start = nn.Linear(config.hidden_size, 1)
- self.linear_end = nn.Linear(config.hidden_size, 1)
- self.sigmoid = nn.Sigmoid()
- self.post_init()
-
- @add_start_docstrings_to_model_forward(ERNIE_M_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple[torch.FloatTensor], QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for position (index) for computing the start_positions loss. Position outside of the sequence are
- not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) for computing the end_positions loss. Position outside of the sequence are not
- taken into account for computing the loss.
- """
-
- result = self.ernie_m(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- if return_dict:
- sequence_output = result.last_hidden_state
- elif not return_dict:
- sequence_output = result[0]
-
- start_logits = self.linear_start(sequence_output)
- start_logits = start_logits.squeeze(-1)
- end_logits = self.linear_end(sequence_output)
- end_logits = end_logits.squeeze(-1)
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = BCEWithLogitsLoss()
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- return tuple(
- i
- for i in [total_loss, start_logits, end_logits, result.hidden_states, result.attentions]
- if i is not None
- )
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=result.hidden_states,
- attentions=result.attentions,
- )
diff --git a/transformers/models/ernie_m/tokenization_ernie_m.py b/transformers/models/ernie_m/tokenization_ernie_m.py
deleted file mode 100644
index 0bd7edea1cab3a573fcd6dc07a12e98860588251..0000000000000000000000000000000000000000
--- a/transformers/models/ernie_m/tokenization_ernie_m.py
+++ /dev/null
@@ -1,405 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Xuan Ouyang, Shuohuan Wang, Chao Pang, Yu Sun, Hao Tian, Hua Wu, Haifeng Wang and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for Ernie-M."""
-
-import io
-import os
-import unicodedata
-from typing import Any, Dict, List, Optional, Tuple
-
-import sentencepiece as spm
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-SPIECE_UNDERLINE = "▁"
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "sentencepiece_model_ckpt": "sentencepiece.bpe.model"}
-
-RESOURCE_FILES_NAMES = {
- "sentencepiece_model_file": "sentencepiece.bpe.model",
- "vocab_file": "vocab.txt",
-}
-
-
-# Adapted from paddlenlp.transformers.ernie_m.tokenizer.ErnieMTokenizer
-class ErnieMTokenizer(PreTrainedTokenizer):
- r"""
- Constructs a Ernie-M tokenizer. It uses the `sentencepiece` tools to cut the words to sub-words.
-
- Args:
- sentencepiece_model_file (`str`):
- The file path of sentencepiece model.
- vocab_file (`str`, *optional*):
- The file path of the vocabulary.
- do_lower_case (`str`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- unk_token (`str`, *optional*, defaults to `"[UNK]"`):
- A special token representing the `unknown (out-of-vocabulary)` token. An unknown token is set to be
- `unk_token` inorder to be converted to an ID.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- A special token separating two different sentences in the same input.
- pad_token (`str`, *optional*, defaults to `"[PAD]"`):
- A special token used to make arrays of tokens the same size for batching purposes.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- A special token used for sequence classification. It is the last token of the sequence when built with
- special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- A special token representing a masked token. This is the token used in the masked language modeling task
- which the model tries to predict the original unmasked ones.
- """
-
- # Ernie-M model doesn't have token_type embedding.
- model_input_names: List[str] = ["input_ids"]
-
- vocab_files_names = VOCAB_FILES_NAMES
- resource_files_names = RESOURCE_FILES_NAMES
-
- def __init__(
- self,
- sentencepiece_model_ckpt,
- vocab_file=None,
- do_lower_case=False,
- encoding="utf8",
- unk_token="[UNK]",
- sep_token="[SEP]",
- pad_token="[PAD]",
- cls_token="[CLS]",
- mask_token="[MASK]",
- sp_model_kwargs: Optional[Dict[str, Any]] = None,
- **kwargs,
- ) -> None:
- # Mask token behave like a normal word, i.e. include the space before it and
- # is included in the raw text, there should be a match in a non-normalized sentence.
-
- self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
-
- self.do_lower_case = do_lower_case
- self.sentencepiece_model_ckpt = sentencepiece_model_ckpt
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(sentencepiece_model_ckpt)
-
- # to mimic paddlenlp.transformers.ernie_m.tokenizer.ErnieMTokenizer functioning
- if vocab_file is not None:
- self.vocab = self.load_vocab(filepath=vocab_file)
- else:
- self.vocab = {self.sp_model.id_to_piece(id): id for id in range(self.sp_model.get_piece_size())}
- self.reverse_vocab = {v: k for k, v in self.vocab.items()}
-
- super().__init__(
- do_lower_case=do_lower_case,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- vocab_file=vocab_file,
- encoding=encoding,
- sp_model_kwargs=self.sp_model_kwargs,
- **kwargs,
- )
-
- def get_offset_mapping(self, text):
- if text is None:
- return None
-
- split_tokens = self.tokenize(text)
- normalized_text, char_mapping = "", []
-
- for i, ch in enumerate(text):
- if ch in self.SP_CHAR_MAPPING:
- ch = self.SP_CHAR_MAPPING.get(ch)
- else:
- ch = unicodedata.normalize("NFKC", ch)
- if self.is_whitespace(ch):
- continue
- normalized_text += ch
- char_mapping.extend([i] * len(ch))
-
- text, token_mapping, offset = normalized_text, [], 0
-
- if self.do_lower_case:
- text = text.lower()
-
- for token in split_tokens:
- if token[:1] == "▁":
- token = token[1:]
- start = text[offset:].index(token) + offset
- end = start + len(token)
-
- token_mapping.append((char_mapping[start], char_mapping[end - 1] + 1))
- offset = end
- return token_mapping
-
- @property
- def vocab_size(self):
- return len(self.vocab)
-
- def get_vocab(self):
- return dict(self.vocab, **self.added_tokens_encoder)
-
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sp_model"] = None
- return state
-
- def __setstate__(self, d):
- self.__dict__ = d
-
- # for backward compatibility
- if not hasattr(self, "sp_model_kwargs"):
- self.sp_model_kwargs = {}
-
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(self.sentencepiece_model_ckpt)
-
- def clean_text(self, text):
- """Performs invalid character removal and whitespace cleanup on text."""
- return "".join((self.SP_CHAR_MAPPING.get(c, c) for c in text))
-
- def _tokenize(self, text, enable_sampling=False, nbest_size=64, alpha=0.1):
- """Tokenize a string."""
-
- if self.sp_model_kwargs.get("enable_sampling") is True:
- enable_sampling = True
- if self.sp_model_kwargs.get("alpha") is not None:
- alpha = self.sp_model_kwargs.get("alpha")
- if self.sp_model_kwargs.get("nbest_size") is not None:
- nbest_size = self.sp_model_kwargs.get("nbest_size")
-
- if not enable_sampling:
- pieces = self.sp_model.EncodeAsPieces(text)
- else:
- pieces = self.sp_model.SampleEncodeAsPieces(text, nbest_size, alpha)
- new_pieces = []
- for pi, piece in enumerate(pieces):
- if piece == SPIECE_UNDERLINE:
- if not pieces[pi + 1].startswith(SPIECE_UNDERLINE) and pi != 0:
- new_pieces.append(SPIECE_UNDERLINE)
- continue
- else:
- continue
- lst_i = 0
- for i, chunk in enumerate(piece):
- if chunk == SPIECE_UNDERLINE:
- continue
- if self.is_ch_char(chunk) or self.is_punct(chunk):
- if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
- new_pieces.append(piece[lst_i:i])
- new_pieces.append(chunk)
- lst_i = i + 1
- elif chunk.isdigit() and i > 0 and not piece[i - 1].isdigit():
- if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
- new_pieces.append(piece[lst_i:i])
- lst_i = i
- elif not chunk.isdigit() and i > 0 and piece[i - 1].isdigit():
- if i > lst_i and piece[lst_i:i] != SPIECE_UNDERLINE:
- new_pieces.append(piece[lst_i:i])
- lst_i = i
- if len(piece) > lst_i:
- new_pieces.append(piece[lst_i:])
- return new_pieces
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (strings for sub-words) in a single string."""
- out_string = "".join(tokens).replace(SPIECE_UNDERLINE, " ").strip()
- return out_string
-
- def convert_ids_to_string(self, ids):
- """
- Converts a sequence of tokens (strings for sub-words) in a single string.
- """
- tokens = self.convert_ids_to_tokens(ids)
- out_string = "".join(tokens).replace(SPIECE_UNDERLINE, " ").strip()
- return out_string
-
- # to mimic paddlenlp.transformers.ernie_m.tokenizer.ErnieMTokenizer functioning
- def _convert_token_to_id(self, token):
- return self.vocab.get(token, self.vocab.get(self.unk_token))
-
- # to mimic paddlenlp.transformers.ernie_m.tokenizer.ErnieMTokenizer functioning
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.reverse_vocab.get(index, self.unk_token)
-
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- r"""
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. An ErnieM sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- Returns:
- `List[int]`: List of input_id with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- _cls = [self.cls_token_id]
- _sep = [self.sep_token_id]
- return _cls + token_ids_0 + _sep + _sep + token_ids_1 + _sep
-
- def build_offset_mapping_with_special_tokens(self, offset_mapping_0, offset_mapping_1=None):
- r"""
- Build offset map from a pair of offset map by concatenating and adding offsets of special tokens. An Ernie-M
- offset_mapping has the following format:
-
- - single sequence: `(0,0) X (0,0)`
- - pair of sequences: `(0,0) A (0,0) (0,0) B (0,0)`
-
- Args:
- offset_mapping_ids_0 (`List[tuple]`):
- List of char offsets to which the special tokens will be added.
- offset_mapping_ids_1 (`List[tuple]`, *optional*):
- Optional second list of wordpiece offsets for offset mapping pairs.
- Returns:
- `List[tuple]`: List of wordpiece offsets with the appropriate offsets of special tokens.
- """
- if offset_mapping_1 is None:
- return [(0, 0)] + offset_mapping_0 + [(0, 0)]
-
- return [(0, 0)] + offset_mapping_0 + [(0, 0), (0, 0)] + offset_mapping_1 + [(0, 0)]
-
- def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
- r"""
- Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `encode` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of ids of the first sequence.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`str`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
- Returns:
- `List[int]`:
- The list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- if token_ids_1 is not None:
- raise ValueError(
- "You should not supply a second sequence if the provided sequence of "
- "ids is already formatted with special tokens for the model."
- )
- return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_0]
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create the token type IDs corresponding to the sequences passed. [What are token type
- IDs?](../glossary#token-type-ids) Should be overridden in a subclass if the model has a special way of
- building: those.
-
- Args:
- token_ids_0 (`List[int]`):
- The first tokenized sequence.
- token_ids_1 (`List[int]`, *optional*):
- The second tokenized sequence.
- Returns:
- `List[int]`: The token type ids.
- """
- # called when `add_special_tokens` is True, so align with `build_inputs_with_special_tokens` method
- if token_ids_1 is None:
- # [CLS] X [SEP]
- return (len(token_ids_0) + 2) * [0]
-
- # [CLS] A [SEP] [SEP] B [SEP]
- return [0] * (len(token_ids_0) + 1) + [1] * (len(token_ids_1) + 3)
-
- def is_ch_char(self, char):
- """
- is_ch_char
- """
- if "\u4e00" <= char <= "\u9fff":
- return True
- return False
-
- def is_alpha(self, char):
- """
- is_alpha
- """
- if ("a" <= char <= "z") or ("A" <= char <= "Z"):
- return True
- return False
-
- def is_punct(self, char):
- """
- is_punct
- """
- if char in ",;:.?!~，；：。？！《》【】":
- return True
- return False
-
- def is_whitespace(self, char):
- """
- is whitespace
- """
- if char == " " or char == "\t" or char == "\n" or char == "\r":
- return True
- if len(char) == 1:
- cat = unicodedata.category(char)
- if cat == "Zs":
- return True
- return False
-
- def load_vocab(self, filepath):
- token_to_idx = {}
- with io.open(filepath, "r", encoding="utf-8") as f:
- for index, line in enumerate(f):
- token = line.rstrip("\n")
- token_to_idx[token] = int(index)
-
- return token_to_idx
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- index = 0
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- else:
- vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(token + "\n")
- index += 1
-
- tokenizer_model_file = os.path.join(save_directory, "sentencepiece.bpe.model")
- with open(tokenizer_model_file, "wb") as fi:
- content_spiece_model = self.sp_model.serialized_model_proto()
- fi.write(content_spiece_model)
-
- return (vocab_file,)
diff --git a/transformers/models/esm/__init__.py b/transformers/models/esm/__init__.py
deleted file mode 100644
index 1b07db5a5eea64b8e5d37cf2c9c89429586ea8fe..0000000000000000000000000000000000000000
--- a/transformers/models/esm/__init__.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# Copyright 2022 Facebook and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available
-
-
-_import_structure = {
- "configuration_esm": ["ESM_PRETRAINED_CONFIG_ARCHIVE_MAP", "EsmConfig"],
- "tokenization_esm": ["EsmTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_esm"] = [
- "ESM_PRETRAINED_MODEL_ARCHIVE_LIST",
- "EsmForMaskedLM",
- "EsmForSequenceClassification",
- "EsmForTokenClassification",
- "EsmModel",
- "EsmPreTrainedModel",
- ]
- _import_structure["modeling_esmfold"] = ["EsmForProteinFolding", "EsmFoldPreTrainedModel"]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_esm"] = [
- "TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFEsmForMaskedLM",
- "TFEsmForSequenceClassification",
- "TFEsmForTokenClassification",
- "TFEsmModel",
- "TFEsmPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_esm import ESM_PRETRAINED_CONFIG_ARCHIVE_MAP, EsmConfig
- from .tokenization_esm import EsmTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_esm import (
- ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
- EsmForMaskedLM,
- EsmForSequenceClassification,
- EsmForTokenClassification,
- EsmModel,
- EsmPreTrainedModel,
- )
- from .modeling_esmfold import EsmFoldPreTrainedModel, EsmForProteinFolding
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_esm import (
- TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFEsmForMaskedLM,
- TFEsmForSequenceClassification,
- TFEsmForTokenClassification,
- TFEsmModel,
- TFEsmPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
diff --git a/transformers/models/esm/__pycache__/__init__.cpython-310.pyc b/transformers/models/esm/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index bfc714242e118e54ebcc398642245c45dd00601c..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/__pycache__/configuration_esm.cpython-310.pyc b/transformers/models/esm/__pycache__/configuration_esm.cpython-310.pyc
deleted file mode 100644
index 358e122508705dfb3468887b17a7b864bc8a9024..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/__pycache__/configuration_esm.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/__pycache__/convert_esm.cpython-310.pyc b/transformers/models/esm/__pycache__/convert_esm.cpython-310.pyc
deleted file mode 100644
index b4a0eeaa822c30d9b9b7a04bb6d78310ed3ecb1f..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/__pycache__/convert_esm.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/__pycache__/modeling_esm.cpython-310.pyc b/transformers/models/esm/__pycache__/modeling_esm.cpython-310.pyc
deleted file mode 100644
index 44b88e99a68958b409901c36cd97b4f8dd88040a..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/__pycache__/modeling_esm.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/__pycache__/modeling_esmfold.cpython-310.pyc b/transformers/models/esm/__pycache__/modeling_esmfold.cpython-310.pyc
deleted file mode 100644
index 04c471997fb990cca062854e2290396c29ab6017..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/__pycache__/modeling_esmfold.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/__pycache__/modeling_tf_esm.cpython-310.pyc b/transformers/models/esm/__pycache__/modeling_tf_esm.cpython-310.pyc
deleted file mode 100644
index d33adceb24a6a1f18298814f5ccde47f1661205a..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/__pycache__/modeling_tf_esm.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/__pycache__/tokenization_esm.cpython-310.pyc b/transformers/models/esm/__pycache__/tokenization_esm.cpython-310.pyc
deleted file mode 100644
index 53145a28f48d625994b36875ef5114a79c0775b1..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/__pycache__/tokenization_esm.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/configuration_esm.py b/transformers/models/esm/configuration_esm.py
deleted file mode 100644
index 31d309cb04a0175d6865d7f79f5f27241a264960..0000000000000000000000000000000000000000
--- a/transformers/models/esm/configuration_esm.py
+++ /dev/null
@@ -1,361 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" ESM model configuration"""
-
-from dataclasses import asdict, dataclass
-from typing import Optional
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-# TODO Update this
-
-from ..deprecated._archive_maps import ESM_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class EsmConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`ESMModel`]. It is used to instantiate a ESM model
- according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the ESM
- [facebook/esm-1b](https://huggingface.co/facebook/esm-1b) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*):
- Vocabulary size of the ESM model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`ESMModel`].
- mask_token_id (`int`, *optional*):
- The index of the mask token in the vocabulary. This must be included in the config because of the
- "mask-dropout" scaling trick, which will scale the inputs depending on the number of masked tokens.
- pad_token_id (`int`, *optional*):
- The index of the padding token in the vocabulary. This must be included in the config because certain parts
- of the ESM code use this instead of the attention mask.
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 1026):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
- Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query", "rotary"`.
- For positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
- [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
- For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
- with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
- is_decoder (`bool`, *optional*, defaults to `False`):
- Whether the model is used as a decoder or not. If `False`, the model is used as an encoder.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- emb_layer_norm_before (`bool`, *optional*):
- Whether to apply layer normalization after embeddings but before the main stem of the network.
- token_dropout (`bool`, defaults to `False`):
- When this is enabled, masked tokens are treated as if they had been dropped out by input dropout.
-
- Examples:
-
- ```python
- >>> from transformers import EsmModel, EsmConfig
-
- >>> # Initializing a ESM facebook/esm-1b style configuration >>> configuration = EsmConfig()
-
- >>> # Initializing a model from the configuration >>> model = ESMModel(configuration)
-
- >>> # Accessing the model configuration >>> configuration = model.config
- ```"""
-
- model_type = "esm"
-
- def __init__(
- self,
- vocab_size=None,
- mask_token_id=None,
- pad_token_id=None,
- hidden_size=768,
- num_hidden_layers=12,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=1026,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- position_embedding_type="absolute",
- use_cache=True,
- emb_layer_norm_before=None,
- token_dropout=False,
- is_folding_model=False,
- esmfold_config=None,
- vocab_list=None,
- **kwargs,
- ):
- super().__init__(pad_token_id=pad_token_id, mask_token_id=mask_token_id, **kwargs)
-
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.position_embedding_type = position_embedding_type
- self.use_cache = use_cache
- self.emb_layer_norm_before = emb_layer_norm_before
- self.token_dropout = token_dropout
- self.is_folding_model = is_folding_model
- if is_folding_model:
- if esmfold_config is None:
- logger.info("No esmfold_config supplied for folding model, using default values.")
- esmfold_config = EsmFoldConfig()
- elif isinstance(esmfold_config, dict):
- esmfold_config = EsmFoldConfig(**esmfold_config)
- self.esmfold_config = esmfold_config
- if vocab_list is None:
- logger.warning("No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!")
- self.vocab_list = get_default_vocab_list()
- else:
- self.vocab_list = vocab_list
- else:
- self.esmfold_config = None
- self.vocab_list = None
- if self.esmfold_config is not None and getattr(self.esmfold_config, "use_esm_attn_map", False):
- raise ValueError("The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!")
-
- def to_dict(self):
- """
- Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
-
- Returns:
- `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
- """
- output = super().to_dict()
- if isinstance(self.esmfold_config, EsmFoldConfig):
- output["esmfold_config"] = self.esmfold_config.to_dict()
- return output
-
-
-@dataclass
-class EsmFoldConfig:
- esm_type: str = None
- fp16_esm: bool = True
- use_esm_attn_map: bool = False
- esm_ablate_pairwise: bool = False
- esm_ablate_sequence: bool = False
- esm_input_dropout: float = 0
-
- embed_aa: bool = True
- bypass_lm: bool = False
-
- lddt_head_hid_dim: int = 128
- trunk: "TrunkConfig" = None
-
- def __post_init__(self):
- if self.trunk is None:
- self.trunk = TrunkConfig()
- elif isinstance(self.trunk, dict):
- self.trunk = TrunkConfig(**self.trunk)
-
- def to_dict(self):
- """
- Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
-
- Returns:
- `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
- """
- output = asdict(self)
- output["trunk"] = self.trunk.to_dict()
- return output
-
-
-@dataclass
-class TrunkConfig:
- num_blocks: int = 48
- sequence_state_dim: int = 1024
- pairwise_state_dim: int = 128
- sequence_head_width: int = 32
- pairwise_head_width: int = 32
- position_bins: int = 32
- dropout: float = 0
- layer_drop: float = 0
- cpu_grad_checkpoint: bool = False
- max_recycles: int = 4
- chunk_size: Optional[int] = 128
- structure_module: "StructureModuleConfig" = None
-
- def __post_init__(self):
- if self.structure_module is None:
- self.structure_module = StructureModuleConfig()
- elif isinstance(self.structure_module, dict):
- self.structure_module = StructureModuleConfig(**self.structure_module)
-
- if self.max_recycles <= 0:
- raise ValueError(f"`max_recycles` should be positive, got {self.max_recycles}.")
- if self.sequence_state_dim % self.sequence_state_dim != 0:
- raise ValueError(
- "`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got"
- f" {self.sequence_state_dim} and {self.sequence_state_dim}."
- )
- if self.pairwise_state_dim % self.pairwise_state_dim != 0:
- raise ValueError(
- "`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got"
- f" {self.pairwise_state_dim} and {self.pairwise_state_dim}."
- )
-
- sequence_num_heads = self.sequence_state_dim // self.sequence_head_width
- pairwise_num_heads = self.pairwise_state_dim // self.pairwise_head_width
-
- if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
- raise ValueError(
- "`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got"
- f" {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}."
- )
- if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
- raise ValueError(
- "`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got"
- f" {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}."
- )
- if self.pairwise_state_dim % 2 != 0:
- raise ValueError(f"`pairwise_state_dim` should be even, got {self.pairwise_state_dim}.")
-
- if self.dropout >= 0.4:
- raise ValueError(f"`dropout` should not be greater than 0.4, got {self.dropout}.")
-
- def to_dict(self):
- """
- Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
-
- Returns:
- `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
- """
- output = asdict(self)
- output["structure_module"] = self.structure_module.to_dict()
- return output
-
-
-@dataclass
-class StructureModuleConfig:
- """
- Args:
- sequence_dim:
- Single representation channel dimension
- pairwise_dim:
- Pair representation channel dimension
- ipa_dim:
- IPA hidden channel dimension
- resnet_dim:
- Angle resnet (Alg. 23 lines 11-14) hidden channel dimension
- num_heads_ipa:
- Number of IPA heads
- num_qk_points:
- Number of query/key points to generate during IPA
- num_v_points:
- Number of value points to generate during IPA
- dropout_rate:
- Dropout rate used throughout the layer
- num_blocks:
- Number of structure module blocks
- num_transition_layers:
- Number of layers in the single representation transition (Alg. 23 lines 8-9)
- num_resnet_blocks:
- Number of blocks in the angle resnet
- num_angles:
- Number of angles to generate in the angle resnet
- trans_scale_factor:
- Scale of single representation transition hidden dimension
- epsilon:
- Small number used in angle resnet normalization
- inf:
- Large number used for attention masking
- """
-
- sequence_dim: int = 384
- pairwise_dim: int = 128
- ipa_dim: int = 16
- resnet_dim: int = 128
- num_heads_ipa: int = 12
- num_qk_points: int = 4
- num_v_points: int = 8
- dropout_rate: float = 0.1
- num_blocks: int = 8
- num_transition_layers: int = 1
- num_resnet_blocks: int = 2
- num_angles: int = 7
- trans_scale_factor: int = 10
- epsilon: float = 1e-8
- inf: float = 1e5
-
- def to_dict(self):
- return asdict(self)
-
-
-def get_default_vocab_list():
- return (
- "",
- "",
- "",
- "",
- "L",
- "A",
- "G",
- "V",
- "S",
- "E",
- "R",
- "T",
- "I",
- "D",
- "P",
- "K",
- "Q",
- "N",
- "F",
- "Y",
- "M",
- "H",
- "W",
- "C",
- "X",
- "B",
- "U",
- "Z",
- "O",
- ".",
- "-",
- "",
- "",
- )
diff --git a/transformers/models/esm/convert_esm.py b/transformers/models/esm/convert_esm.py
deleted file mode 100644
index 22ca3f5392c19d6b1c36a69d0738b8528bfaaa9d..0000000000000000000000000000000000000000
--- a/transformers/models/esm/convert_esm.py
+++ /dev/null
@@ -1,400 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert ESM checkpoint."""
-
-
-import argparse
-import pathlib
-from pathlib import Path
-from tempfile import TemporaryDirectory
-
-import esm as esm_module
-import torch
-from esm.esmfold.v1.misc import batch_encode_sequences as esmfold_encode_sequences
-from esm.esmfold.v1.pretrained import esmfold_v1
-
-from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
-from transformers.models.esm.modeling_esm import (
- EsmForMaskedLM,
- EsmForSequenceClassification,
- EsmIntermediate,
- EsmLayer,
- EsmOutput,
- EsmSelfAttention,
- EsmSelfOutput,
-)
-from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
-from transformers.models.esm.tokenization_esm import EsmTokenizer
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-SAMPLE_DATA = [
- (
- "protein1",
- "MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA",
- ),
- ("protein2", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA"),
- ("protein3", "MKTVRQERLKSIRILERSKEPVSGAQLAEELSSRQVIVQDIAYLRSLGYNVATPRGYVLAGG"),
- ("protein4", "MKTVRQERLKSIRILERSKEPVSGAQLAEELSSRQVIVQDIAYLRSLGYNVATPRGYVLA"),
-]
-
-MODEL_MAPPING = {
- "esm1b_t33_650M_UR50S": esm_module.pretrained.esm1b_t33_650M_UR50S,
- "esm1v_t33_650M_UR90S_1": esm_module.pretrained.esm1v_t33_650M_UR90S_1,
- "esm1v_t33_650M_UR90S_2": esm_module.pretrained.esm1v_t33_650M_UR90S_2,
- "esm1v_t33_650M_UR90S_3": esm_module.pretrained.esm1v_t33_650M_UR90S_3,
- "esm1v_t33_650M_UR90S_4": esm_module.pretrained.esm1v_t33_650M_UR90S_4,
- "esm1v_t33_650M_UR90S_5": esm_module.pretrained.esm1v_t33_650M_UR90S_5,
- "esm2_t48_15B_UR50D": esm_module.pretrained.esm2_t48_15B_UR50D,
- "esm2_t36_3B_UR50D": esm_module.pretrained.esm2_t36_3B_UR50D,
- "esm2_t33_650M_UR50D": esm_module.pretrained.esm2_t33_650M_UR50D,
- "esm2_t30_150M_UR50D": esm_module.pretrained.esm2_t30_150M_UR50D,
- "esm2_t12_35M_UR50D": esm_module.pretrained.esm2_t12_35M_UR50D,
- "esm2_t6_8M_UR50D": esm_module.pretrained.esm2_t6_8M_UR50D,
- "esmfold_v1": esmfold_v1,
-}
-
-restypes = list("ARNDCQEGHILKMFPSTWYV")
-
-restypes_with_x = restypes + ["X"]
-restypes_with_extras = restypes_with_x + ["", "", "", "", ""]
-
-
-def get_esmfold_tokenizer():
- with TemporaryDirectory() as tempdir:
- vocab = "\n".join(restypes_with_extras)
- vocab_file = Path(tempdir) / "vocab.txt"
- vocab_file.write_text(vocab)
- hf_tokenizer = EsmTokenizer(vocab_file=str(vocab_file))
- hf_tokenizer.pad_token_id = 0 # Overlaps with 'A' but that seems to be what they want
- return hf_tokenizer
-
-
-def transfer_and_check_weights(original_module, our_module):
- status = our_module.load_state_dict(original_module.state_dict())
- if status.missing_keys:
- raise ValueError(f"Missing keys: {status.missing_keys}")
- if status.unexpected_keys:
- raise ValueError(f"Unexpected keys: {status.unexpected_keys}")
-
-
-def convert_esm_checkpoint_to_pytorch(
- model: str, pytorch_dump_folder_path: str, classification_head: bool, push_to_repo: str, auth_token: str
-):
- """
- Copy/paste/tweak esm's weights to our BERT structure.
- """
- if model.startswith("esmfold"):
- esm = MODEL_MAPPING[model]()
- else:
- esm, alphabet = MODEL_MAPPING[model]()
- esm.eval() # disable dropout
-
- if model.startswith("esmfold"):
- embed_dim = esm.esm.embed_dim
- num_layers = esm.esm.num_layers
- num_attention_heads = esm.esm.attention_heads
- intermediate_size = 4 * embed_dim
- token_dropout = esm.esm.token_dropout
- emb_layer_norm_before = False # This code path does not exist in ESM-2
- position_embedding_type = "rotary"
- is_folding_model = True
- esmfold_config = EsmFoldConfig()
- for key, val in esm.cfg.items():
- if hasattr(esmfold_config, key) and key != "trunk":
- setattr(esmfold_config, key, val)
- for key, val in esm.cfg.trunk.items():
- if hasattr(esmfold_config.trunk, key) and key != "structure_module":
- setattr(esmfold_config.trunk, key, val)
- for key, val in esm.cfg.trunk.structure_module.items():
- if hasattr(esmfold_config.trunk.structure_module, key):
- setattr(esmfold_config.trunk.structure_module, key, val)
- elif hasattr(esm, "args"):
- # Indicates an ESM-1b or ESM-1v model
- embed_dim = esm.args.embed_dim
- num_layers = esm.args.layers
- num_attention_heads = esm.args.attention_heads
- intermediate_size = esm.args.ffn_embed_dim
- token_dropout = esm.args.token_dropout
- emb_layer_norm_before = True if esm.emb_layer_norm_before else False
- position_embedding_type = "absolute"
- is_folding_model = False
- esmfold_config = None
- else:
- # Indicates an ESM-2 model
- embed_dim = esm.embed_dim
- num_layers = esm.num_layers
- num_attention_heads = esm.attention_heads
- intermediate_size = 4 * embed_dim # This is hardcoded in ESM-2
- token_dropout = esm.token_dropout
- emb_layer_norm_before = False # This code path does not exist in ESM-2
- position_embedding_type = "rotary"
- is_folding_model = False
- esmfold_config = None
-
- if is_folding_model:
- alphabet = esm.esm.alphabet
- vocab_list = tuple(alphabet.all_toks)
- mask_token_id = alphabet.mask_idx
- pad_token_id = alphabet.padding_idx
-
- if is_folding_model:
- original_esm_model = esm.esm
- else:
- original_esm_model = esm
-
- config = EsmConfig(
- vocab_size=original_esm_model.embed_tokens.num_embeddings,
- mask_token_id=mask_token_id,
- hidden_size=embed_dim,
- num_hidden_layers=num_layers,
- num_attention_heads=num_attention_heads,
- intermediate_size=intermediate_size,
- max_position_embeddings=1026,
- layer_norm_eps=1e-5, # PyTorch default used in fairseq
- attention_probs_dropout_prob=0.0,
- hidden_dropout_prob=0.0,
- pad_token_id=pad_token_id,
- emb_layer_norm_before=emb_layer_norm_before,
- token_dropout=token_dropout,
- position_embedding_type=position_embedding_type,
- is_folding_model=is_folding_model,
- esmfold_config=esmfold_config,
- vocab_list=vocab_list,
- )
- if classification_head:
- config.num_labels = esm.classification_heads["mnli"].out_proj.weight.shape[0]
- print("Our ESM config:", config)
-
- if model.startswith("esmfold"):
- model_class = EsmForProteinFolding
- elif classification_head:
- model_class = EsmForSequenceClassification
- else:
- model_class = EsmForMaskedLM
- model = model_class(config)
- model.eval()
-
- # Now let's copy all the weights.
- # Embeddings
- model.esm.embeddings.word_embeddings.weight = original_esm_model.embed_tokens.weight
- if position_embedding_type == "absolute":
- model.esm.embeddings.position_embeddings.weight = original_esm_model.embed_positions.weight
-
- if config.emb_layer_norm_before:
- model.esm.embeddings.layer_norm.weight = original_esm_model.emb_layer_norm_before.weight
- model.esm.embeddings.layer_norm.bias = original_esm_model.emb_layer_norm_before.bias
-
- model.esm.encoder.emb_layer_norm_after.weight = original_esm_model.emb_layer_norm_after.weight
- model.esm.encoder.emb_layer_norm_after.bias = original_esm_model.emb_layer_norm_after.bias
-
- for i in range(config.num_hidden_layers):
- # Encoder: start of layer
- layer: EsmLayer = model.esm.encoder.layer[i]
- # esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
- esm_layer = original_esm_model.layers[i]
-
- # self attention
- self_attn: EsmSelfAttention = layer.attention.self
- assert (
- esm_layer.self_attn.k_proj.weight.data.shape
- == esm_layer.self_attn.q_proj.weight.data.shape
- == esm_layer.self_attn.v_proj.weight.data.shape
- == torch.Size((config.hidden_size, config.hidden_size))
- )
-
- self_attn.query.weight.data = esm_layer.self_attn.q_proj.weight
- self_attn.query.bias.data = esm_layer.self_attn.q_proj.bias
- self_attn.key.weight.data = esm_layer.self_attn.k_proj.weight
- self_attn.key.bias.data = esm_layer.self_attn.k_proj.bias
- self_attn.value.weight.data = esm_layer.self_attn.v_proj.weight
- self_attn.value.bias.data = esm_layer.self_attn.v_proj.bias
-
- if getattr(esm_layer.self_attn, "rot_emb", None) is not None:
- # Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
- # During the training of ESM-2 the model was converted to float16 precision, which also converts
- # the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
- # If we recompute inv_freq without this loss of precision then we will get subtly different rotary
- # embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
- # we make sure the new model copies the data from the old inv_freq.
- self_attn.rotary_embeddings.inv_freq.data = esm_layer.self_attn.rot_emb.inv_freq
-
- # LayerNorm changes for pre-activation
- layer.attention.LayerNorm.weight = esm_layer.self_attn_layer_norm.weight
- layer.attention.LayerNorm.bias = esm_layer.self_attn_layer_norm.bias
- layer.LayerNorm.weight = esm_layer.final_layer_norm.weight
- layer.LayerNorm.bias = esm_layer.final_layer_norm.bias
-
- # self-attention output
- self_output: EsmSelfOutput = layer.attention.output
- assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
- self_output.dense.weight = esm_layer.self_attn.out_proj.weight
- self_output.dense.bias = esm_layer.self_attn.out_proj.bias
-
- # intermediate
- intermediate: EsmIntermediate = layer.intermediate
- assert intermediate.dense.weight.shape == esm_layer.fc1.weight.shape
- intermediate.dense.weight = esm_layer.fc1.weight
- intermediate.dense.bias = esm_layer.fc1.bias
-
- # output
- bert_output: EsmOutput = layer.output
- assert bert_output.dense.weight.shape == esm_layer.fc2.weight.shape
- bert_output.dense.weight = esm_layer.fc2.weight
- bert_output.dense.bias = esm_layer.fc2.bias
- # end of layer
-
- if is_folding_model:
- model.esm_s_combine.data = esm.esm_s_combine.data
- model.af2_to_esm.data = esm.af2_to_esm.data
- transfer_and_check_weights(esm.embedding, model.embedding)
- transfer_and_check_weights(esm.esm_s_mlp, model.esm_s_mlp)
- transfer_and_check_weights(esm.trunk, model.trunk)
- transfer_and_check_weights(esm.distogram_head, model.distogram_head)
- transfer_and_check_weights(esm.ptm_head, model.ptm_head)
- transfer_and_check_weights(esm.lm_head, model.lm_head)
- transfer_and_check_weights(esm.lddt_head, model.lddt_head)
-
- elif classification_head:
- model.classifier.dense.weight = esm.esm.classification_heads["mnli"].dense.weight
- model.classifier.dense.bias = esm.classification_heads["mnli"].dense.bias
- model.classifier.out_proj.weight = esm.classification_heads["mnli"].out_proj.weight
- model.classifier.out_proj.bias = esm.classification_heads["mnli"].out_proj.bias
- else:
- # LM Head
- model.lm_head.dense.weight = esm.lm_head.dense.weight
- model.lm_head.dense.bias = esm.lm_head.dense.bias
- model.lm_head.layer_norm.weight = esm.lm_head.layer_norm.weight
- model.lm_head.layer_norm.bias = esm.lm_head.layer_norm.bias
- model.lm_head.decoder.weight = esm.lm_head.weight
- model.lm_head.bias = esm.lm_head.bias
-
- # Contact prediction head
- transfer_and_check_weights(esm.contact_head, model.esm.contact_head)
-
- # Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
- if is_folding_model:
- # Folding models aren't trained on masked inputs and don't like mask tokens.
- sample_data = SAMPLE_DATA[:2]
- else:
- sample_data = SAMPLE_DATA
-
- if is_folding_model:
- hf_tokenizer = get_esmfold_tokenizer()
- hf_tokens = hf_tokenizer(
- [row[1] for row in sample_data], return_tensors="pt", padding=True, add_special_tokens=False
- )
- esmfold_aas, esmfold_mask, _, _, _ = esmfold_encode_sequences([row[1] for row in sample_data])
- success = torch.all(hf_tokens["input_ids"] == esmfold_aas) and torch.all(
- hf_tokens["attention_mask"] == esmfold_mask
- )
- else:
- # Let's check that we get the same results.
- batch_converter = alphabet.get_batch_converter()
- batch_labels, batch_strs, batch_tokens = batch_converter(sample_data)
- # Prepare tokenizer and make sure it matches
- with TemporaryDirectory() as tempdir:
- vocab = "\n".join(alphabet.all_toks)
- vocab_file = Path(tempdir) / "vocab.txt"
- vocab_file.write_text(vocab)
- hf_tokenizer = EsmTokenizer(vocab_file=str(vocab_file))
-
- hf_tokens = hf_tokenizer([row[1] for row in sample_data], return_tensors="pt", padding=True)
- success = torch.all(hf_tokens["input_ids"] == batch_tokens)
-
- print("Do both models tokenizers output the same tokens?", "🔥" if success else "💩")
- if not success:
- raise Exception("Tokenization does not match!")
-
- with torch.no_grad():
- if is_folding_model:
- # Let's test the model in parts
- # ESMFold always converts the ESM stem to float16, which requires float16 ops
- # that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
- # ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
- # original and the converted model on the GPU at the same time.
- their_output = esm.cuda().infer([row[1] for row in sample_data])
- our_output = model.cuda()(
- input_ids=hf_tokens["input_ids"].cuda(), attention_mask=hf_tokens["attention_mask"].cuda()
- )
- else:
- our_output = model(**hf_tokens, output_hidden_states=True)
- our_output = our_output["logits"]
- if classification_head:
- their_output = esm.model.classification_heads["mnli"](esm.extract_features(batch_tokens))
- else:
- their_output = esm(hf_tokens["input_ids"], repr_layers=list(range(999)))
- their_output = their_output["logits"]
-
- if is_folding_model:
- max_absolute_diff = torch.max(torch.abs(our_output["positions"] - their_output["positions"])).item()
- success = torch.allclose(our_output["positions"], their_output["positions"], atol=1e-5)
- else:
- max_absolute_diff = torch.max(torch.abs(our_output - their_output)).item()
- success = torch.allclose(our_output, their_output, atol=1e-5)
-
- print(f"max_absolute_diff = {max_absolute_diff}") # ~ 1e-5
- print("Do both models output the same tensors?", "🔥" if success else "💩")
-
- if not success:
- raise Exception("Something went wRoNg")
-
- if not is_folding_model:
- # Let's check contact prediction too
- our_output = model.predict_contacts(hf_tokens["input_ids"], hf_tokens["attention_mask"])
- their_output = esm.predict_contacts(hf_tokens["input_ids"])
- max_absolute_diff = torch.max(torch.abs(our_output - their_output)).item()
- success = torch.allclose(our_output, their_output, atol=1e-5)
-
- print("Contact prediction testing:")
- print(f"max_absolute_diff = {max_absolute_diff}") # ~ 1e-5
- print("Do both models output the same tensors?", "🔥" if success else "💩")
-
- if not success:
- raise Exception("Something went wRoNg")
-
- pathlib.Path(pytorch_dump_folder_path).mkdir(parents=True, exist_ok=True)
- print(f"Saving model to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
-
- del esm # Free up some memory before continuing
-
- print(f"Saving tokenizer to {pytorch_dump_folder_path}")
- hf_tokenizer.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_repo:
- model.push_to_hub(repo_id=push_to_repo, token_token=auth_token)
- hf_tokenizer.push_to_hub(repo_id=push_to_repo, token_token=auth_token)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--pytorch_dump_folder_path", type=str, required=True, help="Path to the output PyTorch model."
- )
- parser.add_argument(
- "--classification_head", action="store_true", help="Whether to convert a final classification head."
- )
- parser.add_argument("--model", default=None, type=str, required=True, help="Name of model to convert.")
- parser.add_argument("--push_to_repo", type=str, help="Repo to upload to (including username!).")
- parser.add_argument("--auth_token", type=str, help="HuggingFace auth token.")
- args = parser.parse_args()
- convert_esm_checkpoint_to_pytorch(
- args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
- )
diff --git a/transformers/models/esm/modeling_esm.py b/transformers/models/esm/modeling_esm.py
deleted file mode 100644
index a97ea58d7b81d9969cdac3a6d805b5fe34b9ac3f..0000000000000000000000000000000000000000
--- a/transformers/models/esm/modeling_esm.py
+++ /dev/null
@@ -1,1265 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch ESM model."""
-
-import math
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_outputs import (
- BaseModelOutputWithPastAndCrossAttentions,
- BaseModelOutputWithPoolingAndCrossAttentions,
- MaskedLMOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import logging
-from .configuration_esm import EsmConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "facebook/esm2_t6_8M_UR50D"
-_CONFIG_FOR_DOC = "EsmConfig"
-
-
-from ..deprecated._archive_maps import ESM_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-def rotate_half(x):
- x1, x2 = x.chunk(2, dim=-1)
- return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rotary_pos_emb(x, cos, sin):
- cos = cos[:, :, : x.shape[-2], :]
- sin = sin[:, :, : x.shape[-2], :]
-
- return (x * cos) + (rotate_half(x) * sin)
-
-
-def gelu(x):
- """
- This is the gelu implementation from the original ESM repo. Using F.gelu yields subtly wrong results.
- """
- return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
-
-
-def symmetrize(x):
- "Make layer symmetric in final two dimensions, used for contact prediction."
- return x + x.transpose(-1, -2)
-
-
-def average_product_correct(x):
- "Perform average product correct, used for contact prediction."
- a1 = x.sum(-1, keepdims=True)
- a2 = x.sum(-2, keepdims=True)
- a12 = x.sum((-1, -2), keepdims=True)
-
- avg = a1 * a2
- avg.div_(a12) # in-place to reduce memory
- normalized = x - avg
- return normalized
-
-
-class RotaryEmbedding(torch.nn.Module):
- """
- Rotary position embeddings based on those in
- [RoFormer](https://huggingface.co/docs/transformers/model_doc/roformer). Query and keys are transformed by rotation
- matrices which depend on their relative positions.
- """
-
- def __init__(self, dim: int):
- super().__init__()
- # Generate and save the inverse frequency buffer (non trainable)
- inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
- inv_freq = inv_freq
- self.register_buffer("inv_freq", inv_freq)
-
- self._seq_len_cached = None
- self._cos_cached = None
- self._sin_cached = None
-
- def _update_cos_sin_tables(self, x, seq_dimension=2):
- seq_len = x.shape[seq_dimension]
-
- # Reset the tables if the sequence length has changed,
- # or if we're on a new device (possibly due to tracing for instance)
- if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
- self._seq_len_cached = seq_len
- t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq)
- freqs = torch.outer(t, self.inv_freq)
- emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
-
- self._cos_cached = emb.cos()[None, None, :, :]
- self._sin_cached = emb.sin()[None, None, :, :]
-
- return self._cos_cached, self._sin_cached
-
- def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
- self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k, seq_dimension=-2)
-
- return (
- apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
- apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
- )
-
-
-class EsmContactPredictionHead(nn.Module):
- """Performs symmetrization, apc, and computes a logistic regression on the output features"""
-
- def __init__(
- self,
- in_features: int,
- bias=True,
- eos_idx: int = 2,
- ):
- super().__init__()
- self.in_features = in_features
- self.eos_idx = eos_idx
- self.regression = nn.Linear(in_features, 1, bias)
- self.activation = nn.Sigmoid()
-
- def forward(self, tokens, attentions):
- # remove eos token attentions
- eos_mask = tokens.ne(self.eos_idx).to(attentions)
- eos_mask = eos_mask.unsqueeze(1) * eos_mask.unsqueeze(2)
- attentions = attentions * eos_mask[:, None, None, :, :]
- attentions = attentions[..., :-1, :-1]
- # remove cls token attentions
- attentions = attentions[..., 1:, 1:]
- batch_size, layers, heads, seqlen, _ = attentions.size()
- attentions = attentions.view(batch_size, layers * heads, seqlen, seqlen)
-
- # features: batch x channels x tokens x tokens (symmetric)
- attentions = attentions.to(
- self.regression.weight.device
- ) # attentions always float32, may need to convert to float16
- attentions = average_product_correct(symmetrize(attentions))
- attentions = attentions.permute(0, 2, 3, 1)
- return self.activation(self.regression(attentions).squeeze(3))
-
-
-class EsmEmbeddings(nn.Module):
- """
- Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
- """
-
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
-
- if config.emb_layer_norm_before:
- self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- else:
- self.layer_norm = None
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
-
- self.padding_idx = config.pad_token_id
- self.position_embeddings = nn.Embedding(
- config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
- )
- self.token_dropout = config.token_dropout
- self.mask_token_id = config.mask_token_id
-
- def forward(
- self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
- ):
- if position_ids is None:
- if input_ids is not None:
- # Create the position ids from the input token ids. Any padded tokens remain padded.
- position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)
- else:
- position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
-
- # Note that if we want to support ESM-1 (not 1b!) in future then we need to support an
- # embedding_scale factor here.
- embeddings = inputs_embeds
-
- # Matt: ESM has the option to handle masking in MLM in a slightly unusual way. If the token_dropout
- # flag is False then it is handled in the same was as BERT/RoBERTa. If it is set to True, however,
- # masked tokens are treated as if they were selected for input dropout and zeroed out.
- # This "mask-dropout" is compensated for when masked tokens are not present, by scaling embeddings by
- # a factor of (fraction of unmasked tokens during training) / (fraction of unmasked tokens in sample).
- # This is analogous to the way that dropout layers scale down outputs during evaluation when not
- # actually dropping out values (or, equivalently, scale up their un-dropped outputs in training).
- if self.token_dropout:
- embeddings = embeddings.masked_fill((input_ids == self.mask_token_id).unsqueeze(-1), 0.0)
- mask_ratio_train = 0.15 * 0.8 # Hardcoded as the ratio used in all ESM model training runs
- src_lengths = attention_mask.sum(-1)
- mask_ratio_observed = (input_ids == self.mask_token_id).sum(-1).float() / src_lengths
- embeddings = (embeddings * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]).to(
- embeddings.dtype
- )
-
- if self.position_embedding_type == "absolute":
- position_embeddings = self.position_embeddings(position_ids)
- embeddings = embeddings + position_embeddings
-
- if self.layer_norm is not None:
- embeddings = self.layer_norm(embeddings)
- if attention_mask is not None:
- embeddings = (embeddings * attention_mask.unsqueeze(-1)).to(embeddings.dtype)
- # Matt: I think this line was copied incorrectly from BERT, disabling it for now.
- # embeddings = self.dropout(embeddings)
- return embeddings
-
- def create_position_ids_from_inputs_embeds(self, inputs_embeds):
- """
- We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
-
- Args:
- inputs_embeds: torch.Tensor
-
- Returns: torch.Tensor
- """
- input_shape = inputs_embeds.size()[:-1]
- sequence_length = input_shape[1]
-
- position_ids = torch.arange(
- self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
- )
- return position_ids.unsqueeze(0).expand(input_shape)
-
-
-class EsmSelfAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size)
- self.key = nn.Linear(config.hidden_size, self.all_head_size)
- self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
- self.position_embedding_type = position_embedding_type or getattr(
- config, "position_embedding_type", "absolute"
- )
- self.rotary_embeddings = None
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- self.max_position_embeddings = config.max_position_embeddings
- self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
- elif self.position_embedding_type == "rotary":
- self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size)
-
- self.is_decoder = config.is_decoder
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- mixed_query_layer = self.query(hidden_states)
-
- # If this is instantiated as a cross-attention module, the keys
- # and values come from an encoder; the attention mask needs to be
- # such that the encoder's padding tokens are not attended to.
- is_cross_attention = encoder_hidden_states is not None
-
- if is_cross_attention and past_key_value is not None:
- # reuse k,v, cross_attentions
- key_layer = past_key_value[0]
- value_layer = past_key_value[1]
- attention_mask = encoder_attention_mask
- elif is_cross_attention:
- key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
- value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
- attention_mask = encoder_attention_mask
- elif past_key_value is not None:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
- value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
- else:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Matt: Our BERT model (which this code was derived from) scales attention logits down by sqrt(head_dim).
- # ESM scales the query down by the same factor instead. Modulo numerical stability these are equivalent,
- # but not when rotary embeddings get involved. Therefore, we scale the query here to match the original
- # ESM code and fix rotary embeddings.
- query_layer = query_layer * self.attention_head_size**-0.5
-
- if self.is_decoder:
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_layer, value_layer)
-
- if self.position_embedding_type == "rotary":
- query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- seq_length = hidden_states.size()[1]
- position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
- position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
- distance = position_ids_l - position_ids_r
- positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
- positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
-
- if self.position_embedding_type == "relative_key":
- relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores
- elif self.position_embedding_type == "relative_key_query":
- relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
-
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in EsmModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs.to(value_layer.dtype), value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- if self.is_decoder:
- outputs = outputs + (past_key_value,)
- return outputs
-
-
-class EsmSelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = hidden_states + input_tensor
- return hidden_states
-
-
-class EsmAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.self = EsmSelfAttention(config)
- self.output = EsmSelfOutput(config)
- self.pruned_heads = set()
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- past_key_value=None,
- output_attentions=False,
- ):
- hidden_states_ln = self.LayerNorm(hidden_states)
- self_outputs = self.self(
- hidden_states_ln,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-class EsmIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = gelu(hidden_states)
- return hidden_states
-
-
-class EsmOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = hidden_states + input_tensor
- return hidden_states
-
-
-class EsmLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = EsmAttention(config)
- self.is_decoder = config.is_decoder
- self.add_cross_attention = config.add_cross_attention
- if self.add_cross_attention:
- if not self.is_decoder:
- raise RuntimeError(f"{self} should be used as a decoder model if cross attention is added")
- self.crossattention = EsmAttention(config)
- self.intermediate = EsmIntermediate(config)
- self.output = EsmOutput(config)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- past_key_value=None,
- output_attentions=False,
- ):
- # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
- self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
- self_attention_outputs = self.attention(
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions=output_attentions,
- past_key_value=self_attn_past_key_value,
- )
- attention_output = self_attention_outputs[0]
-
- # if decoder, the last output is tuple of self-attn cache
- if self.is_decoder:
- outputs = self_attention_outputs[1:-1]
- present_key_value = self_attention_outputs[-1]
- else:
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- cross_attn_present_key_value = None
- if self.is_decoder and encoder_hidden_states is not None:
- if not hasattr(self, "crossattention"):
- raise AttributeError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated"
- " with cross-attention layers by setting `config.add_cross_attention=True`"
- )
-
- # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
- cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
- cross_attention_outputs = self.crossattention(
- attention_output,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- cross_attn_past_key_value,
- output_attentions,
- )
- attention_output = cross_attention_outputs[0]
- outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
-
- # add cross-attn cache to positions 3,4 of present_key_value tuple
- cross_attn_present_key_value = cross_attention_outputs[-1]
- present_key_value = present_key_value + cross_attn_present_key_value
-
- layer_output = self.feed_forward_chunk(attention_output)
-
- outputs = (layer_output,) + outputs
-
- # if decoder, return the attn key/values as the last output
- if self.is_decoder:
- outputs = outputs + (present_key_value,)
- return outputs
-
- def feed_forward_chunk(self, attention_output):
- attention_output_ln = self.LayerNorm(attention_output)
- intermediate_output = self.intermediate(attention_output_ln)
- layer_output = self.output(intermediate_output, attention_output)
- return layer_output
-
-
-class EsmEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([EsmLayer(config) for _ in range(config.num_hidden_layers)])
- self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- past_key_values=None,
- use_cache=None,
- output_attentions=False,
- output_hidden_states=False,
- return_dict=True,
- ):
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
- "`use_cache=False`..."
- )
- use_cache = False
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
-
- next_decoder_cache = () if use_cache else None
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
- past_key_value = past_key_values[i] if past_key_values is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- )
-
- hidden_states = layer_outputs[0]
- if use_cache:
- next_decoder_cache = next_decoder_cache + (layer_outputs[-1],)
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
-
- if self.emb_layer_norm_after:
- hidden_states = self.emb_layer_norm_after(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- next_decoder_cache,
- all_hidden_states,
- all_self_attentions,
- all_cross_attentions,
- ]
- if v is not None
- )
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=next_decoder_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPooler
-class EsmPooler(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = nn.Tanh()
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(first_token_tensor)
- pooled_output = self.activation(pooled_output)
- return pooled_output
-
-
-class EsmPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = EsmConfig
- base_model_prefix = "esm"
- supports_gradient_checkpointing = True
- _no_split_modules = ["EsmLayer", "EsmFoldTriangularSelfAttentionBlock", "EsmEmbeddings"]
-
- # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-ESM_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`EsmConfig`]): Model configuration class with all the parameters of the
- model. Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ESM_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare ESM Model transformer outputting raw hidden-states without any specific head on top.",
- ESM_START_DOCSTRING,
-)
-class EsmModel(EsmPreTrainedModel):
- """
-
- The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
- cross-attention is added between the self-attention layers, following the architecture described in [Attention is
- all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
- Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
-
- To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
- to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
- `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
- """
-
- def __init__(self, config, add_pooling_layer=True):
- super().__init__(config)
- self.config = config
-
- self.embeddings = EsmEmbeddings(config)
- self.encoder = EsmEncoder(config)
-
- self.pooler = EsmPooler(config) if add_pooling_layer else None
-
- self.contact_head = EsmContactPredictionHead(
- in_features=config.num_hidden_layers * config.num_attention_heads, bias=True
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(ESM_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPoolingAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
- r"""
- encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if self.config.is_decoder:
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- else:
- use_cache = False
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- # past_key_values_length
- past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
-
- if attention_mask is None:
- attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
-
- # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
- # ourselves in which case we just need to make it broadcastable to all heads.
- extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
-
- # If a 2D or 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- if self.config.is_decoder and encoder_hidden_states is not None:
- encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
- encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
- if encoder_attention_mask is None:
- encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
- encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
- else:
- encoder_extended_attention_mask = None
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- attention_mask=attention_mask,
- inputs_embeds=inputs_embeds,
- past_key_values_length=past_key_values_length,
- )
- encoder_outputs = self.encoder(
- embedding_output,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_extended_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- return (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPoolingAndCrossAttentions(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- past_key_values=encoder_outputs.past_key_values,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- cross_attentions=encoder_outputs.cross_attentions,
- )
-
- def predict_contacts(self, tokens, attention_mask):
- attns = self(tokens, attention_mask=attention_mask, return_dict=True, output_attentions=True).attentions
- attns = torch.stack(attns, dim=1) # Matches the original model layout
- # In the original model, attentions for padding tokens are completely zeroed out.
- # This makes no difference most of the time because the other tokens won't attend to them,
- # but it does for the contact prediction task, which takes attentions as input,
- # so we have to mimic that here.
- attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3)
- attns *= attention_mask.unsqueeze(1).unsqueeze(2).unsqueeze(4)
- return self.contact_head(tokens, attns)
-
-
-@add_start_docstrings("""ESM Model with a `language modeling` head on top.""", ESM_START_DOCSTRING)
-class EsmForMaskedLM(EsmPreTrainedModel):
- _tied_weights_keys = ["lm_head.decoder.weight"]
-
- def __init__(self, config):
- super().__init__(config)
-
- if config.is_decoder:
- logger.warning(
- "If you want to use `EsmForMaskedLM` make sure `config.is_decoder=False` for "
- "bi-directional self-attention."
- )
-
- self.esm = EsmModel(config, add_pooling_layer=False)
- self.lm_head = EsmLMHead(config)
-
- self.init_weights()
-
- def get_output_embeddings(self):
- return self.lm_head.decoder
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(ESM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="",
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.FloatTensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- kwargs (`Dict[str, any]`, optional, defaults to *{}*):
- Used to hide legacy arguments that have been deprecated.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.esm(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = outputs[0]
- prediction_scores = self.lm_head(sequence_output)
-
- masked_lm_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
-
- labels = labels.to(prediction_scores.device)
- masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return MaskedLMOutput(
- loss=masked_lm_loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def predict_contacts(self, tokens, attention_mask):
- return self.esm.predict_contacts(tokens, attention_mask=attention_mask)
-
-
-class EsmLMHead(nn.Module):
- """ESM Head for masked language modeling."""
-
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
- self.bias = nn.Parameter(torch.zeros(config.vocab_size))
-
- def forward(self, features, **kwargs):
- x = self.dense(features)
- x = gelu(x)
- x = self.layer_norm(x)
-
- # project back to size of vocabulary with bias
- x = self.decoder(x) + self.bias
- return x
-
-
-@add_start_docstrings(
- """
- ESM Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
- output) e.g. for GLUE tasks.
- """,
- ESM_START_DOCSTRING,
-)
-class EsmForSequenceClassification(EsmPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
-
- self.esm = EsmModel(config, add_pooling_layer=False)
- self.classifier = EsmClassificationHead(config)
-
- self.init_weights()
-
- @add_start_docstrings_to_model_forward(ESM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.esm(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = outputs[0]
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
-
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- ESM Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- ESM_START_DOCSTRING,
-)
-class EsmForTokenClassification(EsmPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.esm = EsmModel(config, add_pooling_layer=False)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- self.init_weights()
-
- @add_start_docstrings_to_model_forward(ESM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.esm(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
-
- labels = labels.to(logits.device)
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-class EsmClassificationHead(nn.Module):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
-
- def forward(self, features, **kwargs):
- x = features[:, 0, :] # take token (equiv. to [CLS])
- x = self.dropout(x)
- x = self.dense(x)
- x = torch.tanh(x)
- x = self.dropout(x)
- x = self.out_proj(x)
- return x
-
-
-def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
- """
- Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
- are ignored. This is modified from fairseq's `utils.make_positions`.
-
- Args:
- x: torch.Tensor x:
-
- Returns: torch.Tensor
- """
- # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
- mask = input_ids.ne(padding_idx).int()
- incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
- return incremental_indices.long() + padding_idx
diff --git a/transformers/models/esm/modeling_esmfold.py b/transformers/models/esm/modeling_esmfold.py
deleted file mode 100644
index 3aaf811960721b55d5e10a28a4e3be5aaeed1ec7..0000000000000000000000000000000000000000
--- a/transformers/models/esm/modeling_esmfold.py
+++ /dev/null
@@ -1,2322 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import math
-import sys
-from dataclasses import dataclass
-from functools import partial
-from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.nn import LayerNorm
-
-from ...integrations.deepspeed import is_deepspeed_available
-from ...modeling_outputs import ModelOutput
-from ...utils import (
- ContextManagers,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_scipy_available,
- logging,
- replace_return_docstrings,
-)
-from .configuration_esm import EsmConfig
-from .modeling_esm import ESM_START_DOCSTRING, EsmModel, EsmPreTrainedModel
-from .openfold_utils import (
- OFProtein,
- Rigid,
- Rotation,
- atom14_to_atom37,
- chunk_layer,
- compute_predicted_aligned_error,
- compute_tm,
- frames_and_literature_positions_to_atom14_pos,
- make_atom14_masks,
- residue_constants,
- to_pdb,
- torsion_angles_to_frames,
-)
-
-
-logger = logging.get_logger(__name__)
-_CHECKPOINT_FOR_DOC = "facebook/esmfold_v1"
-_CONFIG_FOR_DOC = "EsmConfig"
-
-
-@dataclass
-class EsmForProteinFoldingOutput(ModelOutput):
- """
- Output type of [`EsmForProteinFoldingOutput`].
-
- Args:
- frames (`torch.FloatTensor`):
- Output frames.
- sidechain_frames (`torch.FloatTensor`):
- Output sidechain frames.
- unnormalized_angles (`torch.FloatTensor`):
- Predicted unnormalized backbone and side chain torsion angles.
- angles (`torch.FloatTensor`):
- Predicted backbone and side chain torsion angles.
- positions (`torch.FloatTensor`):
- Predicted positions of the backbone and side chain atoms.
- states (`torch.FloatTensor`):
- Hidden states from the protein folding trunk.
- s_s (`torch.FloatTensor`):
- Per-residue embeddings derived by concatenating the hidden states of each layer of the ESM-2 LM stem.
- s_z (`torch.FloatTensor`):
- Pairwise residue embeddings.
- distogram_logits (`torch.FloatTensor`):
- Input logits to the distogram used to compute residue distances.
- lm_logits (`torch.FloatTensor`):
- Logits output by the ESM-2 protein language model stem.
- aatype (`torch.FloatTensor`):
- Input amino acids (AlphaFold2 indices).
- atom14_atom_exists (`torch.FloatTensor`):
- Whether each atom exists in the atom14 representation.
- residx_atom14_to_atom37 (`torch.FloatTensor`):
- Mapping between atoms in the atom14 and atom37 representations.
- residx_atom37_to_atom14 (`torch.FloatTensor`):
- Mapping between atoms in the atom37 and atom14 representations.
- atom37_atom_exists (`torch.FloatTensor`):
- Whether each atom exists in the atom37 representation.
- residue_index (`torch.FloatTensor`):
- The index of each residue in the protein chain. Unless internal padding tokens are used, this will just be
- a sequence of integers from 0 to `sequence_length`.
- lddt_head (`torch.FloatTensor`):
- Raw outputs from the lddt head used to compute plddt.
- plddt (`torch.FloatTensor`):
- Per-residue confidence scores. Regions of low confidence may indicate areas where the model's prediction is
- uncertain, or where the protein structure is disordered.
- ptm_logits (`torch.FloatTensor`):
- Raw logits used for computing ptm.
- ptm (`torch.FloatTensor`):
- TM-score output representing the model's high-level confidence in the overall structure.
- aligned_confidence_probs (`torch.FloatTensor`):
- Per-residue confidence scores for the aligned structure.
- predicted_aligned_error (`torch.FloatTensor`):
- Predicted error between the model's prediction and the ground truth.
- max_predicted_aligned_error (`torch.FloatTensor`):
- Per-sample maximum predicted error.
- """
-
- frames: torch.FloatTensor = None
- sidechain_frames: torch.FloatTensor = None
- unnormalized_angles: torch.FloatTensor = None
- angles: torch.FloatTensor = None
- positions: torch.FloatTensor = None
- states: torch.FloatTensor = None
- s_s: torch.FloatTensor = None
- s_z: torch.FloatTensor = None
- distogram_logits: torch.FloatTensor = None
- lm_logits: torch.FloatTensor = None
- aatype: torch.FloatTensor = None
- atom14_atom_exists: torch.FloatTensor = None
- residx_atom14_to_atom37: torch.FloatTensor = None
- residx_atom37_to_atom14: torch.FloatTensor = None
- atom37_atom_exists: torch.FloatTensor = None
- residue_index: torch.FloatTensor = None
- lddt_head: torch.FloatTensor = None
- plddt: torch.FloatTensor = None
- ptm_logits: torch.FloatTensor = None
- ptm: torch.FloatTensor = None
- aligned_confidence_probs: torch.FloatTensor = None
- predicted_aligned_error: torch.FloatTensor = None
- max_predicted_aligned_error: torch.FloatTensor = None
-
-
-ESMFOLD_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- masking_pattern (`torch.LongTensor` of shape `({0})`, *optional*):
- Locations of tokens to mask during training as a form of regularization. Mask values selected in `[0, 1]`.
- num_recycles (`int`, *optional*, defaults to `None`):
- Number of times to recycle the input sequence. If `None`, defaults to `config.num_recycles`. "Recycling"
- consists of passing the output of the folding trunk back in as input to the trunk. During training, the
- number of recycles should vary with each batch, to ensure that the model learns to output valid predictions
- after each recycle. During inference, num_recycles should be set to the highest value that the model was
- trained with for maximum accuracy. Accordingly, when this value is set to `None`, config.max_recycles is
- used.
-"""
-
-
-def is_fp16_enabled():
- # Autocast world
- fp16_enabled = torch.get_autocast_gpu_dtype() == torch.float16
- fp16_enabled = fp16_enabled and torch.is_autocast_enabled()
-
- return fp16_enabled
-
-
-def is_deepspeed_initialized():
- if is_deepspeed_available():
- return False
- else:
- try:
- import deepspeed
-
- # This is not available in all DeepSpeed versions.
- return deepspeed.utils.is_initialized()
- except Exception:
- return False
-
-
-def collate_dense_tensors(samples: List[torch.Tensor], pad_v: float = 0) -> torch.Tensor:
- """
- Takes a list of tensors with the following dimensions:
- [(d_11, ..., d_1K),
- (d_21, ..., d_2K), ..., (d_N1, ..., d_NK)]
- and stack + pads them into a single tensor of:
- (N, max_i=1,N { d_i1 }, ..., max_i=1,N {diK})
- """
- if len(samples) == 0:
- return torch.Tensor()
- if len({x.dim() for x in samples}) != 1:
- raise RuntimeError(f"Samples has varying dimensions: {[x.dim() for x in samples]}")
- (device,) = tuple({x.device for x in samples}) # assumes all on same device
- max_shape = [max(lst) for lst in zip(*[x.shape for x in samples])]
- result = torch.empty(len(samples), *max_shape, dtype=samples[0].dtype, device=device)
- result.fill_(pad_v)
- for i in range(len(samples)):
- result_i = result[i]
- t = samples[i]
- result_i[tuple(slice(0, k) for k in t.shape)] = t
- return result
-
-
-def flatten_final_dims(t: torch.Tensor, no_dims: int):
- return t.reshape(t.shape[:-no_dims] + (-1,))
-
-
-def permute_final_dims(tensor: torch.Tensor, inds: List[int]):
- zero_index = -1 * len(inds)
- first_inds = list(range(len(tensor.shape[:zero_index])))
- return tensor.permute(first_inds + [zero_index + i for i in inds])
-
-
-def dict_multimap(fn, dicts):
- first = dicts[0]
- new_dict = {}
- for k, v in first.items():
- all_v = [d[k] for d in dicts]
- if isinstance(v, dict):
- new_dict[k] = dict_multimap(fn, all_v)
- else:
- new_dict[k] = fn(all_v)
-
- return new_dict
-
-
-def trunc_normal_init_(weights, scale=1.0, fan="fan_in"):
- shape = weights.shape
- scale = scale / max(1, shape[1])
-
- if not is_scipy_available():
- logger.warning(
- "This init requires scipy, but scipy was not found, default to an approximation that might not be"
- " equivalent."
- )
- std = math.sqrt(scale)
- torch.nn.init.normal_(weights, std=std).clamp(min=0.0, max=2.0 * std)
-
- else:
- from scipy.stats import truncnorm
-
- std = math.sqrt(scale) / truncnorm.std(a=-2, b=2, loc=0, scale=1)
- samples = truncnorm.rvs(a=-2, b=2, loc=0, scale=std, size=weights.numel())
- samples = np.reshape(samples, shape)
- weights.copy_(torch.tensor(samples, device=weights.device))
-
-
-def ipa_point_weights_init_(weights):
- with torch.no_grad():
- softplus_inverse_1 = 0.541324854612918
- weights.fill_(softplus_inverse_1)
-
-
-class EsmFoldLinear(nn.Linear):
- """
- A Linear layer with built-in nonstandard initializations. Called just like torch.nn.Linear.
-
- Implements the initializers in 1.11.4, plus some additional ones found in the code.
- """
-
- def __init__(
- self,
- in_dim: int,
- out_dim: int,
- bias: bool = True,
- init: str = "default",
- init_fn: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None,
- ):
- """
- Args:
- in_dim:
- The final dimension of inputs to the layer
- out_dim:
- The final dimension of layer outputs
- bias:
- Whether to learn an additive bias. True by default
- init:
- The initializer to use. Choose from:
-
- "default": LeCun fan-in truncated normal initialization "relu": He initialization w/ truncated normal
- distribution "glorot": Fan-average Glorot uniform initialization "gating": Weights=0, Bias=1 "normal":
- Normal initialization with std=1/sqrt(fan_in) "final": Weights=0, Bias=0
-
- Overridden by init_fn if the latter is not None.
- init_fn:
- A custom initializer taking weight and bias as inputs. Overrides init if not None.
- """
- super().__init__(in_dim, out_dim, bias=bias)
-
- if bias:
- with torch.no_grad():
- self.bias.fill_(0)
- self.init = init
- self.init_fn = init_fn
-
- if init not in ["default", "relu", "glorot", "gating", "normal", "final"]:
- raise ValueError("Invalid init string.")
-
-
-class EsmFoldLayerNorm(nn.Module):
- def __init__(self, c_in, eps=1e-5):
- super().__init__()
-
- self.c_in = (c_in,)
- self.eps = eps
-
- self.weight = nn.Parameter(torch.ones(c_in))
- self.bias = nn.Parameter(torch.zeros(c_in))
-
- def forward(self, x):
- d = x.dtype
- if d is torch.bfloat16 and not is_deepspeed_initialized():
- with torch.cuda.amp.autocast(enabled=False):
- out = nn.functional.layer_norm(x, self.c_in, self.weight.to(dtype=d), self.bias.to(dtype=d), self.eps)
- else:
- out = nn.functional.layer_norm(x, self.c_in, self.weight, self.bias, self.eps)
-
- return out
-
-
-@torch.jit.ignore
-def softmax_no_cast(t: torch.Tensor, dim: int = -1) -> torch.Tensor:
- """
- Softmax, but without automatic casting to fp32 when the input is of type bfloat16
- """
- d = t.dtype
- if d is torch.bfloat16 and not is_deepspeed_initialized():
- with torch.cuda.amp.autocast(enabled=False):
- s = torch.nn.functional.softmax(t, dim=dim)
- else:
- s = torch.nn.functional.softmax(t, dim=dim)
-
- return s
-
-
-class EsmFoldAttention(nn.Module):
- """
- Standard multi-head attention using AlphaFold's default layer initialization. Allows multiple bias vectors.
- """
-
- def __init__(
- self,
- c_q: int,
- c_k: int,
- c_v: int,
- c_hidden: int,
- no_heads: int,
- gating: bool = True,
- ):
- """
- Args:
- c_q:
- Input dimension of query data
- c_k:
- Input dimension of key data
- c_v:
- Input dimension of value data
- c_hidden:
- Per-head hidden dimension
- no_heads:
- Number of attention heads
- gating:
- Whether the output should be gated using query data
- """
- super().__init__()
-
- self.c_q = c_q
- self.c_k = c_k
- self.c_v = c_v
- self.c_hidden = c_hidden
- self.no_heads = no_heads
- self.gating = gating
-
- # DISCREPANCY: c_hidden is not the per-head channel dimension, as
- # stated in the supplement, but the overall channel dimension.
-
- self.linear_q = EsmFoldLinear(self.c_q, self.c_hidden * self.no_heads, bias=False, init="glorot")
- self.linear_k = EsmFoldLinear(self.c_k, self.c_hidden * self.no_heads, bias=False, init="glorot")
- self.linear_v = EsmFoldLinear(self.c_v, self.c_hidden * self.no_heads, bias=False, init="glorot")
- self.linear_o = EsmFoldLinear(self.c_hidden * self.no_heads, self.c_q, init="final")
-
- self.linear_g = None
- if self.gating:
- self.linear_g = EsmFoldLinear(self.c_q, self.c_hidden * self.no_heads, init="gating")
-
- self.sigmoid = nn.Sigmoid()
-
- def _prep_qkv(self, q_x: torch.Tensor, kv_x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- # [*, Q/K/V, H * C_hidden]
- q = self.linear_q(q_x)
- k = self.linear_k(kv_x)
- v = self.linear_v(kv_x)
-
- # [*, Q/K, H, C_hidden]
- q = q.view(q.shape[:-1] + (self.no_heads, -1))
- k = k.view(k.shape[:-1] + (self.no_heads, -1))
- v = v.view(v.shape[:-1] + (self.no_heads, -1))
-
- # [*, H, Q/K, C_hidden]
- q = q.transpose(-2, -3)
- k = k.transpose(-2, -3)
- v = v.transpose(-2, -3)
-
- q /= math.sqrt(self.c_hidden)
-
- return q, k, v
-
- def _wrap_up(self, o: torch.Tensor, q_x: torch.Tensor) -> torch.Tensor:
- if self.linear_g is not None:
- g = self.sigmoid(self.linear_g(q_x))
-
- # [*, Q, H, C_hidden]
- g = g.view(g.shape[:-1] + (self.no_heads, -1))
- o = o * g
-
- # [*, Q, H * C_hidden]
- o = flatten_final_dims(o, 2)
-
- # [*, Q, C_q]
- o = self.linear_o(o)
-
- return o
-
- def forward(
- self,
- q_x: torch.Tensor,
- kv_x: torch.Tensor,
- biases: Optional[List[torch.Tensor]] = None,
- use_memory_efficient_kernel: bool = False,
- use_lma: bool = False,
- lma_q_chunk_size: int = 1024,
- lma_kv_chunk_size: int = 4096,
- use_flash: bool = False,
- flash_mask: Optional[torch.Tensor] = None,
- ) -> torch.Tensor:
- """
- Args:
- q_x:
- [*, Q, C_q] query data
- kv_x:
- [*, K, C_k] key data
- biases:
- List of biases that broadcast to [*, H, Q, K]
- use_memory_efficient_kernel:
- Whether to use a custom memory-efficient attention kernel. This should be the default choice for most.
- If none of the "use_<...>" flags are True, a stock PyTorch implementation is used instead
- use_lma:
- Whether to use low-memory attention (Staats & Rabe 2021). If none of the "use_<...>" flags are True, a
- stock PyTorch implementation is used instead
- lma_q_chunk_size:
- Query chunk size (for LMA)
- lma_kv_chunk_size:
- Key/Value chunk size (for LMA)
- Returns
- [*, Q, C_q] attention update
- """
- if use_lma and (lma_q_chunk_size is None or lma_kv_chunk_size is None):
- raise ValueError("If use_lma is specified, lma_q_chunk_size and lma_kv_chunk_size must be provided")
-
- if use_flash and biases is not None:
- raise ValueError("use_flash is incompatible with the bias option. For masking, use flash_mask instead")
-
- attn_options = [use_memory_efficient_kernel, use_lma, use_flash]
- if sum(attn_options) > 1:
- raise ValueError("Choose at most one alternative attention algorithm")
-
- if biases is None:
- biases = []
-
- # [*, H, Q/K, C_hidden]
- query, key, value = self._prep_qkv(q_x, kv_x)
- key = permute_final_dims(key, (1, 0))
-
- # [*, H, Q, K]
- output = torch.matmul(query, key)
- for b in biases:
- output += b
- output = softmax_no_cast(output, -1)
-
- # [*, H, Q, C_hidden]
- output = torch.matmul(output, value)
- output = output.transpose(-2, -3)
- output = self._wrap_up(output, q_x)
-
- return output
-
-
-class EsmFoldTriangleAttention(nn.Module):
- def __init__(self, c_in, c_hidden, no_heads, starting=True, inf=1e9):
- """
- Args:
- c_in:
- Input channel dimension
- c_hidden:
- Overall hidden channel dimension (not per-head)
- no_heads:
- Number of attention heads
- """
- super().__init__()
-
- self.c_in = c_in
- self.c_hidden = c_hidden
- self.no_heads = no_heads
- self.starting = starting
- self.inf = inf
-
- self.layer_norm = LayerNorm(self.c_in)
-
- self.linear = EsmFoldLinear(c_in, self.no_heads, bias=False, init="normal")
-
- self.mha = EsmFoldAttention(self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads)
-
- @torch.jit.ignore
- def _chunk(
- self,
- x: torch.Tensor,
- biases: List[torch.Tensor],
- chunk_size: int,
- use_memory_efficient_kernel: bool = False,
- use_lma: bool = False,
- inplace_safe: bool = False,
- ) -> torch.Tensor:
- "triangle! triangle!"
- mha_inputs = {
- "q_x": x,
- "kv_x": x,
- "biases": biases,
- }
-
- return chunk_layer(
- partial(self.mha, use_memory_efficient_kernel=use_memory_efficient_kernel, use_lma=use_lma),
- mha_inputs,
- chunk_size=chunk_size,
- no_batch_dims=len(x.shape[:-2]),
- _out=x if inplace_safe else None,
- )
-
- def forward(
- self,
- x: torch.Tensor,
- mask: Optional[torch.Tensor] = None,
- chunk_size: Optional[int] = None,
- use_memory_efficient_kernel: bool = False,
- use_lma: bool = False,
- inplace_safe: bool = False,
- ) -> torch.Tensor:
- """
- Args:
- x:
- [*, I, J, C_in] input tensor (e.g. the pair representation)
- Returns:
- [*, I, J, C_in] output tensor
- """
- if mask is None:
- # [*, I, J]
- mask = x.new_ones(
- x.shape[:-1],
- )
-
- if not self.starting:
- x = x.transpose(-2, -3)
- mask = mask.transpose(-1, -2)
-
- # [*, I, J, C_in]
- x = self.layer_norm(x)
-
- # [*, I, 1, 1, J]
- mask_bias = (self.inf * (mask - 1))[..., :, None, None, :]
-
- # [*, H, I, J]
- triangle_bias = permute_final_dims(self.linear(x), (2, 0, 1))
-
- # [*, 1, H, I, J]
- triangle_bias = triangle_bias.unsqueeze(-4)
-
- biases = [mask_bias, triangle_bias]
-
- if chunk_size is not None:
- x = self._chunk(
- x,
- biases,
- chunk_size,
- use_memory_efficient_kernel=use_memory_efficient_kernel,
- use_lma=use_lma,
- inplace_safe=inplace_safe,
- )
- else:
- x = self.mha(
- q_x=x, kv_x=x, biases=biases, use_memory_efficient_kernel=use_memory_efficient_kernel, use_lma=use_lma
- )
-
- if not self.starting:
- x = x.transpose(-2, -3)
-
- return x
-
-
-class EsmFoldTriangleMultiplicativeUpdate(nn.Module):
- """
- Implements Algorithms 11 and 12.
- """
-
- def __init__(self, config, _outgoing=True):
- super().__init__()
- c_hidden = config.pairwise_state_dim
- self._outgoing = _outgoing
-
- self.linear_a_p = EsmFoldLinear(c_hidden, c_hidden)
- self.linear_a_g = EsmFoldLinear(c_hidden, c_hidden, init="gating")
- self.linear_b_p = EsmFoldLinear(c_hidden, c_hidden)
- self.linear_b_g = EsmFoldLinear(c_hidden, c_hidden, init="gating")
- self.linear_g = EsmFoldLinear(c_hidden, c_hidden, init="gating")
- self.linear_z = EsmFoldLinear(c_hidden, c_hidden, init="final")
-
- self.layer_norm_in = LayerNorm(c_hidden)
- self.layer_norm_out = LayerNorm(c_hidden)
-
- self.sigmoid = nn.Sigmoid()
-
- def _combine_projections(
- self, a: torch.Tensor, b: torch.Tensor, _inplace_chunk_size: Optional[int] = None
- ) -> torch.Tensor:
- if self._outgoing:
- a = permute_final_dims(a, (2, 0, 1))
- b = permute_final_dims(b, (2, 1, 0))
- else:
- a = permute_final_dims(a, (2, 1, 0))
- b = permute_final_dims(b, (2, 0, 1))
-
- if _inplace_chunk_size is not None:
- # To be replaced by torch vmap
- for i in range(0, a.shape[-3], _inplace_chunk_size):
- a_chunk = a[..., i : i + _inplace_chunk_size, :, :]
- b_chunk = b[..., i : i + _inplace_chunk_size, :, :]
- a[..., i : i + _inplace_chunk_size, :, :] = torch.matmul(
- a_chunk,
- b_chunk,
- )
-
- p = a
- else:
- p = torch.matmul(a, b)
-
- return permute_final_dims(p, (1, 2, 0))
-
- def _inference_forward(
- self,
- z: torch.Tensor,
- mask: Optional[torch.Tensor] = None,
- inplace_chunk_size: Optional[int] = None,
- with_add: bool = True,
- ):
- """
- Args:
- z:
- A [*, N, N, C_z] pair representation
- mask:
- A [*, N, N] pair mask
- inplace_chunk_size:
- Size of chunks used in the main computation. Increase to trade memory for speed.
- with_add:
- If True, z is overwritten with (z + update). Otherwise, it is overwritten with (update).
- Returns:
- A reference to the overwritten z
-
- More memory-efficient, inference-only version of the forward function. Uses in-place operations, fusion of the
- addition that happens after this module in the Evoformer, a smidge of recomputation, and a cache of overwritten
- values to lower peak memory consumption of this module from 5x the size of the input tensor z to 2.5x its size.
- Useful for inference on extremely long sequences.
-
- It works as follows. We will make reference to variables used in the default forward implementation below.
- Naively, triangle multiplication attention requires the manifestation of 5 tensors the size of z: 1) z, the
- "square" input tensor, 2) a, the first projection of z, 3) b, the second projection of b, 4) g, a z-sized mask,
- and 5) a z-sized tensor for intermediate computations. For large N, this is prohibitively expensive; for
- N=4000, for example, z is more than 8GB alone. To avoid this problem, we compute b, g, and all intermediate
- tensors in small chunks, noting that the chunks required to compute a chunk of the output depend only on the
- tensor a and corresponding vertical and horizontal chunks of z. This suggests an algorithm that loops over
- pairs of chunks of z: hereafter "columns" and "rows" of z, even though each "column" and "row" in fact contains
- inplace_chunk_size contiguous true columns and rows of z. Writing output chunks to a new tensor would bring
- total memory consumption down to 3x the size of z. However, more memory can be saved by writing output chunks
- directly to z in-place. WLOG, we choose to write output chunks vertically, overwriting the ith "column" of z at
- the end of the ith iteration of the main loop. Despite this overwriting, the ith column is always one column
- ahead of previously overwritten columns and can be recovered directly from z. After the first iteration,
- however, the ith row of z is always at least partially overwritten. For this reason, we introduce the z-cache,
- a tensor one-half the size of z. The z-cache initially contains the left half (2nd and 3rd quadrants) of z. For
- 0 < i < N/2, the missing left part of the ith row of z is recovered from this cache at the beginning of the ith
- iteration. Once i exceeds n/2, the cache is "reoriented" to encompass the 3rd and 4th quadrants of z instead.
- Though the 3rd quadrant of the original z is entirely overwritten at this point, it can be recovered from the
- z-cache itself. Thereafter, the ith row of z can be recovered in its entirety from the reoriented z-cache.
- After the final iteration, z has been completely overwritten and contains the triangular multiplicative update.
- If with_add is True, it instead contains the sum of z and the triangular multiplicative update. In either case,
- peak memory consumption is just 2.5x the size of z, disregarding memory used for chunks and other small
- variables.
- """
- if mask is None:
- mask = z.new_ones(z.shape[:-1])
-
- mask = mask.unsqueeze(-1)
-
- def compute_projection_helper(pair, mask, a=True):
- if a:
- linear_g = self.linear_a_g
- linear_p = self.linear_a_p
- else:
- linear_g = self.linear_b_g
- linear_p = self.linear_b_p
-
- pair = self.layer_norm_in(pair)
- p = linear_g(pair)
- p.sigmoid_()
- p *= linear_p(pair)
- p *= mask
- p = permute_final_dims(p, (2, 0, 1))
- return p
-
- def compute_projection(pair, mask, a=True, chunked=True):
- need_transpose = self._outgoing ^ a
- if not chunked:
- p = compute_projection_helper(pair, mask, a)
- if need_transpose:
- p = p.transpose(-1, -2)
- else:
- # This computation is chunked so as not to exceed our 2.5x
- # budget with a large intermediate tensor
- linear_g = self.linear_a_g if a else self.linear_b_g
- c = linear_g.bias.shape[-1]
- out_shape = pair.shape[:-3] + (c,) + pair.shape[-3:-1]
- p = pair.new_zeros(out_shape)
- for i in range(0, pair.shape[-3], inplace_chunk_size):
- pair_chunk = pair[..., i : i + inplace_chunk_size, :, :]
- pair_chunk = compute_projection_helper(
- pair[..., i : i + inplace_chunk_size, :, :],
- mask[..., i : i + inplace_chunk_size, :, :],
- a,
- )
- if need_transpose:
- pair_chunk = pair_chunk.transpose(-1, -2)
- p[..., i : i + inplace_chunk_size] = pair_chunk
- else:
- p[..., i : i + inplace_chunk_size, :] = pair_chunk
-
- del pair_chunk
-
- return p
-
- # We start by fully manifesting a. In addition to the input, this
- # brings total memory consumption to 2x z (disregarding size of chunks)
- # [*, N, N, c]
- a = compute_projection(z, mask, True, chunked=True)
-
- if inplace_chunk_size is not None:
- n = a.shape[-1]
- half_n = n // 2 + n % 2
- row_dim = -3
- col_dim = -2
- b_chunk_dim = row_dim if self._outgoing else col_dim
-
- def empty_slicer(t):
- return [slice(None) for _ in t.shape]
-
- def slice_tensor(t, start, end, dim):
- # Slices start:end from the dim dimension of t
- s = empty_slicer(t)
- s[dim] = slice(start, end)
- return t[s]
-
- def flip_z_cache_(z_cache, z):
- # "Reorient" the z_cache (see below), filling it with quadrants
- # 3---recovered from the z_cache---and 4---recovered from z---
- # of the input tensor z.
- quadrant_3 = slice_tensor(z_cache, half_n, None, row_dim)
- z_cache = z_cache.transpose(row_dim, col_dim)
-
- # If n is odd, we need to shrink the z_cache by one row
- z_cache = z_cache[..., : (n // 2), :, :]
-
- # Move the 3rd quadrant of z into the
- first_half_slicer = empty_slicer(z_cache)
- first_half_slicer[col_dim] = slice(0, half_n)
- z_cache[first_half_slicer] = quadrant_3
-
- # Get the fourth quadrant of z
- quadrant_4 = slice_tensor(z, half_n, None, row_dim)
- quadrant_4 = slice_tensor(quadrant_4, half_n, None, col_dim)
-
- # Insert said quadrant into the rotated z-cache
- quadrant_3_slicer = empty_slicer(z_cache)
- quadrant_3_slicer[col_dim] = slice(half_n, None)
-
- z_cache[quadrant_3_slicer] = quadrant_4
-
- return z_cache
-
- # Initialize the z cache to the left half of z.
- z_cache_shape = list(z.shape)
- z_cache_shape[col_dim] = half_n
- z_cache = z.new_zeros(z_cache_shape)
- z_cache_slicer = empty_slicer(z_cache)
- z_cache_slicer[col_dim] = slice(0, half_n)
- z_cache.copy_(z[z_cache_slicer])
- z_cache_rotated = False
-
- # We need to reorient the z-cache at the halfway point, and we
- # don't want a single chunk to straddle that point. We contract one
- # of the chunks in the middle to address that problem.
- i_range = list(range(0, half_n, inplace_chunk_size))
- initial_offsets = [i_2 - i_1 for i_1, i_2 in zip(i_range, i_range[1:] + [half_n])]
- after_half = list(range(half_n, n, inplace_chunk_size))
- after_half_offsets = [inplace_chunk_size for _ in after_half]
- combined_range_with_offsets = zip(i_range + after_half, initial_offsets + after_half_offsets)
- for i, offset in combined_range_with_offsets:
- if not z_cache_rotated and i >= half_n:
- z_cache = flip_z_cache_(z_cache, z)
- z_cache_rotated = True
-
- z_chunk_b = slice_tensor(z, i, i + offset, b_chunk_dim)
- mask_chunk = slice_tensor(mask, i, i + offset, b_chunk_dim)
-
- z_chunk_b = z_chunk_b.clone()
- if b_chunk_dim == col_dim:
- z_chunk_b = slice_tensor(z, i, i + offset, col_dim)
- else: # b_chunk_dim == row_dim
- # In this case, the b-dimension (b_chunk_dim) is partially
- # overwritten at the end of each iteration. We need to
- # restore the missing component from the z-cache.
- if not z_cache_rotated:
- z_chunk_slicer = empty_slicer(z_chunk_b)
- z_chunk_slicer[col_dim] = slice(0, half_n)
- z_chunk_b[z_chunk_slicer] = slice_tensor(z_cache, i, i + offset, row_dim)
- else:
- z_cache_offset = i - half_n
- z_chunk_b = slice_tensor(z_cache, z_cache_offset, z_cache_offset + offset, row_dim)
-
- b_chunk = compute_projection(z_chunk_b, mask_chunk, a=False, chunked=False)
- del z_chunk_b
-
- x_chunk = torch.matmul(a, b_chunk)
- x_chunk = permute_final_dims(x_chunk, (1, 2, 0))
- x_chunk = self.layer_norm_out(x_chunk)
- x_chunk = self.linear_z(x_chunk)
-
- # The g dimension (col_dim) is parallel to and ahead of the
- # overwrites in z. We can extract the g chunk normally.
- z_chunk_g = slice_tensor(z, i, i + offset, col_dim)
- g_chunk = self.linear_g(self.layer_norm_in(z_chunk_g))
- g_chunk.sigmoid_()
- del z_chunk_g
-
- x_chunk *= g_chunk
-
- # Write the columns into z in-place
- z_slicer = empty_slicer(z)
- z_slicer[col_dim] = slice(i, i + offset)
- if with_add:
- z[z_slicer] += x_chunk
- else:
- z[z_slicer] = x_chunk
- else:
- b = compute_projection(z, mask, False, False)
- x = torch.matmul(a, b)
- x = self.layer_norm_out(x)
- x = self.linear_z(x)
- g = self.linear_g(z)
- g.sigmoid_()
- x *= g
- if with_add:
- z += x
- else:
- z = x
-
- return z
-
- def forward(
- self,
- z: torch.Tensor,
- mask: Optional[torch.Tensor] = None,
- inplace_safe: bool = False,
- _add_with_inplace: bool = False,
- _inplace_chunk_size: Optional[int] = 256,
- ) -> torch.Tensor:
- """
- Args:
- x:
- [*, N_res, N_res, C_z] input tensor
- mask:
- [*, N_res, N_res] input mask
- Returns:
- [*, N_res, N_res, C_z] output tensor
- """
- if inplace_safe:
- x = self._inference_forward(
- z,
- mask,
- inplace_chunk_size=_inplace_chunk_size,
- with_add=_add_with_inplace,
- )
- return x
-
- if mask is None:
- mask = z.new_ones(z.shape[:-1])
-
- mask = mask.unsqueeze(-1)
-
- z = self.layer_norm_in(z)
- a = mask
- a = a * self.sigmoid(self.linear_a_g(z))
- a = a * self.linear_a_p(z)
- b = mask
- b = b * self.sigmoid(self.linear_b_g(z))
- b = b * self.linear_b_p(z)
-
- if is_fp16_enabled():
- with torch.cuda.amp.autocast(enabled=False):
- x = self._combine_projections(a.float(), b.float())
- else:
- x = self._combine_projections(a, b)
-
- del a, b
- x = self.layer_norm_out(x)
- x = self.linear_z(x)
- g = self.sigmoid(self.linear_g(z))
- x = x * g
-
- return x
-
-
-class EsmFoldPreTrainedModel(EsmPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- # Subclass `EsMPreTrainedModel` to deal with special init
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, EsmFoldLinear):
- with torch.no_grad():
- if module.init_fn is not None:
- module.init_fn(module.weight, module.bias)
- elif module.init == "default":
- trunc_normal_init_(module.weight, scale=1.0)
- elif module.init == "relu":
- trunc_normal_init_(module.weight, scale=2.0)
- elif module.init == "glorot":
- nn.init.xavier_uniform_(module.weight, gain=1)
- elif module.init == "gating":
- module.weight.fill_(0.0)
- if module.bias:
- module.bias.fill_(1.0)
- elif module.init == "normal":
- torch.nn.init.kaiming_normal_(module.weight, nonlinearity="linear")
- elif module.init == "final":
- module.weight.fill_(0.0)
- elif isinstance(module, EsmFoldInvariantPointAttention):
- ipa_point_weights_init_(module.head_weights)
- elif isinstance(module, EsmFoldTriangularSelfAttentionBlock):
- torch.nn.init.zeros_(module.tri_mul_in.linear_z.weight)
- torch.nn.init.zeros_(module.tri_mul_in.linear_z.bias)
- torch.nn.init.zeros_(module.tri_mul_out.linear_z.weight)
- torch.nn.init.zeros_(module.tri_mul_out.linear_z.bias)
- torch.nn.init.zeros_(module.tri_att_start.mha.linear_o.weight)
- torch.nn.init.zeros_(module.tri_att_start.mha.linear_o.bias)
- torch.nn.init.zeros_(module.tri_att_end.mha.linear_o.weight)
- torch.nn.init.zeros_(module.tri_att_end.mha.linear_o.bias)
-
- torch.nn.init.zeros_(module.sequence_to_pair.o_proj.weight)
- torch.nn.init.zeros_(module.sequence_to_pair.o_proj.bias)
- torch.nn.init.zeros_(module.pair_to_sequence.linear.weight)
- torch.nn.init.zeros_(module.seq_attention.o_proj.weight)
- torch.nn.init.zeros_(module.seq_attention.o_proj.bias)
- torch.nn.init.zeros_(module.mlp_seq.mlp[-2].weight)
- torch.nn.init.zeros_(module.mlp_seq.mlp[-2].bias)
- torch.nn.init.zeros_(module.mlp_pair.mlp[-2].weight)
- torch.nn.init.zeros_(module.mlp_pair.mlp[-2].bias)
- else:
- super()._init_weights(module)
-
-
-class EsmFoldSelfAttention(nn.Module):
- def __init__(self, embed_dim, num_heads, head_width, gated=False):
- super().__init__()
- assert embed_dim == num_heads * head_width
-
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.head_width = head_width
-
- self.proj = nn.Linear(embed_dim, embed_dim * 3, bias=False)
- self.o_proj = nn.Linear(embed_dim, embed_dim, bias=True)
- self.gated = gated
- if gated:
- self.g_proj = nn.Linear(embed_dim, embed_dim)
- torch.nn.init.zeros_(self.g_proj.weight)
- torch.nn.init.ones_(self.g_proj.bias)
-
- self.rescale_factor = self.head_width**-0.5
-
- torch.nn.init.zeros_(self.o_proj.bias)
-
- def forward(self, x, mask=None, bias=None, indices=None):
- """
- Basic self attention with optional mask and external pairwise bias. To handle sequences of different lengths,
- use mask.
-
- Inputs:
- x: batch of input sequneces (.. x L x C) mask: batch of boolean masks where 1=valid, 0=padding position (..
- x L_k) bias: batch of scalar pairwise attention biases (.. x Lq x Lk x num_heads)
-
- Outputs:
- sequence projection (B x L x embed_dim), attention maps (B x L x L x num_heads)
- """
-
- t = self.proj(x).view(*x.shape[:2], self.num_heads, -1)
- t = t.permute(0, 2, 1, 3)
- q, k, v = t.chunk(3, dim=-1)
-
- q = self.rescale_factor * q
- a = torch.einsum("...qc,...kc->...qk", q, k)
-
- # Add external attention bias.
- if bias is not None:
- a = a + bias.permute(0, 3, 1, 2)
-
- # Do not attend to padding tokens.
- if mask is not None:
- mask = mask[:, None, None]
- a = a.masked_fill(mask == False, -np.inf) # noqa: E712
-
- a = nn.functional.softmax(a, dim=-1)
-
- y = torch.einsum("...hqk,...hkc->...qhc", a, v)
- y = y.reshape(*y.shape[:2], -1)
-
- if self.gated:
- y = self.g_proj(x).sigmoid() * y
- y = self.o_proj(y)
-
- return y, a.permute(0, 3, 1, 2)
-
-
-class EsmFoldDropout(nn.Module):
- """
- Implementation of dropout with the ability to share the dropout mask along a particular dimension.
- """
-
- def __init__(self, r: float, batch_dim: Union[int, List[int]]):
- super().__init__()
-
- self.r = r
- if isinstance(batch_dim, int):
- batch_dim = [batch_dim]
- self.batch_dim = batch_dim
- self.dropout = nn.Dropout(self.r)
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- shape = list(x.shape)
- if self.batch_dim is not None:
- for bd in self.batch_dim:
- shape[bd] = 1
- return x * self.dropout(x.new_ones(shape))
-
-
-class EsmFoldSequenceToPair(nn.Module):
- def __init__(self, sequence_state_dim, inner_dim, pairwise_state_dim):
- super().__init__()
-
- self.layernorm = nn.LayerNorm(sequence_state_dim)
- self.proj = nn.Linear(sequence_state_dim, inner_dim * 2, bias=True)
- self.o_proj = nn.Linear(2 * inner_dim, pairwise_state_dim, bias=True)
-
- torch.nn.init.zeros_(self.proj.bias)
- torch.nn.init.zeros_(self.o_proj.bias)
-
- def forward(self, sequence_state):
- """
- Inputs:
- sequence_state: B x L x sequence_state_dim
-
- Output:
- pairwise_state: B x L x L x pairwise_state_dim
-
- Intermediate state:
- B x L x L x 2*inner_dim
- """
-
- assert len(sequence_state.shape) == 3
-
- s = self.layernorm(sequence_state)
- s = self.proj(s)
- q, k = s.chunk(2, dim=-1)
-
- prod = q[:, None, :, :] * k[:, :, None, :]
- diff = q[:, None, :, :] - k[:, :, None, :]
-
- x = torch.cat([prod, diff], dim=-1)
- x = self.o_proj(x)
-
- return x
-
-
-class EsmFoldPairToSequence(nn.Module):
- def __init__(self, pairwise_state_dim, num_heads):
- super().__init__()
-
- self.layernorm = nn.LayerNorm(pairwise_state_dim)
- self.linear = nn.Linear(pairwise_state_dim, num_heads, bias=False)
-
- def forward(self, pairwise_state):
- """
- Inputs:
- pairwise_state: B x L x L x pairwise_state_dim
-
- Output:
- pairwise_bias: B x L x L x num_heads
- """
- assert len(pairwise_state.shape) == 4
- z = self.layernorm(pairwise_state)
- pairwise_bias = self.linear(z)
- return pairwise_bias
-
-
-class EsmFoldResidueMLP(nn.Module):
- def __init__(self, embed_dim, inner_dim, dropout=0):
- super().__init__()
-
- self.mlp = nn.Sequential(
- nn.LayerNorm(embed_dim),
- nn.Linear(embed_dim, inner_dim),
- nn.ReLU(),
- nn.Linear(inner_dim, embed_dim),
- nn.Dropout(dropout),
- )
-
- def forward(self, x):
- return x + self.mlp(x)
-
-
-class EsmFoldTriangularSelfAttentionBlock(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- sequence_state_dim = config.sequence_state_dim
- pairwise_state_dim = config.pairwise_state_dim
- sequence_num_heads = sequence_state_dim // config.sequence_head_width
- pairwise_num_heads = pairwise_state_dim // config.pairwise_head_width
-
- self.layernorm_1 = nn.LayerNorm(sequence_state_dim)
-
- self.sequence_to_pair = EsmFoldSequenceToPair(sequence_state_dim, pairwise_state_dim // 2, pairwise_state_dim)
- self.pair_to_sequence = EsmFoldPairToSequence(pairwise_state_dim, sequence_num_heads)
-
- self.seq_attention = EsmFoldSelfAttention(
- sequence_state_dim, sequence_num_heads, config.sequence_head_width, gated=True
- )
- self.tri_mul_out = EsmFoldTriangleMultiplicativeUpdate(config, _outgoing=True)
- self.tri_mul_in = EsmFoldTriangleMultiplicativeUpdate(config, _outgoing=False)
-
- self.tri_att_start = EsmFoldTriangleAttention(
- pairwise_state_dim, config.pairwise_head_width, pairwise_num_heads, inf=1e9, starting=True
- )
- self.tri_att_end = EsmFoldTriangleAttention(
- pairwise_state_dim, config.pairwise_head_width, pairwise_num_heads, inf=1e9, starting=False
- )
-
- self.mlp_seq = EsmFoldResidueMLP(sequence_state_dim, 4 * sequence_state_dim, dropout=config.dropout)
- self.mlp_pair = EsmFoldResidueMLP(pairwise_state_dim, 4 * pairwise_state_dim, dropout=config.dropout)
-
- self.drop = nn.Dropout(config.dropout)
- self.row_drop = EsmFoldDropout(config.dropout * 2, 2)
- self.col_drop = EsmFoldDropout(config.dropout * 2, 1)
-
- def forward(self, sequence_state, pairwise_state, mask=None, chunk_size=None, **__kwargs):
- """
- Inputs:
- sequence_state: B x L x sequence_state_dim pairwise_state: B x L x L x pairwise_state_dim mask: B x L boolean
- tensor of valid positions
-
- Output:
- sequence_state: B x L x sequence_state_dim pairwise_state: B x L x L x pairwise_state_dim
- """
- if len(sequence_state.shape) != 3:
- raise ValueError(f"`sequence_state` should be a 3d-tensor, got {len(sequence_state.shape)} dims.")
- if len(pairwise_state.shape) != 4:
- raise ValueError(f"`pairwise_state` should be a 4d-tensor, got {len(pairwise_state.shape)} dims.")
- if mask is not None and len(mask.shape) != 2:
- raise ValueError(f"`mask` should be a 2d-tensor, got {len(mask.shape)} dims.")
-
- batch_dim, seq_dim, sequence_state_dim = sequence_state.shape
- pairwise_state_dim = pairwise_state.shape[3]
-
- if sequence_state_dim != self.config.sequence_state_dim:
- raise ValueError(
- "`sequence_state` last dimension should be equal to `self.sequence_state_dim`. Got "
- f"{sequence_state_dim} != {self.config.sequence_state_dim}."
- )
- if pairwise_state_dim != self.config.pairwise_state_dim:
- raise ValueError(
- "`pairwise_state` last dimension should be equal to `self.pairwise_state_dim`. Got "
- f"{pairwise_state_dim} != {self.config.pairwise_state_dim}."
- )
- if batch_dim != pairwise_state.shape[0]:
- raise ValueError(
- f"`sequence_state` and `pairwise_state` have inconsistent batch size: {batch_dim} != "
- f"{pairwise_state.shape[0]}."
- )
- if seq_dim != pairwise_state.shape[1] or seq_dim != pairwise_state.shape[2]:
- raise ValueError(
- f"`sequence_state` and `pairwise_state` have inconsistent sequence length: {seq_dim} != "
- f"{pairwise_state.shape[1]} or {pairwise_state.shape[2]}."
- )
-
- # Update sequence state
- bias = self.pair_to_sequence(pairwise_state)
-
- # Self attention with bias + mlp.
- y = self.layernorm_1(sequence_state)
- y, _ = self.seq_attention(y, mask=mask, bias=bias)
- sequence_state = sequence_state + self.drop(y)
- sequence_state = self.mlp_seq(sequence_state)
-
- # Update pairwise state
- pairwise_state = pairwise_state + self.sequence_to_pair(sequence_state)
-
- # Axial attention with triangular bias.
- tri_mask = mask.unsqueeze(2) * mask.unsqueeze(1) if mask is not None else None
- pairwise_state = pairwise_state + self.row_drop(self.tri_mul_out(pairwise_state, mask=tri_mask))
- pairwise_state = pairwise_state + self.col_drop(self.tri_mul_in(pairwise_state, mask=tri_mask))
- pairwise_state = pairwise_state + self.row_drop(
- self.tri_att_start(pairwise_state, mask=tri_mask, chunk_size=chunk_size)
- )
- pairwise_state = pairwise_state + self.col_drop(
- self.tri_att_end(pairwise_state, mask=tri_mask, chunk_size=chunk_size)
- )
-
- # MLP over pairs.
- pairwise_state = self.mlp_pair(pairwise_state)
-
- return sequence_state, pairwise_state
-
-
-class EsmCategoricalMixture:
- def __init__(self, param, bins=50, start=0, end=1):
- # All tensors are of shape ..., bins.
- self.logits = param
- bins = torch.linspace(start, end, bins + 1, device=self.logits.device, dtype=self.logits.dtype)
- self.v_bins = (bins[:-1] + bins[1:]) / 2
-
- def log_prob(self, true):
- # Shapes are:
- # self.probs: ... x bins
- # true : ...
- true_index = (true.unsqueeze(-1) - self.v_bins[[None] * true.ndim]).abs().argmin(-1)
- nll = self.logits.log_softmax(-1)
- return torch.take_along_dim(nll, true_index.unsqueeze(-1), dim=-1).squeeze(-1)
-
- def mean(self):
- return (self.logits.softmax(-1) @ self.v_bins.unsqueeze(1)).squeeze(-1)
-
-
-def categorical_lddt(logits, bins=50):
- # Logits are ..., 37, bins.
- return EsmCategoricalMixture(logits, bins=bins).mean()
-
-
-def get_axial_mask(mask):
- """
- Helper to convert B x L mask of valid positions to axial mask used in row column attentions.
-
- Input:
- mask: B x L tensor of booleans
-
- Output:
- mask: B x L x L tensor of booleans
- """
-
- if mask is None:
- return None
-
- if len(mask.shape) != 2:
- raise ValueError(f"`mask` should be a 2d-tensor, got {len(mask.shape)} dims.")
- batch_dim, seq_dim = mask.shape
- m = mask.unsqueeze(1).expand(batch_dim, seq_dim, seq_dim)
- m = m.reshape(batch_dim * seq_dim, seq_dim)
- return m
-
-
-class EsmFoldRelativePosition(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.bins = config.position_bins
-
- # Note an additional offset is used so that the 0th position
- # is reserved for masked pairs.
- self.embedding = torch.nn.Embedding(2 * self.bins + 2, config.pairwise_state_dim)
-
- def forward(self, residue_index, mask=None):
- """
- Input:
- residue_index: B x L tensor of indices (dytpe=torch.long) mask: B x L tensor of booleans
-
- Output:
- pairwise_state: B x L x L x pairwise_state_dim tensor of embeddings
- """
- if residue_index.dtype != torch.long:
- raise ValueError(f"`residue_index` has dtype {residue_index.dtype}, it should be `torch.long`.")
- if mask is not None and residue_index.shape != mask.shape:
- raise ValueError(
- f"`residue_index` and `mask` have inconsistent shapes: {residue_index.shape} != {mask.shape}."
- )
-
- diff = residue_index[:, None, :] - residue_index[:, :, None]
- diff = diff.clamp(-self.bins, self.bins)
- diff = diff + self.bins + 1 # Add 1 to adjust for padding index.
-
- if mask is not None:
- mask = mask[:, None, :] * mask[:, :, None]
- diff[mask == False] = 0 # noqa: E712
-
- output = self.embedding(diff)
- return output
-
-
-class EsmFoldAngleResnetBlock(nn.Module):
- def __init__(self, config):
- super().__init__()
-
- self.linear_1 = EsmFoldLinear(config.resnet_dim, config.resnet_dim, init="relu")
- self.linear_2 = EsmFoldLinear(config.resnet_dim, config.resnet_dim, init="final")
-
- self.relu = nn.ReLU()
-
- def forward(self, a: torch.Tensor) -> torch.Tensor:
- s_initial = a
-
- a = self.relu(a)
- a = self.linear_1(a)
- a = self.relu(a)
- a = self.linear_2(a)
-
- return a + s_initial
-
-
-class EsmFoldAngleResnet(nn.Module):
- """
- Implements Algorithm 20, lines 11-14
- """
-
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- self.linear_in = EsmFoldLinear(config.sequence_dim, config.resnet_dim)
- self.linear_initial = EsmFoldLinear(config.sequence_dim, config.resnet_dim)
-
- self.layers = nn.ModuleList()
- for _ in range(config.num_resnet_blocks):
- layer = EsmFoldAngleResnetBlock(config)
- self.layers.append(layer)
-
- self.linear_out = EsmFoldLinear(config.resnet_dim, config.num_angles * 2)
-
- self.relu = nn.ReLU()
-
- def forward(self, s: torch.Tensor, s_initial: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
- """
- Args:
- s:
- [*, C_hidden] single embedding
- s_initial:
- [*, C_hidden] single embedding as of the start of the StructureModule
- Returns:
- [*, no_angles, 2] predicted angles
- """
- # NOTE: The ReLU's applied to the inputs are absent from the supplement
- # pseudocode but present in the source. For maximal compatibility with
- # the pretrained weights, I'm going with the source.
-
- # [*, C_hidden]
- s_initial = self.relu(s_initial)
- s_initial = self.linear_initial(s_initial)
- s = self.relu(s)
- s = self.linear_in(s)
- s = s + s_initial
-
- for l in self.layers:
- s = l(s)
-
- s = self.relu(s)
-
- # [*, no_angles * 2]
- s = self.linear_out(s)
-
- # [*, no_angles, 2]
- s = s.view(s.shape[:-1] + (-1, 2))
-
- unnormalized_s = s
- norm_denom = torch.sqrt(
- torch.clamp(
- torch.sum(s**2, dim=-1, keepdim=True),
- min=self.config.epsilon,
- )
- )
- s = s / norm_denom
-
- return unnormalized_s, s
-
-
-class EsmFoldInvariantPointAttention(nn.Module):
- """
- Implements Algorithm 22.
- """
-
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- c_s = config.sequence_dim
- c_z = config.pairwise_dim
- self.hidden_dim = config.ipa_dim
- self.num_heads = config.num_heads_ipa
- self.num_qk_points = config.num_qk_points
- self.num_v_points = config.num_v_points
-
- # These linear layers differ from their specifications in the
- # supplement. There, they lack bias and use Glorot initialization.
- # Here as in the official source, they have bias and use the default
- # Lecun initialization.
- hc = config.ipa_dim * config.num_heads_ipa
- self.linear_q = EsmFoldLinear(c_s, hc)
- self.linear_kv = EsmFoldLinear(c_s, 2 * hc)
-
- hpq = config.num_heads_ipa * config.num_qk_points * 3
- self.linear_q_points = EsmFoldLinear(c_s, hpq)
-
- hpkv = config.num_heads_ipa * (config.num_qk_points + config.num_v_points) * 3
- self.linear_kv_points = EsmFoldLinear(c_s, hpkv)
-
- self.linear_b = EsmFoldLinear(c_z, config.num_heads_ipa)
-
- self.head_weights = nn.Parameter(torch.zeros((config.num_heads_ipa)))
-
- concat_out_dim = config.num_heads_ipa * (c_z + config.ipa_dim + config.num_v_points * 4)
- self.linear_out = EsmFoldLinear(concat_out_dim, c_s, init="final")
-
- self.softmax = nn.Softmax(dim=-1)
- self.softplus = nn.Softplus()
-
- def forward(
- self,
- s: torch.Tensor,
- z: Optional[torch.Tensor],
- r: Rigid,
- mask: torch.Tensor,
- _offload_inference: bool = False,
- _z_reference_list: Optional[Sequence[torch.Tensor]] = None,
- ) -> torch.Tensor:
- """
- Args:
- s:
- [*, N_res, C_s] single representation
- z:
- [*, N_res, N_res, C_z] pair representation
- r:
- [*, N_res] transformation object
- mask:
- [*, N_res] mask
- Returns:
- [*, N_res, C_s] single representation update
- """
- z = [z]
-
- #######################################
- # Generate scalar and point activations
- #######################################
- # [*, N_res, H * C_hidden]
- q = self.linear_q(s)
- kv = self.linear_kv(s)
-
- # [*, N_res, H, C_hidden]
- q = q.view(q.shape[:-1] + (self.num_heads, -1))
-
- # [*, N_res, H, 2 * C_hidden]
- kv = kv.view(kv.shape[:-1] + (self.num_heads, -1))
-
- # [*, N_res, H, C_hidden]
- k, v = torch.split(kv, self.hidden_dim, dim=-1)
-
- # [*, N_res, H * P_q * 3]
- q_pts = self.linear_q_points(s)
-
- # This is kind of clunky, but it's how the original does it
- # [*, N_res, H * P_q, 3]
- q_pts = torch.split(q_pts, q_pts.shape[-1] // 3, dim=-1)
- q_pts = torch.stack(q_pts, dim=-1)
- q_pts = r[..., None].apply(q_pts)
-
- # [*, N_res, H, P_q, 3]
- q_pts = q_pts.view(q_pts.shape[:-2] + (self.num_heads, self.num_qk_points, 3))
-
- # [*, N_res, H * (P_q + P_v) * 3]
- kv_pts = self.linear_kv_points(s)
-
- # [*, N_res, H * (P_q + P_v), 3]
- kv_pts = torch.split(kv_pts, kv_pts.shape[-1] // 3, dim=-1)
- kv_pts = torch.stack(kv_pts, dim=-1)
- kv_pts = r[..., None].apply(kv_pts)
-
- # [*, N_res, H, (P_q + P_v), 3]
- kv_pts = kv_pts.view(kv_pts.shape[:-2] + (self.num_heads, -1, 3))
-
- # [*, N_res, H, P_q/P_v, 3]
- k_pts, v_pts = torch.split(kv_pts, [self.num_qk_points, self.num_v_points], dim=-2)
-
- ##########################
- # Compute attention scores
- ##########################
- # [*, N_res, N_res, H]
- b = self.linear_b(z[0])
-
- if _offload_inference:
- assert sys.getrefcount(z[0]) == 2
- z[0] = z[0].cpu()
-
- # [*, H, N_res, N_res]
- if is_fp16_enabled():
- with torch.cuda.amp.autocast(enabled=False):
- a = torch.matmul(
- permute_final_dims(q.float(), (1, 0, 2)), # [*, H, N_res, C_hidden]
- permute_final_dims(k.float(), (1, 2, 0)), # [*, H, C_hidden, N_res]
- )
- else:
- a = torch.matmul(
- permute_final_dims(q, (1, 0, 2)), # [*, H, N_res, C_hidden]
- permute_final_dims(k, (1, 2, 0)), # [*, H, C_hidden, N_res]
- )
-
- a *= math.sqrt(1.0 / (3 * self.hidden_dim))
- a += math.sqrt(1.0 / 3) * permute_final_dims(b, (2, 0, 1))
-
- # [*, N_res, N_res, H, P_q, 3]
- pt_att = q_pts.unsqueeze(-4) - k_pts.unsqueeze(-5)
- pt_att = pt_att**2
-
- # [*, N_res, N_res, H, P_q]
- pt_att = sum(torch.unbind(pt_att, dim=-1))
- head_weights = self.softplus(self.head_weights).view(*((1,) * len(pt_att.shape[:-2]) + (-1, 1)))
- head_weights = head_weights * math.sqrt(1.0 / (3 * (self.num_qk_points * 9.0 / 2)))
- pt_att = pt_att * head_weights
-
- # [*, N_res, N_res, H]
- pt_att = torch.sum(pt_att, dim=-1) * (-0.5)
- # [*, N_res, N_res]
- square_mask = mask.unsqueeze(-1) * mask.unsqueeze(-2)
- square_mask = self.config.inf * (square_mask - 1)
-
- # [*, H, N_res, N_res]
- pt_att = permute_final_dims(pt_att, (2, 0, 1))
-
- a = a + pt_att
- a = a + square_mask.unsqueeze(-3)
- a = self.softmax(a)
-
- ################
- # Compute output
- ################
- # [*, N_res, H, C_hidden]
- o = torch.matmul(a, v.transpose(-2, -3).to(dtype=a.dtype)).transpose(-2, -3)
-
- # [*, N_res, H * C_hidden]
- o = flatten_final_dims(o, 2)
-
- # [*, H, 3, N_res, P_v]
- o_pt = torch.sum(
- (a[..., None, :, :, None] * permute_final_dims(v_pts, (1, 3, 0, 2))[..., None, :, :]),
- dim=-2,
- )
-
- # [*, N_res, H, P_v, 3]
- o_pt = permute_final_dims(o_pt, (2, 0, 3, 1))
- o_pt = r[..., None, None].invert_apply(o_pt)
-
- # [*, N_res, H * P_v]
- o_pt_norm = flatten_final_dims(torch.sqrt(torch.sum(o_pt**2, dim=-1) + self.config.epsilon), 2)
-
- # [*, N_res, H * P_v, 3]
- o_pt = o_pt.reshape(*o_pt.shape[:-3], -1, 3)
-
- if _offload_inference:
- z[0] = z[0].to(o_pt.device)
-
- # [*, N_res, H, C_z]
- o_pair = torch.matmul(a.transpose(-2, -3), z[0].to(dtype=a.dtype))
-
- # [*, N_res, H * C_z]
- o_pair = flatten_final_dims(o_pair, 2)
-
- # [*, N_res, C_s]
- s = self.linear_out(
- torch.cat((o, *torch.unbind(o_pt, dim=-1), o_pt_norm, o_pair), dim=-1).to(dtype=z[0].dtype)
- )
-
- return s
-
-
-class EsmFoldBackboneUpdate(nn.Module):
- """
- Implements part of Algorithm 23.
- """
-
- def __init__(self, config):
- super().__init__()
-
- self.linear = EsmFoldLinear(config.sequence_dim, 6, init="final")
-
- def forward(self, s: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
- """
- Args:
- [*, N_res, C_s] single representation
- Returns:
- [*, N_res, 6] update vector
- """
- # [*, 6]
- update = self.linear(s)
-
- return update
-
-
-class EsmFoldStructureModuleTransitionLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
-
- self.linear_1 = EsmFoldLinear(config.sequence_dim, config.sequence_dim, init="relu")
- self.linear_2 = EsmFoldLinear(config.sequence_dim, config.sequence_dim, init="relu")
- self.linear_3 = EsmFoldLinear(config.sequence_dim, config.sequence_dim, init="final")
-
- self.relu = nn.ReLU()
-
- def forward(self, s):
- s_initial = s
- s = self.linear_1(s)
- s = self.relu(s)
- s = self.linear_2(s)
- s = self.relu(s)
- s = self.linear_3(s)
-
- s = s + s_initial
-
- return s
-
-
-class EsmFoldStructureModuleTransition(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- self.layers = nn.ModuleList()
- for _ in range(config.num_transition_layers):
- l = EsmFoldStructureModuleTransitionLayer(config)
- self.layers.append(l)
-
- self.dropout = nn.Dropout(config.dropout_rate)
- self.layer_norm = LayerNorm(config.sequence_dim)
-
- def forward(self, s):
- for l in self.layers:
- s = l(s)
-
- s = self.dropout(s)
- s = self.layer_norm(s)
-
- return s
-
-
-class EsmFoldStructureModule(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- # Buffers to be lazily initialized later
- # self.default_frames
- # self.group_idx
- # self.atom_mask
- # self.lit_positions
-
- self.layer_norm_s = LayerNorm(config.sequence_dim)
- self.layer_norm_z = LayerNorm(config.pairwise_dim)
-
- self.linear_in = EsmFoldLinear(config.sequence_dim, config.sequence_dim)
-
- self.ipa = EsmFoldInvariantPointAttention(config)
-
- self.ipa_dropout = nn.Dropout(config.dropout_rate)
- self.layer_norm_ipa = LayerNorm(config.sequence_dim)
-
- self.transition = EsmFoldStructureModuleTransition(config)
- self.bb_update = EsmFoldBackboneUpdate(config)
- self.angle_resnet = EsmFoldAngleResnet(config)
-
- def forward(
- self,
- evoformer_output_dict,
- aatype,
- mask=None,
- _offload_inference=False,
- ):
- """
- Args:
- evoformer_output_dict:
- Dictionary containing:
- "single":
- [*, N_res, C_s] single representation
- "pair":
- [*, N_res, N_res, C_z] pair representation
- aatype:
- [*, N_res] amino acid indices
- mask:
- Optional [*, N_res] sequence mask
- Returns:
- A dictionary of outputs
- """
- s = evoformer_output_dict["single"]
-
- if mask is None:
- # [*, N]
- mask = s.new_ones(s.shape[:-1])
-
- # [*, N, C_s]
- s = self.layer_norm_s(s)
-
- # [*, N, N, C_z]
- z = self.layer_norm_z(evoformer_output_dict["pair"])
-
- z_reference_list = None
- if _offload_inference:
- assert sys.getrefcount(evoformer_output_dict["pair"]) == 2
- evoformer_output_dict["pair"] = evoformer_output_dict["pair"].cpu()
- z_reference_list = [z]
- z = None
-
- # [*, N, C_s]
- s_initial = s
- s = self.linear_in(s)
-
- # [*, N]
- rigids = Rigid.identity(
- s.shape[:-1],
- s.dtype,
- s.device,
- self.training,
- fmt="quat",
- )
- outputs = []
- for i in range(self.config.num_blocks):
- # [*, N, C_s]
- s = s + self.ipa(
- s,
- z,
- rigids,
- mask,
- _offload_inference=_offload_inference,
- _z_reference_list=z_reference_list,
- )
- s = self.ipa_dropout(s)
- s = self.layer_norm_ipa(s)
- s = self.transition(s)
-
- # [*, N]
- rigids = rigids.compose_q_update_vec(self.bb_update(s))
-
- # To hew as closely as possible to AlphaFold, we convert our
- # quaternion-based transformations to rotation-matrix ones
- # here
- backb_to_global = Rigid(
- Rotation(rot_mats=rigids.get_rots().get_rot_mats(), quats=None),
- rigids.get_trans(),
- )
-
- backb_to_global = backb_to_global.scale_translation(self.config.trans_scale_factor)
-
- # [*, N, 7, 2]
- unnormalized_angles, angles = self.angle_resnet(s, s_initial)
-
- all_frames_to_global = self.torsion_angles_to_frames(backb_to_global, angles, aatype)
-
- pred_xyz = self.frames_and_literature_positions_to_atom14_pos(all_frames_to_global, aatype)
-
- scaled_rigids = rigids.scale_translation(self.config.trans_scale_factor)
-
- preds = {
- "frames": scaled_rigids.to_tensor_7(),
- "sidechain_frames": all_frames_to_global.to_tensor_4x4(),
- "unnormalized_angles": unnormalized_angles,
- "angles": angles,
- "positions": pred_xyz,
- "states": s,
- }
-
- outputs.append(preds)
-
- rigids = rigids.stop_rot_gradient()
-
- del z, z_reference_list
-
- if _offload_inference:
- evoformer_output_dict["pair"] = evoformer_output_dict["pair"].to(s.device)
-
- outputs = dict_multimap(torch.stack, outputs)
- outputs["single"] = s
-
- return outputs
-
- def _init_residue_constants(self, float_dtype, device):
- if not hasattr(self, "default_frames"):
- self.register_buffer(
- "default_frames",
- torch.tensor(
- residue_constants.restype_rigid_group_default_frame,
- dtype=float_dtype,
- device=device,
- requires_grad=False,
- ),
- persistent=False,
- )
- if not hasattr(self, "group_idx"):
- self.register_buffer(
- "group_idx",
- torch.tensor(
- residue_constants.restype_atom14_to_rigid_group,
- device=device,
- requires_grad=False,
- ),
- persistent=False,
- )
- if not hasattr(self, "atom_mask"):
- self.register_buffer(
- "atom_mask",
- torch.tensor(
- residue_constants.restype_atom14_mask,
- dtype=float_dtype,
- device=device,
- requires_grad=False,
- ),
- persistent=False,
- )
- if not hasattr(self, "lit_positions"):
- self.register_buffer(
- "lit_positions",
- torch.tensor(
- residue_constants.restype_atom14_rigid_group_positions,
- dtype=float_dtype,
- device=device,
- requires_grad=False,
- ),
- persistent=False,
- )
-
- def torsion_angles_to_frames(self, r, alpha, f):
- # Lazily initialize the residue constants on the correct device
- self._init_residue_constants(alpha.dtype, alpha.device)
- # Separated purely to make testing less annoying
- return torsion_angles_to_frames(r, alpha, f, self.default_frames)
-
- def frames_and_literature_positions_to_atom14_pos(self, r, f): # [*, N, 8] # [*, N]
- # Lazily initialize the residue constants on the correct device
- self._init_residue_constants(r.get_rots().dtype, r.get_rots().device)
- return frames_and_literature_positions_to_atom14_pos(
- r,
- f,
- self.default_frames,
- self.group_idx,
- self.atom_mask,
- self.lit_positions,
- )
-
-
-class EsmFoldingTrunk(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- c_s = config.sequence_state_dim
- c_z = config.pairwise_state_dim
-
- self.pairwise_positional_embedding = EsmFoldRelativePosition(config)
-
- self.blocks = nn.ModuleList([EsmFoldTriangularSelfAttentionBlock(config) for _ in range(config.num_blocks)])
-
- self.recycle_bins = 15
- self.recycle_s_norm = nn.LayerNorm(c_s)
- self.recycle_z_norm = nn.LayerNorm(c_z)
- self.recycle_disto = nn.Embedding(self.recycle_bins, c_z)
- self.recycle_disto.weight[0].detach().zero_()
-
- self.structure_module = EsmFoldStructureModule(config.structure_module)
- self.trunk2sm_s = nn.Linear(c_s, config.structure_module.sequence_dim)
- self.trunk2sm_z = nn.Linear(c_z, config.structure_module.pairwise_dim)
-
- self.chunk_size = config.chunk_size
-
- def set_chunk_size(self, chunk_size):
- # This parameter means the axial attention will be computed
- # in a chunked manner. This should make the memory used more or less O(L) instead of O(L^2).
- # It's equivalent to running a for loop over chunks of the dimension we're iterative over,
- # where the chunk_size is the size of the chunks, so 128 would mean to parse 128-length chunks.
- self.chunk_size = chunk_size
-
- def forward(self, seq_feats, pair_feats, true_aa, residx, mask, no_recycles):
- """
- Inputs:
- seq_feats: B x L x C tensor of sequence features pair_feats: B x L x L x C tensor of pair features residx: B
- x L long tensor giving the position in the sequence mask: B x L boolean tensor indicating valid residues
-
- Output:
- predicted_structure: B x L x (num_atoms_per_residue * 3) tensor wrapped in a Coordinates object
- """
-
- device = seq_feats.device
- s_s_0 = seq_feats
- s_z_0 = pair_feats
-
- if no_recycles is None:
- no_recycles = self.config.max_recycles
- else:
- if no_recycles < 0:
- raise ValueError("Number of recycles must not be negative.")
- no_recycles += 1 # First 'recycle' is just the standard forward pass through the model.
-
- def trunk_iter(s, z, residx, mask):
- z = z + self.pairwise_positional_embedding(residx, mask=mask)
-
- for block in self.blocks:
- s, z = block(s, z, mask=mask, residue_index=residx, chunk_size=self.chunk_size)
- return s, z
-
- s_s = s_s_0
- s_z = s_z_0
- recycle_s = torch.zeros_like(s_s)
- recycle_z = torch.zeros_like(s_z)
- recycle_bins = torch.zeros(*s_z.shape[:-1], device=device, dtype=torch.int64)
-
- for recycle_idx in range(no_recycles):
- with ContextManagers([] if recycle_idx == no_recycles - 1 else [torch.no_grad()]):
- # === Recycling ===
- recycle_s = self.recycle_s_norm(recycle_s.detach()).to(device)
- recycle_z = self.recycle_z_norm(recycle_z.detach()).to(device)
- recycle_z += self.recycle_disto(recycle_bins.detach()).to(device)
-
- s_s, s_z = trunk_iter(s_s_0 + recycle_s, s_z_0 + recycle_z, residx, mask)
-
- # === Structure module ===
- structure = self.structure_module(
- {"single": self.trunk2sm_s(s_s), "pair": self.trunk2sm_z(s_z)},
- true_aa,
- mask.float(),
- )
-
- recycle_s = s_s
- recycle_z = s_z
- # Distogram needs the N, CA, C coordinates, and bin constants same as alphafold.
- recycle_bins = EsmFoldingTrunk.distogram(
- structure["positions"][-1][:, :, :3],
- 3.375,
- 21.375,
- self.recycle_bins,
- )
-
- structure["s_s"] = s_s
- structure["s_z"] = s_z
-
- return structure
-
- @staticmethod
- def distogram(coords, min_bin, max_bin, num_bins):
- # Coords are [... L x 3 x 3], where it's [N, CA, C] x 3 coordinates.
- boundaries = torch.linspace(
- min_bin,
- max_bin,
- num_bins - 1,
- device=coords.device,
- )
- boundaries = boundaries**2
- N, CA, C = [x.squeeze(-2) for x in coords.chunk(3, dim=-2)]
- # Infer CB coordinates.
- b = CA - N
- c = C - CA
- a = b.cross(c, dim=-1)
- CB = -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + CA
- dists = (CB[..., None, :, :] - CB[..., :, None, :]).pow(2).sum(dim=-1, keepdims=True)
- bins = torch.sum(dists > boundaries, dim=-1) # [..., L, L]
- return bins
-
-
-# TODO Add information to the docstring about any methods that convert to PDB format, or otherwise prepare
-# the outputs for downstream use.
-
-
-@add_start_docstrings(
- """
- ESMForProteinFolding is the HuggingFace port of the original ESMFold model. It consists of an ESM-2 "stem" followed
- by a protein folding "head", although unlike most other output heads, this "head" is similar in size and runtime to
- the rest of the model combined! It outputs a dictionary containing predicted structural information about the input
- protein(s).
- """,
- ESM_START_DOCSTRING,
-)
-class EsmForProteinFolding(EsmPreTrainedModel):
- _no_split_modules = ["EsmFoldStructureModule", "EsmFoldTriangularSelfAttentionBlock"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.config = config
-
- self.distogram_bins = 64
-
- self.esm = EsmModel(config, add_pooling_layer=False)
-
- self.esm.requires_grad_(False)
- if self.config.esmfold_config.fp16_esm:
- self.esm.half()
-
- self.esm_feats = self.config.hidden_size
- self.esm_attns = self.config.num_hidden_layers * self.config.num_attention_heads
- self.esm_layers = self.config.num_hidden_layers
- self.register_buffer("af2_to_esm", self._af2_to_esm_from_vocab_list(config.vocab_list))
- self.esm_s_combine = nn.Parameter(torch.zeros(self.esm_layers + 1))
-
- trunk_config = self.config.esmfold_config.trunk
- c_s = trunk_config.sequence_state_dim
- c_z = trunk_config.pairwise_state_dim
- self.esm_s_mlp = nn.Sequential(
- LayerNorm(self.esm_feats),
- nn.Linear(self.esm_feats, c_s),
- nn.ReLU(),
- nn.Linear(c_s, c_s),
- )
-
- # 0 is padding, N is unknown residues, N + 1 is mask.
- self.n_tokens_embed = residue_constants.restype_num + 3
- self.pad_idx = 0
- self.unk_idx = self.n_tokens_embed - 2
- self.mask_idx = self.n_tokens_embed - 1
- self.esm_dict_cls_idx = self.config.vocab_list.index("")
- self.esm_dict_mask_idx = self.config.vocab_list.index("")
- self.esm_dict_eos_idx = self.config.vocab_list.index("")
- self.esm_dict_padding_idx = self.config.vocab_list.index("")
- if self.config.esmfold_config.embed_aa:
- self.embedding = nn.Embedding(self.n_tokens_embed, c_s, padding_idx=0)
-
- self.trunk = EsmFoldingTrunk(trunk_config)
-
- self.distogram_head = nn.Linear(c_z, self.distogram_bins)
- self.ptm_head = nn.Linear(c_z, self.distogram_bins)
- self.lm_head = nn.Linear(c_s, self.n_tokens_embed)
- self.lddt_bins = 50
- structure_module_config = trunk_config.structure_module
- self.lddt_head = nn.Sequential(
- nn.LayerNorm(structure_module_config.sequence_dim),
- nn.Linear(structure_module_config.sequence_dim, self.config.esmfold_config.lddt_head_hid_dim),
- nn.Linear(self.config.esmfold_config.lddt_head_hid_dim, self.config.esmfold_config.lddt_head_hid_dim),
- nn.Linear(self.config.esmfold_config.lddt_head_hid_dim, 37 * self.lddt_bins),
- )
-
- @staticmethod
- def _af2_to_esm_from_vocab_list(vocab_list: List[str]) -> torch.Tensor:
- # Remember that t is shifted from residue_constants by 1 (0 is padding).
- esm_reorder = [vocab_list.index("")] + [vocab_list.index(v) for v in residue_constants.restypes_with_x]
- return torch.tensor(esm_reorder)
-
- @add_start_docstrings_to_model_forward(ESMFOLD_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=EsmForProteinFoldingOutput, config_class=EsmConfig)
- def forward(
- self,
- input_ids: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- masking_pattern: Optional[torch.Tensor] = None,
- num_recycles: Optional[int] = None,
- ) -> EsmForProteinFoldingOutput:
- r"""
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, EsmForProteinFolding
-
- >>> model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")
- >>> tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
- >>> inputs = tokenizer(["MLKNVQVQLV"], return_tensors="pt", add_special_tokens=False) # A tiny random peptide
- >>> outputs = model(**inputs)
- >>> folded_positions = outputs.positions
- ```
-
- """
- cfg = self.config.esmfold_config
-
- aa = input_ids # B x L
- B = aa.shape[0]
- L = aa.shape[1]
- device = input_ids.device
- if attention_mask is None:
- attention_mask = torch.ones_like(aa, device=device)
- if position_ids is None:
- position_ids = torch.arange(L, device=device).expand_as(input_ids)
-
- # === ESM ===
- esmaa = self.af2_idx_to_esm_idx(aa, attention_mask)
-
- if masking_pattern is not None:
- masked_aa, esmaa, mlm_targets = self.bert_mask(aa, esmaa, attention_mask, masking_pattern)
- else:
- masked_aa = aa
- mlm_targets = None
-
- # We get sequence and pair representations from whatever version of ESM /
- # configuration we are using. The sequence representation esm_s is always
- # present. The pair embedding esm_z may be present depending on the
- # configuration of the model. If esm_z is not used by the model then it
- # is returned as None here.
- esm_s = self.compute_language_model_representations(esmaa)
-
- # Convert esm_s and esm_z, if present, to the precision used by the trunk and
- # the structure module. These tensors may be a lower precision if, for example,
- # we're running the language model in fp16 precision.
- esm_s = esm_s.to(self.esm_s_combine.dtype)
-
- if cfg.esm_ablate_sequence:
- esm_s = esm_s * 0
-
- esm_s = esm_s.detach()
-
- # === preprocessing ===
- esm_s = (self.esm_s_combine.softmax(0).unsqueeze(0) @ esm_s).squeeze(2)
- s_s_0 = self.esm_s_mlp(esm_s)
-
- s_z_0 = s_s_0.new_zeros(B, L, L, cfg.trunk.pairwise_state_dim)
-
- if self.config.esmfold_config.embed_aa:
- s_s_0 += self.embedding(masked_aa)
-
- structure: dict = self.trunk(s_s_0, s_z_0, aa, position_ids, attention_mask, no_recycles=num_recycles)
- # Documenting what we expect:
- structure = {
- k: v
- for k, v in structure.items()
- if k
- in [
- "s_z",
- "s_s",
- "frames",
- "sidechain_frames",
- "unnormalized_angles",
- "angles",
- "positions",
- "states",
- ]
- }
-
- # Add BERT mask for the loss to use, if available.
- if mlm_targets:
- structure["mlm_targets"] = mlm_targets
-
- disto_logits = self.distogram_head(structure["s_z"])
- disto_logits = (disto_logits + disto_logits.transpose(1, 2)) / 2
- structure["distogram_logits"] = disto_logits
-
- lm_logits = self.lm_head(structure["s_s"])
- structure["lm_logits"] = lm_logits
-
- structure["aatype"] = aa
- make_atom14_masks(structure)
- # Of course, this doesn't respect the true mask because it doesn't know about it...
- # We're not going to properly mask change of index tensors:
- # "residx_atom14_to_atom37",
- # "residx_atom37_to_atom14",
- for k in [
- "atom14_atom_exists",
- "atom37_atom_exists",
- ]:
- structure[k] *= attention_mask.unsqueeze(-1)
- structure["residue_index"] = position_ids
-
- lddt_head = self.lddt_head(structure["states"]).reshape(structure["states"].shape[0], B, L, -1, self.lddt_bins)
- structure["lddt_head"] = lddt_head
- plddt = categorical_lddt(lddt_head[-1], bins=self.lddt_bins)
- structure["plddt"] = plddt
-
- ptm_logits = self.ptm_head(structure["s_z"])
- structure["ptm_logits"] = ptm_logits
- structure["ptm"] = compute_tm(ptm_logits, max_bin=31, no_bins=self.distogram_bins)
- structure.update(compute_predicted_aligned_error(ptm_logits, max_bin=31, no_bins=self.distogram_bins))
-
- return EsmForProteinFoldingOutput(**structure)
-
- def af2_idx_to_esm_idx(self, aa, mask):
- # avoid indexing on different devices
- if self.af2_to_esm.device != aa.device:
- self.af2_to_esm = self.af2_to_esm.to(aa.device)
- aa = (aa + 1).masked_fill(mask != 1, 0)
- return self.af2_to_esm[aa]
-
- def compute_language_model_representations(self, esmaa: torch.Tensor) -> torch.Tensor:
- device = next(self.parameters()).device
- B, L = esmaa.shape # B = batch size, L = sequence length.
-
- if self.config.esmfold_config.bypass_lm:
- esm_s = torch.zeros(B, L, self.esm_s_combine.size[0], -1, self.esm_feats, device=device)
- return esm_s
-
- bosi, eosi = self.esm_dict_cls_idx, self.esm_dict_eos_idx
- bos = esmaa.new_full((B, 1), bosi)
- eos = esmaa.new_full((B, 1), self.esm_dict_padding_idx)
- esmaa = torch.cat([bos, esmaa, eos], dim=1)
- # Use the first padding index as eos during inference.
- esmaa[range(B), (esmaa != 1).sum(1)] = eosi
-
- # _, esm_z, esm_s = self.esm(esmaa, return_pairs=self.config.esmfold_config.use_esm_attn_map)
- # Because we do not support use_esm_attn_map in the HF port as it is not used in any public models,
- # esm_z is always None
- esm_hidden_states = self.esm(esmaa, attention_mask=esmaa != 1, output_hidden_states=True)["hidden_states"]
- esm_s = torch.stack(esm_hidden_states, dim=2)
-
- esm_s = esm_s[:, 1:-1] # B, L, nLayers, C
-
- return esm_s
-
- def bert_mask(self, aa, esmaa, mask, pattern):
- new_aa = aa.clone()
- target = aa.clone()
- new_esmaa = esmaa.clone()
- new_aa[pattern == 1] = self.mask_idx
- target[pattern != 1] = 0
- new_esmaa[pattern == 1] = self.esm_dict_mask_idx
- return new_aa, new_esmaa, target
-
- @torch.no_grad()
- def infer(
- self,
- seqs: Union[str, List[str]],
- position_ids=None,
- ):
- if isinstance(seqs, str):
- lst = [seqs]
- else:
- lst = seqs
- # Returns the raw outputs of the model given an input sequence.
- device = next(self.parameters()).device
- aatype = collate_dense_tensors(
- [
- torch.from_numpy(
- residue_constants.sequence_to_onehot(
- sequence=seq,
- mapping=residue_constants.restype_order_with_x,
- map_unknown_to_x=True,
- )
- )
- .to(device)
- .argmax(dim=1)
- for seq in lst
- ]
- ) # B=1 x L
- mask = collate_dense_tensors([aatype.new_ones(len(seq)) for seq in lst])
- position_ids = (
- torch.arange(aatype.shape[1], device=device).expand(len(lst), -1)
- if position_ids is None
- else position_ids.to(device)
- )
- if position_ids.ndim == 1:
- position_ids = position_ids.unsqueeze(0)
- return self.forward(
- aatype,
- mask,
- position_ids=position_ids,
- )
-
- @staticmethod
- def output_to_pdb(output: Dict) -> List[str]:
- """Returns the pbd (file) string from the model given the model output."""
- output = {k: v.to("cpu").numpy() for k, v in output.items()}
- pdbs = []
- final_atom_positions = atom14_to_atom37(output["positions"][-1], output)
- final_atom_mask = output["atom37_atom_exists"]
- for i in range(output["aatype"].shape[0]):
- aa = output["aatype"][i]
- pred_pos = final_atom_positions[i]
- mask = final_atom_mask[i]
- resid = output["residue_index"][i] + 1
- pred = OFProtein(
- aatype=aa,
- atom_positions=pred_pos,
- atom_mask=mask,
- residue_index=resid,
- b_factors=output["plddt"][i],
- )
- pdbs.append(to_pdb(pred))
- return pdbs
-
- def infer_pdb(self, seqs, *args, **kwargs) -> str:
- """Returns the pdb (file) string from the model given an input sequence."""
- assert isinstance(seqs, str)
- output = self.infer(seqs, *args, **kwargs)
- return self.output_to_pdb(output)[0]
-
- def infer_pdbs(self, seqs: List[str], *args, **kwargs) -> List[str]:
- """Returns the pdb (file) string from the model given an input sequence."""
- output = self.infer(seqs, *args, **kwargs)
- return self.output_to_pdb(output)
diff --git a/transformers/models/esm/modeling_tf_esm.py b/transformers/models/esm/modeling_tf_esm.py
deleted file mode 100644
index 2688c207b0adaca4ee79c37c8529694f608490b6..0000000000000000000000000000000000000000
--- a/transformers/models/esm/modeling_tf_esm.py
+++ /dev/null
@@ -1,1567 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch ESM model."""
-
-
-from __future__ import annotations
-
-import os
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_tf_outputs import (
- TFBaseModelOutputWithPastAndCrossAttentions,
- TFBaseModelOutputWithPoolingAndCrossAttentions,
- TFMaskedLMOutput,
- TFSequenceClassifierOutput,
- TFTokenClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFMaskedLanguageModelingLoss,
- TFModelInputType,
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- TFTokenClassificationLoss,
- get_initializer,
- keras,
- shape_list,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, stable_softmax
-from ...utils import logging
-from .configuration_esm import EsmConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "facebook/esm2_t6_8M_UR50D"
-_CONFIG_FOR_DOC = "EsmConfig"
-
-
-def rotate_half(x):
- x1, x2 = tf.split(x, 2, axis=-1)
- return tf.concat((-x2, x1), axis=-1)
-
-
-def apply_rotary_pos_emb(x, cos, sin):
- cos = cos[:, :, : tf.shape(x)[-2], :]
- sin = sin[:, :, : tf.shape(x)[-2], :]
-
- return (x * cos) + (rotate_half(x) * sin)
-
-
-def symmetrize(x):
- "Make layer symmetric in final two dimensions, used for contact prediction."
- return x + tf.linalg.matrix_transpose(x) # Transposes last two dimensions only
-
-
-def average_product_correct(x):
- "Perform average product correct, used for contact prediction."
- a1 = tf.reduce_sum(x, -1, keepdims=True)
- a2 = tf.reduce_sum(x, -2, keepdims=True)
- a12 = tf.reduce_sum(x, (-1, -2), keepdims=True)
-
- avg = a1 * a2
- avg = avg / a12
- normalized = x - avg
- return normalized
-
-
-class TFRotaryEmbedding(keras.layers.Layer):
- """
- Rotary position embeddings based on those in
- [RoFormer](https://huggingface.co/docs/transformers/model_doc/roformer). Query and keys are transformed by rotation
- matrices which depend on their relative positions.
- """
-
- def __init__(self, dim: int, name=None):
- super().__init__(name=name)
- # Matt: The PyTorch version of this layer does a lot of work to cache values, but we just rely on TF compilation
- # and/or XLA to sort out constants like that. It actually may not seem like this layer needs to be stateful at
- # all when we benefit from TF compilation, but it does. The reason is that self.inv_freq is a buffer in the
- # original implementation, but all the shared ESM checkpoints were trained with fp16 params. This means that
- # the inv_freq tensor was stored as a float16, and we need to replicate those lower-precision values or our
- # models give different outputs from the original.
- self.dim = dim
-
- def build(self, input_shape):
- super().build(input_shape)
- self.inv_freq = self.add_weight(
- "inv_freq", shape=(self.dim // 2,), dtype=tf.float32, initializer=get_initializer(1.0), trainable=False
- )
- self.inv_freq.assign(
- 1.0 / (10000 ** (tf.range(start=0, limit=self.dim, delta=2, dtype=tf.float32) / self.dim))
- )
-
- def _compute_cos_sin(self, x, seq_dimension=2):
- seq_len = tf.shape(x)[seq_dimension]
-
- t = tf.range(seq_len, dtype=self.inv_freq.dtype)
- freqs = tf.einsum("i, j -> ij", t, self.inv_freq) # Outer multiplication
- emb = tf.concat((freqs, freqs), axis=-1)[None, None, :, :]
-
- return tf.cos(emb), tf.sin(emb)
-
- def call(self, q: tf.Tensor, k: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
- cos_emb, sin_emb = self._compute_cos_sin(k, seq_dimension=-2)
-
- return (
- apply_rotary_pos_emb(q, cos_emb, sin_emb),
- apply_rotary_pos_emb(k, cos_emb, sin_emb),
- )
-
-
-class TFEsmContactPredictionHead(keras.layers.Layer):
- """Performs symmetrization, apc, and computes a logistic regression on the output features"""
-
- def __init__(
- self,
- in_features: int,
- bias=True,
- eos_idx: int = 2,
- name=None,
- ):
- super().__init__(name=name)
- self.eos_idx = eos_idx
- self.in_features = in_features
- self.regression = keras.layers.Dense(1, use_bias=bias, activation="sigmoid", name="regression")
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "regression", None) is not None:
- with tf.name_scope(self.regression.name):
- self.regression.build((None, self.in_features))
-
- def call(self, tokens, attentions):
- # remove eos token attentions
- eos_mask = tf.cast(tokens != self.eos_idx, attentions.dtype)
- eos_mask = tf.expand_dims(eos_mask, 1) * tf.expand_dims(eos_mask, 2)
- attentions = attentions * eos_mask[:, None, None, :, :]
- attentions = attentions[..., :-1, :-1]
- # remove cls token attentions
- attentions = attentions[..., 1:, 1:]
- batch_size, layers, heads, seqlen, _ = shape_list(attentions)
- attentions = tf.reshape(attentions, (batch_size, layers * heads, seqlen, seqlen))
-
- # features: batch x channels x tokens x tokens (symmetric)
- attentions = average_product_correct(symmetrize(attentions))
- attentions = tf.transpose(attentions, perm=(0, 2, 3, 1))
- return tf.squeeze(self.regression(attentions), 3)
-
-
-class TFEsmEmbeddings(keras.layers.Layer):
- """
- Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
- """
-
- def __init__(self, config, name=None):
- super().__init__(name=name)
- self.word_embeddings = keras.layers.Embedding(
- config.vocab_size,
- config.hidden_size,
- embeddings_initializer=get_initializer(config.initializer_range),
- name="word_embeddings",
- )
- self.position_embeddings = keras.layers.Embedding(
- config.max_position_embeddings,
- config.hidden_size,
- embeddings_initializer=get_initializer(config.initializer_range),
- name="position_embeddings",
- )
-
- if config.emb_layer_norm_before:
- self.layer_norm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
- else:
- self.layer_norm = None
- # Matt: I think this line was copied incorrectly from BERT, disabling for now
- # self.dropout = Dropout(config.hidden_dropout_prob)
- self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
-
- self.position_ids = tf.range(config.max_position_embeddings)[None, :]
-
- self.padding_idx = config.pad_token_id
- self.token_dropout = config.token_dropout
- self.mask_token_id = config.mask_token_id
- self.config = config
-
- def call(
- self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
- ):
- if position_ids is None:
- if input_ids is not None:
- # Create the position ids from the input token ids. Any padded tokens remain padded.
- position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)
- else:
- position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)
-
- if inputs_embeds is None:
- check_embeddings_within_bounds(input_ids, self.config.vocab_size)
- inputs_embeds = self.word_embeddings(input_ids)
-
- # Note that if we want to support ESM-1 (not 1b!) in future then we need to support an
- # embedding_scale factor here.
- embeddings = inputs_embeds
-
- # Matt: ESM has the option to handle masking in MLM in a slightly unusual way. If the token_dropout
- # flag is False then it is handled in the same was as BERT/RoBERTa. If it is set to True, however,
- # masked tokens are treated as if they were selected for input dropout and zeroed out.
- # This "mask-dropout" is compensated for when masked tokens are not present, by scaling embeddings by
- # a factor of (fraction of unmasked tokens during training) / (fraction of unmasked tokens in sample).
- # This is analogous to the way that dropout layers scale down outputs during evaluation when not
- # actually dropping out values (or, equivalently, scale up their un-dropped outputs in training).
- if self.token_dropout:
- embeddings = tf.where((input_ids == self.mask_token_id)[:, :, None], 0.0, embeddings)
- mask_ratio_train = 0.15 * 0.8 # Hardcoded as the ratio used in all ESM model training runs
- src_lengths = tf.cast(tf.reduce_sum(attention_mask, axis=-1), tf.float32)
- masked_tokens = input_ids == self.mask_token_id
- mask_ratio_observed = tf.math.count_nonzero(masked_tokens, dtype=tf.float32, axis=-1) / src_lengths
- embeddings = embeddings * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]
-
- if self.position_embedding_type == "absolute":
- position_embeddings = self.position_embeddings(position_ids)
- embeddings += position_embeddings
-
- if self.layer_norm is not None:
- embeddings = self.layer_norm(embeddings)
- if attention_mask is not None:
- embeddings = embeddings * tf.cast(tf.expand_dims(attention_mask, -1), embeddings.dtype)
- # Matt: I think this line was copied incorrectly from BERT, disabling it for now.
- # embeddings = self.dropout(embeddings)
- return embeddings
-
- def create_position_ids_from_inputs_embeds(self, inputs_embeds):
- """
- We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
-
- Args:
- inputs_embeds: tf.Tensor
-
- Returns: tf.Tensor
- """
- input_shape = shape_list(inputs_embeds)[:-1]
- sequence_length = input_shape[1]
-
- position_ids = tf.range(
- start=self.padding_idx + 1, limit=sequence_length + self.padding_idx + 1, dtype=tf.int64
- )
- return tf.broadcast_to(tf.expand_dims(position_ids, 0), input_shape)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "word_embeddings", None) is not None:
- with tf.name_scope(self.word_embeddings.name):
- self.word_embeddings.build(None)
- if getattr(self, "position_embeddings", None) is not None:
- with tf.name_scope(self.position_embeddings.name):
- self.position_embeddings.build(None)
- if getattr(self, "layer_norm", None) is not None:
- with tf.name_scope(self.layer_norm.name):
- self.layer_norm.build([None, None, self.config.hidden_size])
-
-
-class TFEsmSelfAttention(keras.layers.Layer):
- def __init__(self, config, position_embedding_type=None, name=None):
- super().__init__(name=name)
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = keras.layers.Dense(
- self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"
- )
- self.key = keras.layers.Dense(
- self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key"
- )
- self.value = keras.layers.Dense(
- self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"
- )
-
- self.dropout = keras.layers.Dropout(config.attention_probs_dropout_prob)
- self.position_embedding_type = position_embedding_type or getattr(
- config, "position_embedding_type", "absolute"
- )
- self.rotary_embeddings = None
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- self.max_position_embeddings = config.max_position_embeddings
- self.distance_embedding = keras.layers.Embedding(
- 2 * config.max_position_embeddings - 1,
- self.attention_head_size,
- embeddings_initializer=get_initializer(config.initializer_range),
- )
- elif self.position_embedding_type == "rotary":
- self.rotary_embeddings = TFRotaryEmbedding(dim=self.attention_head_size, name="rotary_embeddings")
-
- self.is_decoder = config.is_decoder
- self.config = config
-
- def transpose_for_scores(self, x: tf.Tensor) -> tf.Tensor:
- new_x_shape = shape_list(x)[:-1] + [self.num_attention_heads, self.attention_head_size]
- x = tf.reshape(x, new_x_shape)
- return tf.transpose(x, perm=(0, 2, 1, 3))
-
- def call(
- self,
- hidden_states: tf.Tensor,
- attention_mask: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- encoder_hidden_states: tf.Tensor | None = None,
- encoder_attention_mask: tf.Tensor | None = None,
- past_key_value: Tuple[Tuple[tf.Tensor]] | None = None,
- output_attentions: Optional[bool] = False,
- training: bool = False,
- ) -> Tuple[tf.Tensor]:
- mixed_query_layer = self.query(hidden_states)
-
- # If this is instantiated as a cross-attention module, the keys
- # and values come from an encoder; the attention mask needs to be
- # such that the encoder's padding tokens are not attended to.
- is_cross_attention = encoder_hidden_states is not None
-
- if is_cross_attention and past_key_value is not None:
- # reuse k,v, cross_attentions
- key_layer = past_key_value[0]
- value_layer = past_key_value[1]
- attention_mask = encoder_attention_mask
- elif is_cross_attention:
- key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
- value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
- attention_mask = encoder_attention_mask
- elif past_key_value is not None:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- key_layer = tf.concat([past_key_value[0], key_layer], axis=2)
- value_layer = tf.concat([past_key_value[1], value_layer], axis=2)
- else:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Matt: Our BERT model (which this code was derived from) scales attention logits down by sqrt(head_dim).
- # ESM scales the query down by the same factor instead. Modulo numerical stability these are equivalent,
- # but not when rotary embeddings get involved. Therefore, we scale the query here to match the original
- # ESM code and fix rotary embeddings.
- query_layer = query_layer * self.attention_head_size**-0.5
-
- if self.is_decoder:
- # if cross_attention save Tuple(tf.Tensor, tf.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(tf.Tensor, tf.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_layer, value_layer)
-
- if self.position_embedding_type == "rotary":
- query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
-
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- seq_length = shape_list(hidden_states)[1]
- position_ids_l = tf.expand_dims(tf.range(seq_length, dtype=tf.int64), -1)
- position_ids_r = tf.expand_dims(tf.range(seq_length, dtype=tf.int64), 0)
- distance = position_ids_l - position_ids_r
- positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
- positional_embedding = tf.cast(positional_embedding, query_layer.dtype) # fp16 compatibility
-
- if self.position_embedding_type == "relative_key":
- relative_position_scores = tf.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores
- elif self.position_embedding_type == "relative_key_query":
- relative_position_scores_query = tf.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- relative_position_scores_key = tf.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
-
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in EsmModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = stable_softmax(attention_scores, axis=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs, training=training)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = attention_probs @ value_layer
-
- context_layer = tf.transpose(context_layer, perm=(0, 2, 1, 3))
- new_context_layer_shape = shape_list(context_layer)[:-2] + [self.all_head_size]
- context_layer = tf.reshape(context_layer, new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- if self.is_decoder:
- outputs = outputs + (past_key_value,)
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "query", None) is not None:
- with tf.name_scope(self.query.name):
- self.query.build([None, None, self.config.hidden_size])
- if getattr(self, "key", None) is not None:
- with tf.name_scope(self.key.name):
- self.key.build([None, None, self.config.hidden_size])
- if getattr(self, "value", None) is not None:
- with tf.name_scope(self.value.name):
- self.value.build([None, None, self.config.hidden_size])
- if getattr(self, "rotary_embeddings", None) is not None:
- with tf.name_scope(self.rotary_embeddings.name):
- self.rotary_embeddings.build(None)
-
-
-class TFEsmSelfOutput(keras.layers.Layer):
- def __init__(self, config, name=None):
- super().__init__(name=name)
- self.dense = keras.layers.Dense(
- config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states, input_tensor, training=False):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states, training=training)
- hidden_states += input_tensor
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFEsmAttention(keras.layers.Layer):
- def __init__(self, config, name=None):
- super().__init__(name=name)
- self.self = TFEsmSelfAttention(config, name="self")
- self.output_layer = TFEsmSelfOutput(config, name="output")
- self.pruned_heads = set()
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.config = config
-
- def prune_heads(self, heads):
- raise NotImplementedError
-
- def call(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- past_key_value=None,
- output_attentions=False,
- training=False,
- ):
- hidden_states_ln = self.LayerNorm(hidden_states)
- self_outputs = self.self(
- hidden_states_ln,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- training,
- )
- attention_output = self.output_layer(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "self", None) is not None:
- with tf.name_scope(self.self.name):
- self.self.build(None)
- if getattr(self, "output_layer", None) is not None:
- with tf.name_scope(self.output_layer.name):
- self.output_layer.build(None)
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
-
-
-class TFEsmIntermediate(keras.layers.Layer):
- def __init__(self, config: EsmConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.intermediate_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="dense",
- )
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.dense(inputs=hidden_states)
- hidden_states = tf.nn.gelu(hidden_states)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFEsmOutput(keras.layers.Layer):
- def __init__(self, config, name=None):
- super().__init__(name=name)
- self.dense = keras.layers.Dense(
- config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
- self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
- self.config = config
-
- def call(self, hidden_states, input_tensor, training=False):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states, training=training)
- hidden_states += input_tensor
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.intermediate_size])
-
-
-class TFEsmLayer(keras.layers.Layer):
- def __init__(self, config, name=None):
- super().__init__(name=name)
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = TFEsmAttention(config, name="attention")
- self.is_decoder = config.is_decoder
- self.add_cross_attention = config.add_cross_attention
- if self.add_cross_attention:
- if not self.is_decoder:
- raise RuntimeError(f"{self} should be used as a decoder model if cross attention is added")
- self.crossattention = TFEsmAttention(config)
- self.intermediate = TFEsmIntermediate(config, name="intermediate")
- self.output_layer = TFEsmOutput(config, name="output")
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
- self.config = config
-
- def call(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- past_key_value=None,
- output_attentions=False,
- training=False,
- ):
- # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
- self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
- self_attention_outputs = self.attention(
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions=output_attentions,
- past_key_value=self_attn_past_key_value,
- training=training,
- )
- attention_output = self_attention_outputs[0]
-
- # if decoder, the last output is tuple of self-attn cache
- if self.is_decoder:
- outputs = self_attention_outputs[1:-1]
- present_key_value = self_attention_outputs[-1]
- else:
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- cross_attn_present_key_value = None
- if self.is_decoder and encoder_hidden_states is not None:
- if not hasattr(self, "crossattention"):
- raise AttributeError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated"
- " with cross-attention layers by setting `config.add_cross_attention=True`"
- )
-
- # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
- cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
- cross_attention_outputs = self.crossattention(
- attention_output,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- cross_attn_past_key_value,
- output_attentions,
- training=training,
- )
- attention_output = cross_attention_outputs[0]
- outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
-
- # add cross-attn cache to positions 3,4 of present_key_value tuple
- cross_attn_present_key_value = cross_attention_outputs[-1]
- present_key_value = present_key_value + cross_attn_present_key_value
-
- layernorm_output = self.LayerNorm(attention_output)
- intermediate_output = self.intermediate(hidden_states=layernorm_output)
- layer_output = self.output_layer(
- hidden_states=intermediate_output, input_tensor=attention_output, training=training
- )
- outputs = (layer_output,) + outputs # add attentions if we output them
-
- # if decoder, return the attn key/values as the last output
- if self.is_decoder:
- outputs = outputs + (present_key_value,)
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "intermediate", None) is not None:
- with tf.name_scope(self.intermediate.name):
- self.intermediate.build(None)
- if getattr(self, "output_layer", None) is not None:
- with tf.name_scope(self.output_layer.name):
- self.output_layer.build(None)
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.hidden_size])
-
-
-class TFEsmEncoder(keras.layers.Layer):
- def __init__(self, config, name=None):
- super().__init__(name=name)
- self.config = config
- self.layer = [TFEsmLayer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]
- self.emb_layer_norm_after = keras.layers.LayerNormalization(
- epsilon=config.layer_norm_eps, name="emb_layer_norm_after"
- )
-
- def call(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- past_key_values=None,
- use_cache=None,
- output_attentions=False,
- output_hidden_states=False,
- return_dict=True,
- training=False,
- ):
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
-
- next_decoder_cache = () if use_cache else None
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
- past_key_value = past_key_values[i] if past_key_values is not None else None
-
- layer_outputs = layer_module(
- hidden_states,
- attention_mask,
- layer_head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- past_key_value,
- output_attentions,
- training,
- )
-
- hidden_states = layer_outputs[0]
- if use_cache:
- next_decoder_cache += (layer_outputs[-1],)
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
-
- if self.emb_layer_norm_after:
- hidden_states = self.emb_layer_norm_after(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- next_decoder_cache,
- all_hidden_states,
- all_self_attentions,
- all_cross_attentions,
- ]
- if v is not None
- )
- return TFBaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=next_decoder_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "emb_layer_norm_after", None) is not None:
- with tf.name_scope(self.emb_layer_norm_after.name):
- self.emb_layer_norm_after.build([None, None, self.config.hidden_size])
- if getattr(self, "layer", None) is not None:
- for layer in self.layer:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-# Copied from transformers.models.bert.modeling_tf_bert.TFBertPooler with Bert->Esm
-class TFEsmPooler(keras.layers.Layer):
- def __init__(self, config: EsmConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.dense = keras.layers.Dense(
- units=config.hidden_size,
- kernel_initializer=get_initializer(config.initializer_range),
- activation="tanh",
- name="dense",
- )
- self.config = config
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(inputs=first_token_tensor)
-
- return pooled_output
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
-
-
-class TFEsmPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = EsmConfig
- base_model_prefix = "esm"
-
-
-ESM_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a Keras [Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it as a
- regular Keras model and refer to the TF/Keras documentation for all matters related to general usage and behavior.
-
- Parameters:
- config ([`EsmConfig`]): Model configuration class with all the parameters of the
- model. Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-ESM_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`tf.Tensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`tf.Tensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`tf.Tensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare ESM Model transformer outputting raw hidden-states without any specific head on top.",
- ESM_START_DOCSTRING,
-)
-class TFEsmMainLayer(keras.layers.Layer):
- """
-
- The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
- cross-attention is added between the self-attention layers, following the architecture described in [Attention is
- all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
- Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
-
- To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
- to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
- `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
- """
-
- _keys_to_ignore_on_load_missing = [r"position_ids"]
-
- def __init__(self, config, add_pooling_layer=True, name=None, **kwargs):
- super().__init__(name=name, **kwargs)
-
- self.config = config
- self.is_decoder = config.is_decoder
-
- self.embeddings = TFEsmEmbeddings(config, name="embeddings")
- self.encoder = TFEsmEncoder(config, name="encoder")
- self.pooler = TFEsmPooler(config, name="pooler") if add_pooling_layer else None
-
- self.contact_head = TFEsmContactPredictionHead(
- in_features=self.config.num_hidden_layers * self.config.num_attention_heads, bias=True, name="contact_head"
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "pooler", None) is not None:
- with tf.name_scope(self.pooler.name):
- self.pooler.build(None)
- if getattr(self, "contact_head", None) is not None:
- with tf.name_scope(self.contact_head.name):
- self.contact_head.build(None)
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value: tf.Variable):
- self.embeddings.word_embeddings.weight = value
- self.embeddings.vocab_size = shape_list(value)[0]
-
- def _prune_heads(self, heads_to_prune):
- raise NotImplementedError
-
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- encoder_hidden_states: np.ndarray | tf.Tensor | None = None,
- encoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFBaseModelOutputWithPoolingAndCrossAttentions, Tuple[tf.Tensor]]:
- if not self.config.is_decoder:
- use_cache = False
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
-
- if past_key_values is None:
- past_key_values_length = 0
- past_key_values = [None] * len(self.encoder.layer)
- else:
- past_key_values_length = shape_list(past_key_values[0][0])[-2]
-
- if attention_mask is None:
- attention_mask = tf.fill(dims=(batch_size, seq_length + past_key_values_length), value=1)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- past_key_values_length=past_key_values_length,
- training=training,
- )
-
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask_shape = shape_list(attention_mask)
-
- mask_seq_length = seq_length + past_key_values_length
- # Copied from `modeling_tf_t5.py`
- # Provided a padding mask of dimensions [batch_size, mask_seq_length]
- # - if the model is a decoder, apply a causal mask in addition to the padding mask
- # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]
- if self.is_decoder:
- seq_ids = tf.range(mask_seq_length)
- causal_mask = tf.less_equal(
- tf.tile(seq_ids[None, None, :], (batch_size, mask_seq_length, 1)),
- seq_ids[None, :, None],
- )
- causal_mask = tf.cast(causal_mask, dtype=attention_mask.dtype)
- extended_attention_mask = causal_mask * attention_mask[:, None, :]
- attention_mask_shape = shape_list(extended_attention_mask)
- extended_attention_mask = tf.reshape(
- extended_attention_mask, (attention_mask_shape[0], 1, attention_mask_shape[1], attention_mask_shape[2])
- )
- if past_key_values[0] is not None:
- # attention_mask needs to be sliced to the shape `[batch_size, 1, from_seq_length - cached_seq_length, to_seq_length]
- extended_attention_mask = extended_attention_mask[:, :, -seq_length:, :]
- else:
- extended_attention_mask = tf.reshape(
- attention_mask, (attention_mask_shape[0], 1, 1, attention_mask_shape[1])
- )
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and -10000.0 for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- extended_attention_mask = tf.cast(extended_attention_mask, dtype=embedding_output.dtype)
- one_cst = tf.constant(1.0, dtype=embedding_output.dtype)
- ten_thousand_cst = tf.constant(-10000.0, dtype=embedding_output.dtype)
- extended_attention_mask = tf.multiply(tf.subtract(one_cst, extended_attention_mask), ten_thousand_cst)
-
- # Copied from `modeling_tf_t5.py` with -1e9 -> -10000
- if self.is_decoder and encoder_attention_mask is not None:
- # If a 2D ou 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- encoder_attention_mask = tf.cast(encoder_attention_mask, dtype=extended_attention_mask.dtype)
- num_dims_encoder_attention_mask = len(shape_list(encoder_attention_mask))
- if num_dims_encoder_attention_mask == 3:
- encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
- if num_dims_encoder_attention_mask == 2:
- encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
-
- # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
- # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow/transformer/transformer_layers.py#L270
- # encoder_extended_attention_mask = tf.math.equal(encoder_extended_attention_mask,
- # tf.transpose(encoder_extended_attention_mask, perm=(-1, -2)))
-
- encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -10000.0
- else:
- encoder_extended_attention_mask = None
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.config.num_hidden_layers
-
- encoder_outputs = self.encoder(
- hidden_states=embedding_output,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_extended_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = encoder_outputs[0]
- pooled_output = self.pooler(hidden_states=sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- return (
- sequence_output,
- pooled_output,
- ) + encoder_outputs[1:]
-
- return TFBaseModelOutputWithPoolingAndCrossAttentions(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- past_key_values=encoder_outputs.past_key_values,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- cross_attentions=encoder_outputs.cross_attentions,
- )
-
- def predict_contacts(self, tokens, attention_mask):
- attns = self(tokens, attention_mask=attention_mask, return_dict=True, output_attentions=True).attentions
- attns = tf.stack(attns, axis=1) # Matches the original model layout
- # In the original model, attentions for padding tokens are completely zeroed out.
- # This makes no difference most of the time because the other tokens won't attend to them,
- # but it does for the contact prediction task, which takes attentions as input,
- # so we have to mimic that here.
- attention_mask = tf.cast(attention_mask, attns.dtype)
- attns *= attention_mask[:, None, None, None]
- attns *= attention_mask[:, None, None, :, None]
- return self.contact_head(tokens, attns)
-
-
-@add_start_docstrings(
- "The bare ESM Model transformer outputting raw hidden-states without any specific head on top.",
- ESM_START_DOCSTRING,
-)
-class TFEsmModel(TFEsmPreTrainedModel):
- def __init__(self, config: EsmConfig, add_pooling_layer=True, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.esm = TFEsmMainLayer(config, add_pooling_layer=add_pooling_layer, name="esm")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ESM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutputWithPoolingAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- encoder_hidden_states: np.ndarray | tf.Tensor | None = None,
- encoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithPoolingAndCrossAttentions, Tuple[tf.Tensor]]:
- r"""
- encoder_hidden_states (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- past_key_values (`Tuple[Tuple[tf.Tensor]]` of length `config.n_layers`)
- contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*, defaults to `True`):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`). Set to `False` during training, `True` during generation
- """
- outputs = self.esm(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- return outputs
-
- def predict_contacts(self, tokens, attention_mask):
- return self.esm.predict_contacts(tokens, attention_mask)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "esm", None) is not None:
- with tf.name_scope(self.esm.name):
- self.esm.build(None)
-
-
-@add_start_docstrings("""ESM Model with a `language modeling` head on top.""", ESM_START_DOCSTRING)
-class TFEsmForMaskedLM(TFEsmPreTrainedModel, TFMaskedLanguageModelingLoss):
- _keys_to_ignore_on_load_missing = [r"position_ids"]
- _keys_to_ignore_on_load_unexpected = [r"pooler"]
-
- def __init__(self, config):
- super().__init__(config)
-
- if config.is_decoder:
- logger.warning(
- "If you want to use `EsmForMaskedLM` make sure `config.is_decoder=False` for "
- "bi-directional self-attention."
- )
-
- self.esm = TFEsmMainLayer(config, add_pooling_layer=False, name="esm")
- self.lm_head = TFEsmLMHead(config, name="lm_head")
- if config.tie_word_embeddings:
- # Ensure word embeddings are built so that we actually have something to tie
- with tf.name_scope(os.path.join(self._name_scope(), "esm", "embeddings", "word_embeddings")):
- self.esm.embeddings.word_embeddings.build((None, None))
- self.lm_head.decoder = self.esm.embeddings.word_embeddings.weights[0]
-
- def get_output_embeddings(self):
- return self.lm_head.decoder
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head.decoder = new_embeddings
-
- def get_lm_head(self):
- return self.lm_head
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ESM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="",
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- encoder_hidden_states: np.ndarray | tf.Tensor | None = None,
- encoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- labels: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFMaskedLMOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- kwargs (`Dict[str, any]`, optional, defaults to *{}*):
- Used to hide legacy arguments that have been deprecated.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.esm(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- prediction_scores = self.lm_head(sequence_output)
-
- masked_lm_loss = None
- if labels is not None:
- masked_lm_loss = self.hf_compute_loss(labels=labels, logits=prediction_scores)
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return TFMaskedLMOutput(
- loss=masked_lm_loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def predict_contacts(self, tokens, attention_mask):
- return self.esm.predict_contacts(tokens, attention_mask)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "esm", None) is not None:
- with tf.name_scope(self.esm.name):
- self.esm.build(None)
- if getattr(self, "lm_head", None) is not None:
- with tf.name_scope(self.lm_head.name):
- self.lm_head.build(None)
-
-
-class TFEsmLMHead(keras.layers.Layer):
- """ESM Head for masked language modeling."""
-
- def __init__(self, config, name=None):
- super().__init__(name=name)
- self.dense = keras.layers.Dense(
- config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
- )
-
- self.layer_norm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
- if config.tie_word_embeddings:
- self.decoder = None
- else:
- self.decoder = keras.layers.Dense(
- config.vocab_size,
- kernel_initializer=get_initializer(config.initializer_range),
- name="decoder",
- use_bias=False,
- )
- self.config = config
-
- def build(self, input_shape=None):
- # Separate bias to match the PT model and allow weight cross-loading to work
- # Put it in the build so it gets the right name when adding it as a weight
- if self.built:
- return
- self.built = True
- self.bias = self.add_weight("bias", shape=(self.config.vocab_size,), initializer="zeros", trainable=True)
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "layer_norm", None) is not None:
- with tf.name_scope(self.layer_norm.name):
- self.layer_norm.build([None, None, self.config.hidden_size])
- if getattr(self, "decoder", None) is not None and not self.config.tie_word_embeddings:
- with tf.name_scope(self.decoder.name):
- self.decoder.build([None, None, self.config.hidden_size])
-
- def get_bias(self):
- return {"bias": self.bias}
-
- def call(self, features):
- x = self.dense(features)
- x = tf.nn.gelu(x)
- x = self.layer_norm(x)
-
- # project back to size of vocabulary with bias
- if self.config.tie_word_embeddings:
- x = tf.matmul(x, self.decoder, transpose_b=True) + self.bias
- else:
- x = self.decoder(x) + self.bias
- return x
-
-
-@add_start_docstrings(
- """
- ESM Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
- output) e.g. for GLUE tasks.
- """,
- ESM_START_DOCSTRING,
-)
-class TFEsmForSequenceClassification(TFEsmPreTrainedModel, TFSequenceClassificationLoss):
- _keys_to_ignore_on_load_missing = [r"position_ids"]
-
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
-
- self.esm = TFEsmMainLayer(config, add_pooling_layer=False, name="esm")
- self.classifier = TFEsmClassificationHead(config, name="classifier")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ESM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- labels: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.esm(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- logits = self.classifier(sequence_output)
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "esm", None) is not None:
- with tf.name_scope(self.esm.name):
- self.esm.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build(None)
-
-
-@add_start_docstrings(
- """
- ESM Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- ESM_START_DOCSTRING,
-)
-class TFEsmForTokenClassification(TFEsmPreTrainedModel, TFTokenClassificationLoss):
- _keys_to_ignore_on_load_unexpected = [r"pooler"]
- _keys_to_ignore_on_load_missing = [r"position_ids"]
-
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.esm = TFEsmMainLayer(config, add_pooling_layer=False, name="esm")
- self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
- self.classifier = keras.layers.Dense(config.num_labels, name="classifier")
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(ESM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- labels: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[TFTokenClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.esm(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output, training=training)
- logits = self.classifier(sequence_output)
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFTokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "esm", None) is not None:
- with tf.name_scope(self.esm.name):
- self.esm.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-class TFEsmClassificationHead(keras.layers.Layer):
- """Head for sentence-level classification tasks."""
-
- def __init__(self, config, name=None):
- super().__init__(name=name)
- self.dense = keras.layers.Dense(
- config.hidden_size,
- kernel_initializer=get_initializer(config.initializer_range),
- activation="tanh",
- name="dense",
- )
- self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)
- self.out_proj = keras.layers.Dense(
- config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- activation="linear",
- name="out_proj",
- )
- self.config = config
-
- def call(self, features, training=False):
- x = features[:, 0, :] # take token (equiv. to [CLS])
- x = self.dropout(x, training=training)
- x = self.dense(x)
- x = self.dropout(x, training=training)
- x = self.out_proj(x)
- return x
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.hidden_size])
- if getattr(self, "out_proj", None) is not None:
- with tf.name_scope(self.out_proj.name):
- self.out_proj.build([None, None, self.config.hidden_size])
-
-
-def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
- """
- Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
- are ignored. This is modified from fairseq's `utils.make_positions`.
-
- Args:
- x: tf.Tensor x:
-
- Returns: tf.Tensor
- """
- # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
- mask = tf.cast(input_ids != padding_idx, tf.int64)
- incremental_indices = (tf.cumsum(mask, axis=1) + past_key_values_length) * mask
- return incremental_indices + padding_idx
diff --git a/transformers/models/esm/openfold_utils/__init__.py b/transformers/models/esm/openfold_utils/__init__.py
deleted file mode 100644
index 02a8c149ae320dd9b045edc5df31760a4eebefd9..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/__init__.py
+++ /dev/null
@@ -1,8 +0,0 @@
-from .chunk_utils import chunk_layer
-from .data_transforms import make_atom14_masks
-from .feats import atom14_to_atom37, frames_and_literature_positions_to_atom14_pos, torsion_angles_to_frames
-from .loss import compute_predicted_aligned_error, compute_tm
-from .protein import Protein as OFProtein
-from .protein import to_pdb
-from .rigid_utils import Rigid, Rotation
-from .tensor_utils import dict_multimap, flatten_final_dims, permute_final_dims
diff --git a/transformers/models/esm/openfold_utils/__pycache__/__init__.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index b3eb02531f6d466b4fc5fcac2f785ed66e6f4141..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/__pycache__/chunk_utils.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/chunk_utils.cpython-310.pyc
deleted file mode 100644
index 9d3b8ec0b284e54d3a415033d84540aab0c497ab..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/chunk_utils.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/__pycache__/data_transforms.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/data_transforms.cpython-310.pyc
deleted file mode 100644
index 4bd422669d008dddd168ee446c778e043e7c3fd3..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/data_transforms.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/__pycache__/feats.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/feats.cpython-310.pyc
deleted file mode 100644
index b04353312131b67a0b0260a9b99097ed62001e88..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/feats.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/__pycache__/loss.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/loss.cpython-310.pyc
deleted file mode 100644
index 5dea89a482252c1fba4e1cefe77f3bbcad0d5333..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/loss.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/__pycache__/protein.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/protein.cpython-310.pyc
deleted file mode 100644
index da928f7e65487fafa651a96d39a1efaaa15b9207..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/protein.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/__pycache__/residue_constants.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/residue_constants.cpython-310.pyc
deleted file mode 100644
index a98b2f120b3f321ce023ee0745395646e969a192..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/residue_constants.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/__pycache__/rigid_utils.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/rigid_utils.cpython-310.pyc
deleted file mode 100644
index 22b598ff1b3a90656bed2c0cfed486d83c17f7ee..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/rigid_utils.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/__pycache__/tensor_utils.cpython-310.pyc b/transformers/models/esm/openfold_utils/__pycache__/tensor_utils.cpython-310.pyc
deleted file mode 100644
index 7e8bc0f181df3ca28caaca817dea7db4dd1fe5da..0000000000000000000000000000000000000000
Binary files a/transformers/models/esm/openfold_utils/__pycache__/tensor_utils.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/esm/openfold_utils/chunk_utils.py b/transformers/models/esm/openfold_utils/chunk_utils.py
deleted file mode 100644
index 301721d135ee4d63ff111d45c06471c50c89e925..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/chunk_utils.py
+++ /dev/null
@@ -1,397 +0,0 @@
-# Copyright 2021 AlQuraishi Laboratory
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import logging
-import math
-from functools import partial
-from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union
-
-import torch
-
-from .tensor_utils import tensor_tree_map, tree_map
-
-
-def _fetch_dims(tree: Union[dict, list, tuple, torch.Tensor]) -> List[Tuple[int, ...]]:
- shapes = []
- if isinstance(tree, dict):
- for v in tree.values():
- shapes.extend(_fetch_dims(v))
- elif isinstance(tree, (list, tuple)):
- for t in tree:
- shapes.extend(_fetch_dims(t))
- elif isinstance(tree, torch.Tensor):
- shapes.append(tree.shape)
- else:
- raise ValueError("Not supported")
-
- return shapes
-
-
-@torch.jit.ignore
-def _flat_idx_to_idx(flat_idx: int, dims: Tuple[int, ...]) -> Tuple[int, ...]:
- idx = []
- for d in reversed(dims):
- idx.append(flat_idx % d)
- flat_idx = flat_idx // d
-
- return tuple(reversed(idx))
-
-
-@torch.jit.ignore
-def _get_minimal_slice_set(
- start: Sequence[int],
- end: Sequence[int],
- dims: Sequence[int],
- start_edges: Optional[Sequence[bool]] = None,
- end_edges: Optional[Sequence[bool]] = None,
-) -> List[Tuple[slice, ...]]:
- """
- Produces an ordered sequence of tensor slices that, when used in sequence on a tensor with shape dims, yields
- tensors that contain every leaf in the contiguous range [start, end]. Care is taken to yield a short sequence of
- slices, and perhaps even the shortest possible (I'm pretty sure it's the latter).
-
- end is INCLUSIVE.
- """
-
- # start_edges and end_edges both indicate whether, starting from any given
- # dimension, the start/end index is at the top/bottom edge of the
- # corresponding tensor, modeled as a tree
- def reduce_edge_list(l: List[bool]) -> None:
- tally = True
- for i in range(len(l)):
- reversed_idx = -1 * (i + 1)
- l[reversed_idx] &= tally
- tally = l[reversed_idx]
-
- if start_edges is None:
- start_edges = [s == 0 for s in start]
- reduce_edge_list(start_edges)
- if end_edges is None:
- end_edges = [e == (d - 1) for e, d in zip(end, dims)]
- reduce_edge_list(end_edges)
-
- # Base cases. Either start/end are empty and we're done, or the final,
- # one-dimensional tensor can be simply sliced
- if len(start) == 0:
- return [()]
- elif len(start) == 1:
- return [(slice(start[0], end[0] + 1),)]
-
- slices: List[Tuple[slice, ...]] = []
- path_list: List[slice] = []
-
- # Dimensions common to start and end can be selected directly
- for s, e in zip(start, end):
- if s == e:
- path_list.append(slice(s, s + 1))
- else:
- break
-
- path: Tuple[slice, ...] = tuple(path_list)
- divergence_idx = len(path)
-
- # start == end, and we're done
- if divergence_idx == len(dims):
- return [path]
-
- def upper() -> Tuple[Tuple[slice, ...], ...]:
- assert start_edges is not None
- assert end_edges is not None
-
- sdi = start[divergence_idx]
- return tuple(
- path + (slice(sdi, sdi + 1),) + s
- for s in _get_minimal_slice_set(
- start[divergence_idx + 1 :],
- [d - 1 for d in dims[divergence_idx + 1 :]],
- dims[divergence_idx + 1 :],
- start_edges=start_edges[divergence_idx + 1 :],
- end_edges=[True for _ in end_edges[divergence_idx + 1 :]],
- )
- )
-
- def lower() -> Tuple[Tuple[slice, ...], ...]:
- assert start_edges is not None
- assert end_edges is not None
-
- edi = end[divergence_idx]
- return tuple(
- path + (slice(edi, edi + 1),) + s
- for s in _get_minimal_slice_set(
- [0 for _ in start[divergence_idx + 1 :]],
- end[divergence_idx + 1 :],
- dims[divergence_idx + 1 :],
- start_edges=[True for _ in start_edges[divergence_idx + 1 :]],
- end_edges=end_edges[divergence_idx + 1 :],
- )
- )
-
- # If both start and end are at the edges of the subtree rooted at
- # divergence_idx, we can just select the whole subtree at once
- if start_edges[divergence_idx] and end_edges[divergence_idx]:
- slices.append(path + (slice(start[divergence_idx], end[divergence_idx] + 1),))
- # If just start is at the edge, we can grab almost all of the subtree,
- # treating only the ragged bottom edge as an edge case
- elif start_edges[divergence_idx]:
- slices.append(path + (slice(start[divergence_idx], end[divergence_idx]),))
- slices.extend(lower())
- # Analogous to the previous case, but the top is ragged this time
- elif end_edges[divergence_idx]:
- slices.extend(upper())
- slices.append(path + (slice(start[divergence_idx] + 1, end[divergence_idx] + 1),))
- # If both sides of the range are ragged, we need to handle both sides
- # separately. If there's contiguous meat in between them, we can index it
- # in one big chunk
- else:
- slices.extend(upper())
- middle_ground = end[divergence_idx] - start[divergence_idx]
- if middle_ground > 1:
- slices.append(path + (slice(start[divergence_idx] + 1, end[divergence_idx]),))
- slices.extend(lower())
-
- return slices
-
-
-@torch.jit.ignore
-def _chunk_slice(t: torch.Tensor, flat_start: int, flat_end: int, no_batch_dims: int) -> torch.Tensor:
- """
- Equivalent to
-
- t.reshape((-1,) + t.shape[no_batch_dims:])[flat_start:flat_end]
-
- but without the need for the initial reshape call, which can be memory-intensive in certain situations. The only
- reshape operations in this function are performed on sub-tensors that scale with (flat_end - flat_start), the chunk
- size.
- """
-
- batch_dims = t.shape[:no_batch_dims]
- start_idx = list(_flat_idx_to_idx(flat_start, batch_dims))
- # _get_minimal_slice_set is inclusive
- end_idx = list(_flat_idx_to_idx(flat_end - 1, batch_dims))
-
- # Get an ordered list of slices to perform
- slices = _get_minimal_slice_set(
- start_idx,
- end_idx,
- batch_dims,
- )
-
- sliced_tensors = [t[s] for s in slices]
-
- return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]) for s in sliced_tensors])
-
-
-def chunk_layer(
- layer: Callable,
- inputs: Dict[str, Any],
- chunk_size: int,
- no_batch_dims: int,
- low_mem: bool = False,
- _out: Any = None,
- _add_into_out: bool = False,
-) -> Any:
- """
- Implements the "chunking" procedure described in section 1.11.8.
-
- Layer outputs and inputs are assumed to be simple "pytrees," consisting only of (arbitrarily nested) lists, tuples,
- and dicts with torch.Tensor leaves.
-
- Args:
- layer:
- The layer to be applied chunk-wise
- inputs:
- A (non-nested) dictionary of keyworded inputs. All leaves must be tensors and must share the same batch
- dimensions.
- chunk_size:
- The number of sub-batches per chunk. If multiple batch dimensions are specified, a "sub-batch" is defined
- as a single indexing of all batch dimensions simultaneously (s.t. the number of sub-batches is the product
- of the batch dimensions).
- no_batch_dims:
- How many of the initial dimensions of each input tensor can be considered batch dimensions.
- low_mem:
- Avoids flattening potentially large input tensors. Unnecessary in most cases, and is ever so slightly
- slower than the default setting.
- Returns:
- The reassembled output of the layer on the inputs.
- """
- if not (len(inputs) > 0):
- raise ValueError("Must provide at least one input")
-
- initial_dims = [shape[:no_batch_dims] for shape in _fetch_dims(inputs)]
- orig_batch_dims = tuple([max(s) for s in zip(*initial_dims)])
-
- def _prep_inputs(t: torch.Tensor) -> torch.Tensor:
- if not low_mem:
- if not sum(t.shape[:no_batch_dims]) == no_batch_dims:
- t = t.expand(orig_batch_dims + t.shape[no_batch_dims:])
- t = t.reshape(-1, *t.shape[no_batch_dims:])
- else:
- t = t.expand(orig_batch_dims + t.shape[no_batch_dims:])
- return t
-
- prepped_inputs: Dict[str, Any] = tensor_tree_map(_prep_inputs, inputs)
- prepped_outputs = None
- if _out is not None:
- prepped_outputs = tensor_tree_map(lambda t: t.view([-1] + list(t.shape[no_batch_dims:])), _out)
-
- flat_batch_dim = 1
- for d in orig_batch_dims:
- flat_batch_dim *= d
-
- no_chunks = flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)
-
- def _select_chunk(t: torch.Tensor) -> torch.Tensor:
- return t[i : i + chunk_size] if t.shape[0] != 1 else t
-
- i = 0
- out = prepped_outputs
- for _ in range(no_chunks):
- # Chunk the input
- if not low_mem:
- select_chunk = _select_chunk
- else:
- select_chunk = partial(
- _chunk_slice,
- flat_start=i,
- flat_end=min(flat_batch_dim, i + chunk_size),
- no_batch_dims=len(orig_batch_dims),
- )
-
- chunks: Dict[str, Any] = tensor_tree_map(select_chunk, prepped_inputs)
-
- # Run the layer on the chunk
- output_chunk = layer(**chunks)
-
- # Allocate space for the output
- if out is None:
- out = tensor_tree_map(lambda t: t.new_zeros((flat_batch_dim,) + t.shape[1:]), output_chunk)
-
- # Put the chunk in its pre-allocated space
- if isinstance(output_chunk, dict):
-
- def assign(d1: dict, d2: dict) -> None:
- for k, v in d1.items():
- if isinstance(v, dict):
- assign(v, d2[k])
- else:
- if _add_into_out:
- v[i : i + chunk_size] += d2[k]
- else:
- v[i : i + chunk_size] = d2[k]
-
- assign(out, output_chunk)
- elif isinstance(output_chunk, tuple):
- for x1, x2 in zip(out, output_chunk):
- if _add_into_out:
- x1[i : i + chunk_size] += x2
- else:
- x1[i : i + chunk_size] = x2
- elif isinstance(output_chunk, torch.Tensor):
- if _add_into_out:
- out[i : i + chunk_size] += output_chunk
- else:
- out[i : i + chunk_size] = output_chunk
- else:
- raise ValueError("Not supported")
-
- i += chunk_size
-
- out = tensor_tree_map(lambda t: t.view(orig_batch_dims + t.shape[1:]), out)
-
- return out
-
-
-class ChunkSizeTuner:
- def __init__(
- self,
- # Heuristically, runtimes for most of the modules in the network
- # plateau earlier than this on all GPUs I've run the model on.
- max_chunk_size: int = 512,
- ):
- self.max_chunk_size = max_chunk_size
- self.cached_chunk_size: Optional[int] = None
- self.cached_arg_data: Optional[tuple] = None
-
- def _determine_favorable_chunk_size(self, fn: Callable, args: tuple, min_chunk_size: int) -> int:
- logging.info("Tuning chunk size...")
-
- if min_chunk_size >= self.max_chunk_size:
- return min_chunk_size
-
- candidates: List[int] = [2**l for l in range(int(math.log(self.max_chunk_size, 2)) + 1)]
- candidates = [c for c in candidates if c > min_chunk_size]
- candidates = [min_chunk_size] + candidates
- candidates[-1] += 4
-
- def test_chunk_size(chunk_size: int) -> bool:
- try:
- with torch.no_grad():
- fn(*args, chunk_size=chunk_size)
- return True
- except RuntimeError:
- return False
-
- min_viable_chunk_size_index = 0
- i = len(candidates) - 1
- while i > min_viable_chunk_size_index:
- viable = test_chunk_size(candidates[i])
- if not viable:
- i = (min_viable_chunk_size_index + i) // 2
- else:
- min_viable_chunk_size_index = i
- i = (i + len(candidates) - 1) // 2
-
- return candidates[min_viable_chunk_size_index]
-
- def _compare_arg_caches(self, ac1: Iterable, ac2: Iterable) -> bool:
- consistent = True
- for a1, a2 in zip(ac1, ac2):
- assert type(ac1) == type(ac2)
- if isinstance(ac1, (list, tuple)):
- consistent &= self._compare_arg_caches(a1, a2)
- elif isinstance(ac1, dict):
- a1_items = [v for _, v in sorted(a1.items(), key=lambda x: x[0])]
- a2_items = [v for _, v in sorted(a2.items(), key=lambda x: x[0])]
- consistent &= self._compare_arg_caches(a1_items, a2_items)
- else:
- consistent &= a1 == a2
-
- return consistent
-
- def tune_chunk_size(
- self,
- representative_fn: Callable,
- args: tuple,
- min_chunk_size: int,
- ) -> int:
- consistent = True
- arg_data: tuple = tree_map(lambda a: a.shape if isinstance(a, torch.Tensor) else a, args, object)
- if self.cached_arg_data is not None:
- # If args have changed shape/value, we need to re-tune
- assert len(self.cached_arg_data) == len(arg_data)
- consistent = self._compare_arg_caches(self.cached_arg_data, arg_data)
- else:
- # Otherwise, we can reuse the precomputed value
- consistent = False
-
- if not consistent:
- self.cached_chunk_size = self._determine_favorable_chunk_size(
- representative_fn,
- args,
- min_chunk_size,
- )
- self.cached_arg_data = arg_data
-
- assert self.cached_chunk_size is not None
-
- return self.cached_chunk_size
diff --git a/transformers/models/esm/openfold_utils/data_transforms.py b/transformers/models/esm/openfold_utils/data_transforms.py
deleted file mode 100644
index 8d4c17589ae66df2a8fd0ccfe8d6e335004eed9a..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/data_transforms.py
+++ /dev/null
@@ -1,93 +0,0 @@
-# Copyright 2021 AlQuraishi Laboratory
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Dict
-
-import numpy as np
-import torch
-
-from . import residue_constants as rc
-from .tensor_utils import tensor_tree_map, tree_map
-
-
-def make_atom14_masks(protein: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
- """Construct denser atom positions (14 dimensions instead of 37)."""
- restype_atom14_to_atom37_list = []
- restype_atom37_to_atom14_list = []
- restype_atom14_mask_list = []
-
- for rt in rc.restypes:
- atom_names = rc.restype_name_to_atom14_names[rc.restype_1to3[rt]]
- restype_atom14_to_atom37_list.append([(rc.atom_order[name] if name else 0) for name in atom_names])
- atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)}
- restype_atom37_to_atom14_list.append(
- [(atom_name_to_idx14[name] if name in atom_name_to_idx14 else 0) for name in rc.atom_types]
- )
-
- restype_atom14_mask_list.append([(1.0 if name else 0.0) for name in atom_names])
-
- # Add dummy mapping for restype 'UNK'
- restype_atom14_to_atom37_list.append([0] * 14)
- restype_atom37_to_atom14_list.append([0] * 37)
- restype_atom14_mask_list.append([0.0] * 14)
-
- restype_atom14_to_atom37 = torch.tensor(
- restype_atom14_to_atom37_list,
- dtype=torch.int32,
- device=protein["aatype"].device,
- )
- restype_atom37_to_atom14 = torch.tensor(
- restype_atom37_to_atom14_list,
- dtype=torch.int32,
- device=protein["aatype"].device,
- )
- restype_atom14_mask = torch.tensor(
- restype_atom14_mask_list,
- dtype=torch.float32,
- device=protein["aatype"].device,
- )
- protein_aatype = protein["aatype"].to(torch.long)
-
- # create the mapping for (residx, atom14) --> atom37, i.e. an array
- # with shape (num_res, 14) containing the atom37 indices for this protein
- residx_atom14_to_atom37 = restype_atom14_to_atom37[protein_aatype]
- residx_atom14_mask = restype_atom14_mask[protein_aatype]
-
- protein["atom14_atom_exists"] = residx_atom14_mask
- protein["residx_atom14_to_atom37"] = residx_atom14_to_atom37.long()
-
- # create the gather indices for mapping back
- residx_atom37_to_atom14 = restype_atom37_to_atom14[protein_aatype]
- protein["residx_atom37_to_atom14"] = residx_atom37_to_atom14.long()
-
- # create the corresponding mask
- restype_atom37_mask = torch.zeros([21, 37], dtype=torch.float32, device=protein["aatype"].device)
- for restype, restype_letter in enumerate(rc.restypes):
- restype_name = rc.restype_1to3[restype_letter]
- atom_names = rc.residue_atoms[restype_name]
- for atom_name in atom_names:
- atom_type = rc.atom_order[atom_name]
- restype_atom37_mask[restype, atom_type] = 1
-
- residx_atom37_mask = restype_atom37_mask[protein_aatype]
- protein["atom37_atom_exists"] = residx_atom37_mask
-
- return protein
-
-
-def make_atom14_masks_np(batch: Dict[str, torch.Tensor]) -> Dict[str, np.ndarray]:
- batch = tree_map(lambda n: torch.tensor(n, device=batch["aatype"].device), batch, np.ndarray)
- out = tensor_tree_map(lambda t: np.array(t), make_atom14_masks(batch))
- return out
diff --git a/transformers/models/esm/openfold_utils/feats.py b/transformers/models/esm/openfold_utils/feats.py
deleted file mode 100644
index 18b01a1fecaccfaafd93f8a269eff6ede752ccb1..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/feats.py
+++ /dev/null
@@ -1,255 +0,0 @@
-# Copyright 2021 AlQuraishi Laboratory
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Dict, Tuple, overload
-
-import torch
-import torch.types
-from torch import nn
-
-from . import residue_constants as rc
-from .rigid_utils import Rigid, Rotation
-from .tensor_utils import batched_gather
-
-
-@overload
-def pseudo_beta_fn(aatype: torch.Tensor, all_atom_positions: torch.Tensor, all_atom_masks: None) -> torch.Tensor:
- ...
-
-
-@overload
-def pseudo_beta_fn(
- aatype: torch.Tensor, all_atom_positions: torch.Tensor, all_atom_masks: torch.Tensor
-) -> Tuple[torch.Tensor, torch.Tensor]:
- ...
-
-
-def pseudo_beta_fn(aatype, all_atom_positions, all_atom_masks):
- is_gly = aatype == rc.restype_order["G"]
- ca_idx = rc.atom_order["CA"]
- cb_idx = rc.atom_order["CB"]
- pseudo_beta = torch.where(
- is_gly[..., None].expand(*((-1,) * len(is_gly.shape)), 3),
- all_atom_positions[..., ca_idx, :],
- all_atom_positions[..., cb_idx, :],
- )
-
- if all_atom_masks is not None:
- pseudo_beta_mask = torch.where(
- is_gly,
- all_atom_masks[..., ca_idx],
- all_atom_masks[..., cb_idx],
- )
- return pseudo_beta, pseudo_beta_mask
- else:
- return pseudo_beta
-
-
-def atom14_to_atom37(atom14: torch.Tensor, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
- atom37_data = batched_gather(
- atom14,
- batch["residx_atom37_to_atom14"],
- dim=-2,
- no_batch_dims=len(atom14.shape[:-2]),
- )
-
- atom37_data = atom37_data * batch["atom37_atom_exists"][..., None]
-
- return atom37_data
-
-
-def build_template_angle_feat(template_feats: Dict[str, torch.Tensor]) -> torch.Tensor:
- template_aatype = template_feats["template_aatype"]
- torsion_angles_sin_cos = template_feats["template_torsion_angles_sin_cos"]
- alt_torsion_angles_sin_cos = template_feats["template_alt_torsion_angles_sin_cos"]
- torsion_angles_mask = template_feats["template_torsion_angles_mask"]
- template_angle_feat = torch.cat(
- [
- nn.functional.one_hot(template_aatype, 22),
- torsion_angles_sin_cos.reshape(*torsion_angles_sin_cos.shape[:-2], 14),
- alt_torsion_angles_sin_cos.reshape(*alt_torsion_angles_sin_cos.shape[:-2], 14),
- torsion_angles_mask,
- ],
- dim=-1,
- )
-
- return template_angle_feat
-
-
-def build_template_pair_feat(
- batch: Dict[str, torch.Tensor],
- min_bin: torch.types.Number,
- max_bin: torch.types.Number,
- no_bins: int,
- use_unit_vector: bool = False,
- eps: float = 1e-20,
- inf: float = 1e8,
-) -> torch.Tensor:
- template_mask = batch["template_pseudo_beta_mask"]
- template_mask_2d = template_mask[..., None] * template_mask[..., None, :]
-
- # Compute distogram (this seems to differ slightly from Alg. 5)
- tpb = batch["template_pseudo_beta"]
- dgram = torch.sum((tpb[..., None, :] - tpb[..., None, :, :]) ** 2, dim=-1, keepdim=True)
- lower = torch.linspace(min_bin, max_bin, no_bins, device=tpb.device) ** 2
- upper = torch.cat([lower[1:], lower.new_tensor([inf])], dim=-1)
- dgram = ((dgram > lower) * (dgram < upper)).type(dgram.dtype)
-
- to_concat = [dgram, template_mask_2d[..., None]]
-
- aatype_one_hot: torch.LongTensor = nn.functional.one_hot(
- batch["template_aatype"],
- rc.restype_num + 2,
- )
-
- n_res = batch["template_aatype"].shape[-1]
- to_concat.append(aatype_one_hot[..., None, :, :].expand(*aatype_one_hot.shape[:-2], n_res, -1, -1))
- to_concat.append(aatype_one_hot[..., None, :].expand(*aatype_one_hot.shape[:-2], -1, n_res, -1))
-
- n, ca, c = [rc.atom_order[a] for a in ["N", "CA", "C"]]
- rigids = Rigid.make_transform_from_reference(
- n_xyz=batch["template_all_atom_positions"][..., n, :],
- ca_xyz=batch["template_all_atom_positions"][..., ca, :],
- c_xyz=batch["template_all_atom_positions"][..., c, :],
- eps=eps,
- )
- points = rigids.get_trans()[..., None, :, :]
- rigid_vec = rigids[..., None].invert_apply(points)
-
- inv_distance_scalar = torch.rsqrt(eps + torch.sum(rigid_vec**2, dim=-1))
-
- t_aa_masks = batch["template_all_atom_mask"]
- template_mask = t_aa_masks[..., n] * t_aa_masks[..., ca] * t_aa_masks[..., c]
- template_mask_2d = template_mask[..., None] * template_mask[..., None, :]
-
- inv_distance_scalar = inv_distance_scalar * template_mask_2d
- unit_vector = rigid_vec * inv_distance_scalar[..., None]
-
- if not use_unit_vector:
- unit_vector = unit_vector * 0.0
-
- to_concat.extend(torch.unbind(unit_vector[..., None, :], dim=-1))
- to_concat.append(template_mask_2d[..., None])
-
- act = torch.cat(to_concat, dim=-1)
- act = act * template_mask_2d[..., None]
-
- return act
-
-
-def build_extra_msa_feat(batch: Dict[str, torch.Tensor]) -> torch.Tensor:
- msa_1hot: torch.LongTensor = nn.functional.one_hot(batch["extra_msa"], 23)
- msa_feat = [
- msa_1hot,
- batch["extra_has_deletion"].unsqueeze(-1),
- batch["extra_deletion_value"].unsqueeze(-1),
- ]
- return torch.cat(msa_feat, dim=-1)
-
-
-def torsion_angles_to_frames(
- r: Rigid,
- alpha: torch.Tensor,
- aatype: torch.Tensor,
- rrgdf: torch.Tensor,
-) -> Rigid:
- # [*, N, 8, 4, 4]
- default_4x4 = rrgdf[aatype, ...]
-
- # [*, N, 8] transformations, i.e.
- # One [*, N, 8, 3, 3] rotation matrix and
- # One [*, N, 8, 3] translation matrix
- default_r = r.from_tensor_4x4(default_4x4)
-
- bb_rot = alpha.new_zeros((*((1,) * len(alpha.shape[:-1])), 2))
- bb_rot[..., 1] = 1
-
- # [*, N, 8, 2]
- alpha = torch.cat([bb_rot.expand(*alpha.shape[:-2], -1, -1), alpha], dim=-2)
-
- # [*, N, 8, 3, 3]
- # Produces rotation matrices of the form:
- # [
- # [1, 0 , 0 ],
- # [0, a_2,-a_1],
- # [0, a_1, a_2]
- # ]
- # This follows the original code rather than the supplement, which uses
- # different indices.
-
- all_rots = alpha.new_zeros(default_r.get_rots().get_rot_mats().shape)
- all_rots[..., 0, 0] = 1
- all_rots[..., 1, 1] = alpha[..., 1]
- all_rots[..., 1, 2] = -alpha[..., 0]
- all_rots[..., 2, 1:] = alpha
-
- all_frames = default_r.compose(Rigid(Rotation(rot_mats=all_rots), None))
-
- chi2_frame_to_frame = all_frames[..., 5]
- chi3_frame_to_frame = all_frames[..., 6]
- chi4_frame_to_frame = all_frames[..., 7]
-
- chi1_frame_to_bb = all_frames[..., 4]
- chi2_frame_to_bb = chi1_frame_to_bb.compose(chi2_frame_to_frame)
- chi3_frame_to_bb = chi2_frame_to_bb.compose(chi3_frame_to_frame)
- chi4_frame_to_bb = chi3_frame_to_bb.compose(chi4_frame_to_frame)
-
- all_frames_to_bb = Rigid.cat(
- [
- all_frames[..., :5],
- chi2_frame_to_bb.unsqueeze(-1),
- chi3_frame_to_bb.unsqueeze(-1),
- chi4_frame_to_bb.unsqueeze(-1),
- ],
- dim=-1,
- )
-
- all_frames_to_global = r[..., None].compose(all_frames_to_bb)
-
- return all_frames_to_global
-
-
-def frames_and_literature_positions_to_atom14_pos(
- r: Rigid,
- aatype: torch.Tensor,
- default_frames: torch.Tensor,
- group_idx: torch.Tensor,
- atom_mask: torch.Tensor,
- lit_positions: torch.Tensor,
-) -> torch.Tensor:
- # [*, N, 14]
- group_mask = group_idx[aatype, ...]
-
- # [*, N, 14, 8]
- group_mask_one_hot: torch.LongTensor = nn.functional.one_hot(
- group_mask,
- num_classes=default_frames.shape[-3],
- )
-
- # [*, N, 14, 8]
- t_atoms_to_global = r[..., None, :] * group_mask_one_hot
-
- # [*, N, 14]
- t_atoms_to_global = t_atoms_to_global.map_tensor_fn(lambda x: torch.sum(x, dim=-1))
-
- # [*, N, 14, 1]
- atom_mask = atom_mask[aatype, ...].unsqueeze(-1)
-
- # [*, N, 14, 3]
- lit_positions = lit_positions[aatype, ...]
- pred_positions = t_atoms_to_global.apply(lit_positions)
- pred_positions = pred_positions * atom_mask
-
- return pred_positions
diff --git a/transformers/models/esm/openfold_utils/loss.py b/transformers/models/esm/openfold_utils/loss.py
deleted file mode 100644
index 8c442786dc82ba2ebe243923509ed76a40de2a01..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/loss.py
+++ /dev/null
@@ -1,105 +0,0 @@
-# Copyright 2021 AlQuraishi Laboratory
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Dict, Optional, Tuple
-
-import torch
-
-
-def _calculate_bin_centers(boundaries: torch.Tensor) -> torch.Tensor:
- step = boundaries[1] - boundaries[0]
- bin_centers = boundaries + step / 2
- bin_centers = torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1)], dim=0)
- return bin_centers
-
-
-def _calculate_expected_aligned_error(
- alignment_confidence_breaks: torch.Tensor,
- aligned_distance_error_probs: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
- bin_centers = _calculate_bin_centers(alignment_confidence_breaks)
- return (
- torch.sum(aligned_distance_error_probs * bin_centers, dim=-1),
- bin_centers[-1],
- )
-
-
-def compute_predicted_aligned_error(
- logits: torch.Tensor,
- max_bin: int = 31,
- no_bins: int = 64,
- **kwargs,
-) -> Dict[str, torch.Tensor]:
- """Computes aligned confidence metrics from logits.
-
- Args:
- logits: [*, num_res, num_res, num_bins] the logits output from
- PredictedAlignedErrorHead.
- max_bin: Maximum bin value
- no_bins: Number of bins
- Returns:
- aligned_confidence_probs: [*, num_res, num_res, num_bins] the predicted
- aligned error probabilities over bins for each residue pair.
- predicted_aligned_error: [*, num_res, num_res] the expected aligned distance
- error for each pair of residues.
- max_predicted_aligned_error: [*] the maximum predicted error possible.
- """
- boundaries = torch.linspace(0, max_bin, steps=(no_bins - 1), device=logits.device)
-
- aligned_confidence_probs = torch.nn.functional.softmax(logits, dim=-1)
- predicted_aligned_error, max_predicted_aligned_error = _calculate_expected_aligned_error(
- alignment_confidence_breaks=boundaries,
- aligned_distance_error_probs=aligned_confidence_probs,
- )
-
- return {
- "aligned_confidence_probs": aligned_confidence_probs,
- "predicted_aligned_error": predicted_aligned_error,
- "max_predicted_aligned_error": max_predicted_aligned_error,
- }
-
-
-def compute_tm(
- logits: torch.Tensor,
- residue_weights: Optional[torch.Tensor] = None,
- max_bin: int = 31,
- no_bins: int = 64,
- eps: float = 1e-8,
- **kwargs,
-) -> torch.Tensor:
- if residue_weights is None:
- residue_weights = logits.new_ones(logits.shape[-2])
-
- boundaries = torch.linspace(0, max_bin, steps=(no_bins - 1), device=logits.device)
-
- bin_centers = _calculate_bin_centers(boundaries)
- torch.sum(residue_weights)
- n = logits.shape[-2]
- clipped_n = max(n, 19)
-
- d0 = 1.24 * (clipped_n - 15) ** (1.0 / 3) - 1.8
-
- probs = torch.nn.functional.softmax(logits, dim=-1)
-
- tm_per_bin = 1.0 / (1 + (bin_centers**2) / (d0**2))
- predicted_tm_term = torch.sum(probs * tm_per_bin, dim=-1)
-
- normed_residue_mask = residue_weights / (eps + residue_weights.sum())
- per_alignment = torch.sum(predicted_tm_term * normed_residue_mask, dim=-1)
-
- weighted = per_alignment * residue_weights
-
- argmax = (weighted == torch.max(weighted)).nonzero()[0]
- return per_alignment[tuple(argmax)]
diff --git a/transformers/models/esm/openfold_utils/protein.py b/transformers/models/esm/openfold_utils/protein.py
deleted file mode 100644
index 32e01571715c1b0c806e9cb764b2dec8aaab6068..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/protein.py
+++ /dev/null
@@ -1,329 +0,0 @@
-# Copyright 2021 AlQuraishi Laboratory
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Protein data type."""
-import dataclasses
-import re
-import string
-from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple
-
-import numpy as np
-
-from . import residue_constants
-
-
-FeatureDict = Mapping[str, np.ndarray]
-ModelOutput = Mapping[str, Any] # Is a nested dict.
-PICO_TO_ANGSTROM = 0.01
-
-
-@dataclasses.dataclass(frozen=True)
-class Protein:
- """Protein structure representation."""
-
- # Cartesian coordinates of atoms in angstroms. The atom types correspond to
- # residue_constants.atom_types, i.e. the first three are N, CA, CB.
- atom_positions: np.ndarray # [num_res, num_atom_type, 3]
-
- # Amino-acid type for each residue represented as an integer between 0 and
- # 20, where 20 is 'X'.
- aatype: np.ndarray # [num_res]
-
- # Binary float mask to indicate presence of a particular atom. 1.0 if an atom
- # is present and 0.0 if not. This should be used for loss masking.
- atom_mask: np.ndarray # [num_res, num_atom_type]
-
- # Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
- residue_index: np.ndarray # [num_res]
-
- # B-factors, or temperature factors, of each residue (in sq. angstroms units),
- # representing the displacement of the residue from its ground truth mean
- # value.
- b_factors: np.ndarray # [num_res, num_atom_type]
-
- # Chain indices for multi-chain predictions
- chain_index: Optional[np.ndarray] = None
-
- # Optional remark about the protein. Included as a comment in output PDB
- # files
- remark: Optional[str] = None
-
- # Templates used to generate this protein (prediction-only)
- parents: Optional[Sequence[str]] = None
-
- # Chain corresponding to each parent
- parents_chain_index: Optional[Sequence[int]] = None
-
-
-def from_proteinnet_string(proteinnet_str: str) -> Protein:
- tag_re = r"(\[[A-Z]+\]\n)"
- tags: List[str] = [tag.strip() for tag in re.split(tag_re, proteinnet_str) if len(tag) > 0]
- groups: Iterator[Tuple[str, List[str]]] = zip(tags[0::2], [l.split("\n") for l in tags[1::2]])
-
- atoms: List[str] = ["N", "CA", "C"]
- aatype = None
- atom_positions = None
- atom_mask = None
- for g in groups:
- if "[PRIMARY]" == g[0]:
- seq = g[1][0].strip()
- for i in range(len(seq)):
- if seq[i] not in residue_constants.restypes:
- seq[i] = "X" # FIXME: strings are immutable
- aatype = np.array(
- [residue_constants.restype_order.get(res_symbol, residue_constants.restype_num) for res_symbol in seq]
- )
- elif "[TERTIARY]" == g[0]:
- tertiary: List[List[float]] = []
- for axis in range(3):
- tertiary.append(list(map(float, g[1][axis].split())))
- tertiary_np = np.array(tertiary)
- atom_positions = np.zeros((len(tertiary[0]) // 3, residue_constants.atom_type_num, 3)).astype(np.float32)
- for i, atom in enumerate(atoms):
- atom_positions[:, residue_constants.atom_order[atom], :] = np.transpose(tertiary_np[:, i::3])
- atom_positions *= PICO_TO_ANGSTROM
- elif "[MASK]" == g[0]:
- mask = np.array(list(map({"-": 0, "+": 1}.get, g[1][0].strip())))
- atom_mask = np.zeros(
- (
- len(mask),
- residue_constants.atom_type_num,
- )
- ).astype(np.float32)
- for i, atom in enumerate(atoms):
- atom_mask[:, residue_constants.atom_order[atom]] = 1
- atom_mask *= mask[..., None]
-
- assert aatype is not None
-
- return Protein(
- atom_positions=atom_positions,
- atom_mask=atom_mask,
- aatype=aatype,
- residue_index=np.arange(len(aatype)),
- b_factors=None,
- )
-
-
-def get_pdb_headers(prot: Protein, chain_id: int = 0) -> List[str]:
- pdb_headers: List[str] = []
-
- remark = prot.remark
- if remark is not None:
- pdb_headers.append(f"REMARK {remark}")
-
- parents = prot.parents
- parents_chain_index = prot.parents_chain_index
- if parents is not None and parents_chain_index is not None:
- parents = [p for i, p in zip(parents_chain_index, parents) if i == chain_id]
-
- if parents is None or len(parents) == 0:
- parents = ["N/A"]
-
- pdb_headers.append(f"PARENT {' '.join(parents)}")
-
- return pdb_headers
-
-
-def add_pdb_headers(prot: Protein, pdb_str: str) -> str:
- """Add pdb headers to an existing PDB string. Useful during multi-chain
- recycling
- """
- out_pdb_lines: List[str] = []
- lines = pdb_str.split("\n")
-
- remark = prot.remark
- if remark is not None:
- out_pdb_lines.append(f"REMARK {remark}")
-
- parents_per_chain: List[List[str]]
- if prot.parents is not None and len(prot.parents) > 0:
- parents_per_chain = []
- if prot.parents_chain_index is not None:
- parent_dict: Dict[str, List[str]] = {}
- for p, i in zip(prot.parents, prot.parents_chain_index):
- parent_dict.setdefault(str(i), [])
- parent_dict[str(i)].append(p)
-
- max_idx = max([int(chain_idx) for chain_idx in parent_dict])
- for i in range(max_idx + 1):
- chain_parents = parent_dict.get(str(i), ["N/A"])
- parents_per_chain.append(chain_parents)
- else:
- parents_per_chain.append(list(prot.parents))
- else:
- parents_per_chain = [["N/A"]]
-
- def make_parent_line(p: Sequence[str]) -> str:
- return f"PARENT {' '.join(p)}"
-
- out_pdb_lines.append(make_parent_line(parents_per_chain[0]))
-
- chain_counter = 0
- for i, l in enumerate(lines):
- if "PARENT" not in l and "REMARK" not in l:
- out_pdb_lines.append(l)
- if "TER" in l and "END" not in lines[i + 1]:
- chain_counter += 1
- if not chain_counter >= len(parents_per_chain):
- chain_parents = parents_per_chain[chain_counter]
- else:
- chain_parents = ["N/A"]
-
- out_pdb_lines.append(make_parent_line(chain_parents))
-
- return "\n".join(out_pdb_lines)
-
-
-def to_pdb(prot: Protein) -> str:
- """Converts a `Protein` instance to a PDB string.
-
- Args:
- prot: The protein to convert to PDB.
-
- Returns:
- PDB string.
- """
- restypes = residue_constants.restypes + ["X"]
-
- def res_1to3(r: int) -> str:
- return residue_constants.restype_1to3.get(restypes[r], "UNK")
-
- atom_types = residue_constants.atom_types
-
- pdb_lines: List[str] = []
-
- atom_mask = prot.atom_mask
- aatype = prot.aatype
- atom_positions = prot.atom_positions
- residue_index = prot.residue_index.astype(np.int32)
- b_factors = prot.b_factors
- chain_index = prot.chain_index
-
- if np.any(aatype > residue_constants.restype_num):
- raise ValueError("Invalid aatypes.")
-
- headers = get_pdb_headers(prot)
- if len(headers) > 0:
- pdb_lines.extend(headers)
-
- n = aatype.shape[0]
- atom_index = 1
- prev_chain_index = 0
- chain_tags = string.ascii_uppercase
- chain_tag = None
- # Add all atom sites.
- for i in range(n):
- res_name_3 = res_1to3(aatype[i])
- for atom_name, pos, mask, b_factor in zip(atom_types, atom_positions[i], atom_mask[i], b_factors[i]):
- if mask < 0.5:
- continue
-
- record_type = "ATOM"
- name = atom_name if len(atom_name) == 4 else f" {atom_name}"
- alt_loc = ""
- insertion_code = ""
- occupancy = 1.00
- element = atom_name[0] # Protein supports only C, N, O, S, this works.
- charge = ""
-
- chain_tag = "A"
- if chain_index is not None:
- chain_tag = chain_tags[chain_index[i]]
-
- # PDB is a columnar format, every space matters here!
- atom_line = (
- f"{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"
- f"{res_name_3:>3} {chain_tag:>1}"
- f"{residue_index[i]:>4}{insertion_code:>1} "
- f"{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"
- f"{occupancy:>6.2f}{b_factor:>6.2f} "
- f"{element:>2}{charge:>2}"
- )
- pdb_lines.append(atom_line)
- atom_index += 1
-
- should_terminate = i == n - 1
- if chain_index is not None:
- if i != n - 1 and chain_index[i + 1] != prev_chain_index:
- should_terminate = True
- prev_chain_index = chain_index[i + 1]
-
- if should_terminate:
- # Close the chain.
- chain_end = "TER"
- chain_termination_line = (
- f"{chain_end:<6}{atom_index:>5} {res_1to3(aatype[i]):>3} {chain_tag:>1}{residue_index[i]:>4}"
- )
- pdb_lines.append(chain_termination_line)
- atom_index += 1
-
- if i != n - 1:
- # "prev" is a misnomer here. This happens at the beginning of
- # each new chain.
- pdb_lines.extend(get_pdb_headers(prot, prev_chain_index))
-
- pdb_lines.append("END")
- pdb_lines.append("")
- return "\n".join(pdb_lines)
-
-
-def ideal_atom_mask(prot: Protein) -> np.ndarray:
- """Computes an ideal atom mask.
-
- `Protein.atom_mask` typically is defined according to the atoms that are reported in the PDB. This function
- computes a mask according to heavy atoms that should be present in the given sequence of amino acids.
-
- Args:
- prot: `Protein` whose fields are `numpy.ndarray` objects.
-
- Returns:
- An ideal atom mask.
- """
- return residue_constants.STANDARD_ATOM_MASK[prot.aatype]
-
-
-def from_prediction(
- features: FeatureDict,
- result: ModelOutput,
- b_factors: Optional[np.ndarray] = None,
- chain_index: Optional[np.ndarray] = None,
- remark: Optional[str] = None,
- parents: Optional[Sequence[str]] = None,
- parents_chain_index: Optional[Sequence[int]] = None,
-) -> Protein:
- """Assembles a protein from a prediction.
-
- Args:
- features: Dictionary holding model inputs.
- result: Dictionary holding model outputs.
- b_factors: (Optional) B-factors to use for the protein.
- chain_index: (Optional) Chain indices for multi-chain predictions
- remark: (Optional) Remark about the prediction
- parents: (Optional) List of template names
- Returns:
- A protein instance.
- """
- return Protein(
- aatype=features["aatype"],
- atom_positions=result["final_atom_positions"],
- atom_mask=result["final_atom_mask"],
- residue_index=features["residue_index"] + 1,
- b_factors=b_factors if b_factors is not None else np.zeros_like(result["final_atom_mask"]),
- chain_index=chain_index,
- remark=remark,
- parents=parents,
- parents_chain_index=parents_chain_index,
- )
diff --git a/transformers/models/esm/openfold_utils/residue_constants.py b/transformers/models/esm/openfold_utils/residue_constants.py
deleted file mode 100644
index 8f0ad3b50c65050a4ffd4370e9b4f3a3312fc723..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/residue_constants.py
+++ /dev/null
@@ -1,983 +0,0 @@
-# Copyright 2021 AlQuraishi Laboratory
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Constants used in AlphaFold."""
-
-import collections
-import copy
-import functools
-from importlib import resources
-from typing import Dict, List, Mapping, Sequence, Tuple
-
-import numpy as np
-
-
-# Internal import (35fd).
-
-
-# Distance from one CA to next CA [trans configuration: omega = 180].
-ca_ca = 3.80209737096
-
-# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
-# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
-# chi angles so their chi angle lists are empty.
-chi_angles_atoms: Dict[str, List[List[str]]] = {
- "ALA": [],
- # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
- "ARG": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"], ["CB", "CG", "CD", "NE"], ["CG", "CD", "NE", "CZ"]],
- "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
- "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
- "CYS": [["N", "CA", "CB", "SG"]],
- "GLN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"], ["CB", "CG", "CD", "OE1"]],
- "GLU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"], ["CB", "CG", "CD", "OE1"]],
- "GLY": [],
- "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]],
- "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]],
- "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
- "LYS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"], ["CB", "CG", "CD", "CE"], ["CG", "CD", "CE", "NZ"]],
- "MET": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "SD"], ["CB", "CG", "SD", "CE"]],
- "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
- "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]],
- "SER": [["N", "CA", "CB", "OG"]],
- "THR": [["N", "CA", "CB", "OG1"]],
- "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
- "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
- "VAL": [["N", "CA", "CB", "CG1"]],
-}
-
-# If chi angles given in fixed-length array, this matrix determines how to mask
-# them for each AA type. The order is as per restype_order (see below).
-chi_angles_mask: List[List[float]] = [
- [0.0, 0.0, 0.0, 0.0], # ALA
- [1.0, 1.0, 1.0, 1.0], # ARG
- [1.0, 1.0, 0.0, 0.0], # ASN
- [1.0, 1.0, 0.0, 0.0], # ASP
- [1.0, 0.0, 0.0, 0.0], # CYS
- [1.0, 1.0, 1.0, 0.0], # GLN
- [1.0, 1.0, 1.0, 0.0], # GLU
- [0.0, 0.0, 0.0, 0.0], # GLY
- [1.0, 1.0, 0.0, 0.0], # HIS
- [1.0, 1.0, 0.0, 0.0], # ILE
- [1.0, 1.0, 0.0, 0.0], # LEU
- [1.0, 1.0, 1.0, 1.0], # LYS
- [1.0, 1.0, 1.0, 0.0], # MET
- [1.0, 1.0, 0.0, 0.0], # PHE
- [1.0, 1.0, 0.0, 0.0], # PRO
- [1.0, 0.0, 0.0, 0.0], # SER
- [1.0, 0.0, 0.0, 0.0], # THR
- [1.0, 1.0, 0.0, 0.0], # TRP
- [1.0, 1.0, 0.0, 0.0], # TYR
- [1.0, 0.0, 0.0, 0.0], # VAL
-]
-
-# The following chi angles are pi periodic: they can be rotated by a multiple
-# of pi without affecting the structure.
-chi_pi_periodic: List[List[float]] = [
- [0.0, 0.0, 0.0, 0.0], # ALA
- [0.0, 0.0, 0.0, 0.0], # ARG
- [0.0, 0.0, 0.0, 0.0], # ASN
- [0.0, 1.0, 0.0, 0.0], # ASP
- [0.0, 0.0, 0.0, 0.0], # CYS
- [0.0, 0.0, 0.0, 0.0], # GLN
- [0.0, 0.0, 1.0, 0.0], # GLU
- [0.0, 0.0, 0.0, 0.0], # GLY
- [0.0, 0.0, 0.0, 0.0], # HIS
- [0.0, 0.0, 0.0, 0.0], # ILE
- [0.0, 0.0, 0.0, 0.0], # LEU
- [0.0, 0.0, 0.0, 0.0], # LYS
- [0.0, 0.0, 0.0, 0.0], # MET
- [0.0, 1.0, 0.0, 0.0], # PHE
- [0.0, 0.0, 0.0, 0.0], # PRO
- [0.0, 0.0, 0.0, 0.0], # SER
- [0.0, 0.0, 0.0, 0.0], # THR
- [0.0, 0.0, 0.0, 0.0], # TRP
- [0.0, 1.0, 0.0, 0.0], # TYR
- [0.0, 0.0, 0.0, 0.0], # VAL
- [0.0, 0.0, 0.0, 0.0], # UNK
-]
-
-# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
-# psi and chi angles:
-# 0: 'backbone group',
-# 1: 'pre-omega-group', (empty)
-# 2: 'phi-group', (currently empty, because it defines only hydrogens)
-# 3: 'psi-group',
-# 4,5,6,7: 'chi1,2,3,4-group'
-# The atom positions are relative to the axis-end-atom of the corresponding
-# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
-# is defined such that the dihedral-angle-definiting atom (the last entry in
-# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
-# format: [atomname, group_idx, rel_position]
-rigid_group_atom_positions: Dict[str, List[Tuple[str, int, Tuple[float, float, float]]]] = {
- "ALA": [
- ("N", 0, (-0.525, 1.363, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.526, -0.000, -0.000)),
- ("CB", 0, (-0.529, -0.774, -1.205)),
- ("O", 3, (0.627, 1.062, 0.000)),
- ],
- "ARG": [
- ("N", 0, (-0.524, 1.362, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.525, -0.000, -0.000)),
- ("CB", 0, (-0.524, -0.778, -1.209)),
- ("O", 3, (0.626, 1.062, 0.000)),
- ("CG", 4, (0.616, 1.390, -0.000)),
- ("CD", 5, (0.564, 1.414, 0.000)),
- ("NE", 6, (0.539, 1.357, -0.000)),
- ("NH1", 7, (0.206, 2.301, 0.000)),
- ("NH2", 7, (2.078, 0.978, -0.000)),
- ("CZ", 7, (0.758, 1.093, -0.000)),
- ],
- "ASN": [
- ("N", 0, (-0.536, 1.357, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.526, -0.000, -0.000)),
- ("CB", 0, (-0.531, -0.787, -1.200)),
- ("O", 3, (0.625, 1.062, 0.000)),
- ("CG", 4, (0.584, 1.399, 0.000)),
- ("ND2", 5, (0.593, -1.188, 0.001)),
- ("OD1", 5, (0.633, 1.059, 0.000)),
- ],
- "ASP": [
- ("N", 0, (-0.525, 1.362, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.527, 0.000, -0.000)),
- ("CB", 0, (-0.526, -0.778, -1.208)),
- ("O", 3, (0.626, 1.062, -0.000)),
- ("CG", 4, (0.593, 1.398, -0.000)),
- ("OD1", 5, (0.610, 1.091, 0.000)),
- ("OD2", 5, (0.592, -1.101, -0.003)),
- ],
- "CYS": [
- ("N", 0, (-0.522, 1.362, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.524, 0.000, 0.000)),
- ("CB", 0, (-0.519, -0.773, -1.212)),
- ("O", 3, (0.625, 1.062, -0.000)),
- ("SG", 4, (0.728, 1.653, 0.000)),
- ],
- "GLN": [
- ("N", 0, (-0.526, 1.361, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.526, 0.000, 0.000)),
- ("CB", 0, (-0.525, -0.779, -1.207)),
- ("O", 3, (0.626, 1.062, -0.000)),
- ("CG", 4, (0.615, 1.393, 0.000)),
- ("CD", 5, (0.587, 1.399, -0.000)),
- ("NE2", 6, (0.593, -1.189, -0.001)),
- ("OE1", 6, (0.634, 1.060, 0.000)),
- ],
- "GLU": [
- ("N", 0, (-0.528, 1.361, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.526, -0.000, -0.000)),
- ("CB", 0, (-0.526, -0.781, -1.207)),
- ("O", 3, (0.626, 1.062, 0.000)),
- ("CG", 4, (0.615, 1.392, 0.000)),
- ("CD", 5, (0.600, 1.397, 0.000)),
- ("OE1", 6, (0.607, 1.095, -0.000)),
- ("OE2", 6, (0.589, -1.104, -0.001)),
- ],
- "GLY": [
- ("N", 0, (-0.572, 1.337, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.517, -0.000, -0.000)),
- ("O", 3, (0.626, 1.062, -0.000)),
- ],
- "HIS": [
- ("N", 0, (-0.527, 1.360, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.525, 0.000, 0.000)),
- ("CB", 0, (-0.525, -0.778, -1.208)),
- ("O", 3, (0.625, 1.063, 0.000)),
- ("CG", 4, (0.600, 1.370, -0.000)),
- ("CD2", 5, (0.889, -1.021, 0.003)),
- ("ND1", 5, (0.744, 1.160, -0.000)),
- ("CE1", 5, (2.030, 0.851, 0.002)),
- ("NE2", 5, (2.145, -0.466, 0.004)),
- ],
- "ILE": [
- ("N", 0, (-0.493, 1.373, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.527, -0.000, -0.000)),
- ("CB", 0, (-0.536, -0.793, -1.213)),
- ("O", 3, (0.627, 1.062, -0.000)),
- ("CG1", 4, (0.534, 1.437, -0.000)),
- ("CG2", 4, (0.540, -0.785, -1.199)),
- ("CD1", 5, (0.619, 1.391, 0.000)),
- ],
- "LEU": [
- ("N", 0, (-0.520, 1.363, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.525, -0.000, -0.000)),
- ("CB", 0, (-0.522, -0.773, -1.214)),
- ("O", 3, (0.625, 1.063, -0.000)),
- ("CG", 4, (0.678, 1.371, 0.000)),
- ("CD1", 5, (0.530, 1.430, -0.000)),
- ("CD2", 5, (0.535, -0.774, 1.200)),
- ],
- "LYS": [
- ("N", 0, (-0.526, 1.362, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.526, 0.000, 0.000)),
- ("CB", 0, (-0.524, -0.778, -1.208)),
- ("O", 3, (0.626, 1.062, -0.000)),
- ("CG", 4, (0.619, 1.390, 0.000)),
- ("CD", 5, (0.559, 1.417, 0.000)),
- ("CE", 6, (0.560, 1.416, 0.000)),
- ("NZ", 7, (0.554, 1.387, 0.000)),
- ],
- "MET": [
- ("N", 0, (-0.521, 1.364, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.525, 0.000, 0.000)),
- ("CB", 0, (-0.523, -0.776, -1.210)),
- ("O", 3, (0.625, 1.062, -0.000)),
- ("CG", 4, (0.613, 1.391, -0.000)),
- ("SD", 5, (0.703, 1.695, 0.000)),
- ("CE", 6, (0.320, 1.786, -0.000)),
- ],
- "PHE": [
- ("N", 0, (-0.518, 1.363, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.524, 0.000, -0.000)),
- ("CB", 0, (-0.525, -0.776, -1.212)),
- ("O", 3, (0.626, 1.062, -0.000)),
- ("CG", 4, (0.607, 1.377, 0.000)),
- ("CD1", 5, (0.709, 1.195, -0.000)),
- ("CD2", 5, (0.706, -1.196, 0.000)),
- ("CE1", 5, (2.102, 1.198, -0.000)),
- ("CE2", 5, (2.098, -1.201, -0.000)),
- ("CZ", 5, (2.794, -0.003, -0.001)),
- ],
- "PRO": [
- ("N", 0, (-0.566, 1.351, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.527, -0.000, 0.000)),
- ("CB", 0, (-0.546, -0.611, -1.293)),
- ("O", 3, (0.621, 1.066, 0.000)),
- ("CG", 4, (0.382, 1.445, 0.0)),
- # ('CD', 5, (0.427, 1.440, 0.0)),
- ("CD", 5, (0.477, 1.424, 0.0)), # manually made angle 2 degrees larger
- ],
- "SER": [
- ("N", 0, (-0.529, 1.360, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.525, -0.000, -0.000)),
- ("CB", 0, (-0.518, -0.777, -1.211)),
- ("O", 3, (0.626, 1.062, -0.000)),
- ("OG", 4, (0.503, 1.325, 0.000)),
- ],
- "THR": [
- ("N", 0, (-0.517, 1.364, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.526, 0.000, -0.000)),
- ("CB", 0, (-0.516, -0.793, -1.215)),
- ("O", 3, (0.626, 1.062, 0.000)),
- ("CG2", 4, (0.550, -0.718, -1.228)),
- ("OG1", 4, (0.472, 1.353, 0.000)),
- ],
- "TRP": [
- ("N", 0, (-0.521, 1.363, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.525, -0.000, 0.000)),
- ("CB", 0, (-0.523, -0.776, -1.212)),
- ("O", 3, (0.627, 1.062, 0.000)),
- ("CG", 4, (0.609, 1.370, -0.000)),
- ("CD1", 5, (0.824, 1.091, 0.000)),
- ("CD2", 5, (0.854, -1.148, -0.005)),
- ("CE2", 5, (2.186, -0.678, -0.007)),
- ("CE3", 5, (0.622, -2.530, -0.007)),
- ("NE1", 5, (2.140, 0.690, -0.004)),
- ("CH2", 5, (3.028, -2.890, -0.013)),
- ("CZ2", 5, (3.283, -1.543, -0.011)),
- ("CZ3", 5, (1.715, -3.389, -0.011)),
- ],
- "TYR": [
- ("N", 0, (-0.522, 1.362, 0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.524, -0.000, -0.000)),
- ("CB", 0, (-0.522, -0.776, -1.213)),
- ("O", 3, (0.627, 1.062, -0.000)),
- ("CG", 4, (0.607, 1.382, -0.000)),
- ("CD1", 5, (0.716, 1.195, -0.000)),
- ("CD2", 5, (0.713, -1.194, -0.001)),
- ("CE1", 5, (2.107, 1.200, -0.002)),
- ("CE2", 5, (2.104, -1.201, -0.003)),
- ("OH", 5, (4.168, -0.002, -0.005)),
- ("CZ", 5, (2.791, -0.001, -0.003)),
- ],
- "VAL": [
- ("N", 0, (-0.494, 1.373, -0.000)),
- ("CA", 0, (0.000, 0.000, 0.000)),
- ("C", 0, (1.527, -0.000, -0.000)),
- ("CB", 0, (-0.533, -0.795, -1.213)),
- ("O", 3, (0.627, 1.062, -0.000)),
- ("CG1", 4, (0.540, 1.429, -0.000)),
- ("CG2", 4, (0.533, -0.776, 1.203)),
- ],
-}
-
-# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
-residue_atoms: Dict[str, List[str]] = {
- "ALA": ["C", "CA", "CB", "N", "O"],
- "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"],
- "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"],
- "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"],
- "CYS": ["C", "CA", "CB", "N", "O", "SG"],
- "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"],
- "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"],
- "GLY": ["C", "CA", "N", "O"],
- "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"],
- "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"],
- "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"],
- "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"],
- "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"],
- "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"],
- "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"],
- "SER": ["C", "CA", "CB", "N", "O", "OG"],
- "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"],
- "TRP": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE2", "CE3", "CZ2", "CZ3", "CH2", "N", "NE1", "O"],
- "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"],
- "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"],
-}
-
-# Naming swaps for ambiguous atom names.
-# Due to symmetries in the amino acids the naming of atoms is ambiguous in
-# 4 of the 20 amino acids.
-# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
-# in LEU, VAL and ARG can be resolved by using the 3d constellations of
-# the 'ambiguous' atoms and their neighbours)
-# TODO: ^ interpret this
-residue_atom_renaming_swaps: Dict[str, Dict[str, str]] = {
- "ASP": {"OD1": "OD2"},
- "GLU": {"OE1": "OE2"},
- "PHE": {"CD1": "CD2", "CE1": "CE2"},
- "TYR": {"CD1": "CD2", "CE1": "CE2"},
-}
-
-# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
-van_der_waals_radius: Dict[str, float] = {
- "C": 1.7,
- "N": 1.55,
- "O": 1.52,
- "S": 1.8,
-}
-
-Bond = collections.namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"])
-BondAngle = collections.namedtuple(
- "BondAngle",
- ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"],
-)
-
-
-def map_structure_with_atom_order(in_list: list, first_call: bool = True) -> list:
- # Maps strings in a nested list structure to their corresponding index in atom_order
- if first_call:
- in_list = copy.deepcopy(in_list)
- for i in range(len(in_list)):
- if isinstance(in_list[i], list):
- in_list[i] = map_structure_with_atom_order(in_list[i], first_call=False)
- elif isinstance(in_list[i], str):
- in_list[i] = atom_order[in_list[i]]
- else:
- raise ValueError("Unexpected type when mapping nested lists!")
- return in_list
-
-
-@functools.lru_cache(maxsize=None)
-def load_stereo_chemical_props() -> (
- Tuple[
- Mapping[str, List[Bond]],
- Mapping[str, List[Bond]],
- Mapping[str, List[BondAngle]],
- ]
-):
- """Load stereo_chemical_props.txt into a nice structure.
-
- Load literature values for bond lengths and bond angles and translate bond angles into the length of the opposite
- edge of the triangle ("residue_virtual_bonds").
-
- Returns:
- residue_bonds: dict that maps resname --> list of Bond tuples residue_virtual_bonds: dict that maps resname -->
- list of Bond tuples residue_bond_angles: dict that maps resname --> list of BondAngle tuples
- """
- # TODO: this file should be downloaded in a setup script
- stereo_chemical_props = resources.read_text("openfold.resources", "stereo_chemical_props.txt")
-
- lines_iter = iter(stereo_chemical_props.splitlines())
- # Load bond lengths.
- residue_bonds: Dict[str, List[Bond]] = {}
- next(lines_iter) # Skip header line.
- for line in lines_iter:
- if line.strip() == "-":
- break
- bond, resname, bond_length, stddev = line.split()
- atom1, atom2 = bond.split("-")
- if resname not in residue_bonds:
- residue_bonds[resname] = []
- residue_bonds[resname].append(Bond(atom1, atom2, float(bond_length), float(stddev)))
- residue_bonds["UNK"] = []
-
- # Load bond angles.
- residue_bond_angles: Dict[str, List[BondAngle]] = {}
- next(lines_iter) # Skip empty line.
- next(lines_iter) # Skip header line.
- for line in lines_iter:
- if line.strip() == "-":
- break
- bond, resname, angle_degree, stddev_degree = line.split()
- atom1, atom2, atom3 = bond.split("-")
- if resname not in residue_bond_angles:
- residue_bond_angles[resname] = []
- residue_bond_angles[resname].append(
- BondAngle(
- atom1,
- atom2,
- atom3,
- float(angle_degree) / 180.0 * np.pi,
- float(stddev_degree) / 180.0 * np.pi,
- )
- )
- residue_bond_angles["UNK"] = []
-
- def make_bond_key(atom1_name: str, atom2_name: str) -> str:
- """Unique key to lookup bonds."""
- return "-".join(sorted([atom1_name, atom2_name]))
-
- # Translate bond angles into distances ("virtual bonds").
- residue_virtual_bonds: Dict[str, List[Bond]] = {}
- for resname, bond_angles in residue_bond_angles.items():
- # Create a fast lookup dict for bond lengths.
- bond_cache: Dict[str, Bond] = {}
- for b in residue_bonds[resname]:
- bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
- residue_virtual_bonds[resname] = []
- for ba in bond_angles:
- bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
- bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]
-
- # Compute distance between atom1 and atom3 using the law of cosines
- # c^2 = a^2 + b^2 - 2ab*cos(gamma).
- gamma = ba.angle_rad
- length = np.sqrt(bond1.length**2 + bond2.length**2 - 2 * bond1.length * bond2.length * np.cos(gamma))
-
- # Propagation of uncertainty assuming uncorrelated errors.
- dl_outer = 0.5 / length
- dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
- dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
- dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
- stddev = np.sqrt(
- (dl_dgamma * ba.stddev) ** 2 + (dl_db1 * bond1.stddev) ** 2 + (dl_db2 * bond2.stddev) ** 2
- )
- residue_virtual_bonds[resname].append(Bond(ba.atom1_name, ba.atom3name, length, stddev))
-
- return (residue_bonds, residue_virtual_bonds, residue_bond_angles)
-
-
-# Between-residue bond lengths for general bonds (first element) and for Proline
-# (second element).
-between_res_bond_length_c_n: Tuple[float, float] = (1.329, 1.341)
-between_res_bond_length_stddev_c_n: Tuple[float, float] = (0.014, 0.016)
-
-# Between-residue cos_angles.
-between_res_cos_angles_c_n_ca: Tuple[float, float] = (-0.5203, 0.0353) # degrees: 121.352 +- 2.315
-between_res_cos_angles_ca_c_n: Tuple[float, float] = (-0.4473, 0.0311) # degrees: 116.568 +- 1.995
-
-# This mapping is used when we need to store atom data in a format that requires
-# fixed atom data size for every residue (e.g. a numpy array).
-atom_types: List[str] = [
- "N",
- "CA",
- "C",
- "CB",
- "O",
- "CG",
- "CG1",
- "CG2",
- "OG",
- "OG1",
- "SG",
- "CD",
- "CD1",
- "CD2",
- "ND1",
- "ND2",
- "OD1",
- "OD2",
- "SD",
- "CE",
- "CE1",
- "CE2",
- "CE3",
- "NE",
- "NE1",
- "NE2",
- "OE1",
- "OE2",
- "CH2",
- "NH1",
- "NH2",
- "OH",
- "CZ",
- "CZ2",
- "CZ3",
- "NZ",
- "OXT",
-]
-atom_order: Dict[str, int] = {atom_type: i for i, atom_type in enumerate(atom_types)}
-atom_type_num = len(atom_types) # := 37.
-
-# A compact atom encoding with 14 columns
-# pylint: disable=line-too-long
-# pylint: disable=bad-whitespace
-restype_name_to_atom14_names: Dict[str, List[str]] = {
- "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
- "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2", "", "", ""],
- "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""],
- "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""],
- "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""],
- "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""],
- "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""],
- "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""],
- "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2", "", "", "", ""],
- "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""],
- "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""],
- "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""],
- "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""],
- "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "", "", ""],
- "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""],
- "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""],
- "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""],
- "TRP": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "NE1", "CE2", "CE3", "CZ2", "CZ3", "CH2"],
- "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH", "", ""],
- "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""],
- "UNK": ["", "", "", "", "", "", "", "", "", "", "", "", "", ""],
-}
-# pylint: enable=line-too-long
-# pylint: enable=bad-whitespace
-
-
-# This is the standard residue order when coding AA type as a number.
-# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
-restypes: List[str] = [
- "A",
- "R",
- "N",
- "D",
- "C",
- "Q",
- "E",
- "G",
- "H",
- "I",
- "L",
- "K",
- "M",
- "F",
- "P",
- "S",
- "T",
- "W",
- "Y",
- "V",
-]
-restype_order: Dict[str, int] = {restype: i for i, restype in enumerate(restypes)}
-restype_num = len(restypes) # := 20.
-unk_restype_index = restype_num # Catch-all index for unknown restypes.
-
-restypes_with_x: List[str] = restypes + ["X"]
-restype_order_with_x: Dict[str, int] = {restype: i for i, restype in enumerate(restypes_with_x)}
-
-
-def sequence_to_onehot(sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False) -> np.ndarray:
- """Maps the given sequence into a one-hot encoded matrix.
-
- Args:
- sequence: An amino acid sequence.
- mapping: A dictionary mapping amino acids to integers.
- map_unknown_to_x: If True, any amino acid that is not in the mapping will be
- mapped to the unknown amino acid 'X'. If the mapping doesn't contain amino acid 'X', an error will be thrown.
- If False, any amino acid not in the mapping will throw an error.
-
- Returns:
- A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of the sequence.
-
- Raises:
- ValueError: If the mapping doesn't contain values from 0 to
- num_unique_aas - 1 without any gaps.
- """
- num_entries = max(mapping.values()) + 1
-
- if sorted(set(mapping.values())) != list(range(num_entries)):
- raise ValueError(
- "The mapping must have values from 0 to num_unique_aas-1 without any gaps. Got: %s"
- % sorted(mapping.values())
- )
-
- one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)
-
- for aa_index, aa_type in enumerate(sequence):
- if map_unknown_to_x:
- if aa_type.isalpha() and aa_type.isupper():
- aa_id = mapping.get(aa_type, mapping["X"])
- else:
- raise ValueError(f"Invalid character in the sequence: {aa_type}")
- else:
- aa_id = mapping[aa_type]
- one_hot_arr[aa_index, aa_id] = 1
-
- return one_hot_arr
-
-
-restype_1to3: Dict[str, str] = {
- "A": "ALA",
- "R": "ARG",
- "N": "ASN",
- "D": "ASP",
- "C": "CYS",
- "Q": "GLN",
- "E": "GLU",
- "G": "GLY",
- "H": "HIS",
- "I": "ILE",
- "L": "LEU",
- "K": "LYS",
- "M": "MET",
- "F": "PHE",
- "P": "PRO",
- "S": "SER",
- "T": "THR",
- "W": "TRP",
- "Y": "TYR",
- "V": "VAL",
-}
-
-
-# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
-# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
-# many more, and less common, three letter names as keys and maps many of these
-# to the same one letter name (including 'X' and 'U' which we don't use here).
-restype_3to1: Dict[str, str] = {v: k for k, v in restype_1to3.items()}
-
-# Define a restype name for all unknown residues.
-unk_restype = "UNK"
-
-resnames: List[str] = [restype_1to3[r] for r in restypes] + [unk_restype]
-resname_to_idx: Dict[str, int] = {resname: i for i, resname in enumerate(resnames)}
-
-
-# The mapping here uses hhblits convention, so that B is mapped to D, J and O
-# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
-# remaining 20 amino acids are kept in alphabetical order.
-# There are 2 non-amino acid codes, X (representing any amino acid) and
-# "-" representing a missing amino acid in an alignment. The id for these
-# codes is put at the end (20 and 21) so that they can easily be ignored if
-# desired.
-HHBLITS_AA_TO_ID: Dict[str, int] = {
- "A": 0,
- "B": 2,
- "C": 1,
- "D": 2,
- "E": 3,
- "F": 4,
- "G": 5,
- "H": 6,
- "I": 7,
- "J": 20,
- "K": 8,
- "L": 9,
- "M": 10,
- "N": 11,
- "O": 20,
- "P": 12,
- "Q": 13,
- "R": 14,
- "S": 15,
- "T": 16,
- "U": 1,
- "V": 17,
- "W": 18,
- "X": 20,
- "Y": 19,
- "Z": 3,
- "-": 21,
-}
-
-# Partial inversion of HHBLITS_AA_TO_ID.
-ID_TO_HHBLITS_AA: Dict[int, str] = {
- 0: "A",
- 1: "C", # Also U.
- 2: "D", # Also B.
- 3: "E", # Also Z.
- 4: "F",
- 5: "G",
- 6: "H",
- 7: "I",
- 8: "K",
- 9: "L",
- 10: "M",
- 11: "N",
- 12: "P",
- 13: "Q",
- 14: "R",
- 15: "S",
- 16: "T",
- 17: "V",
- 18: "W",
- 19: "Y",
- 20: "X", # Includes J and O.
- 21: "-",
-}
-
-restypes_with_x_and_gap: List[str] = restypes + ["X", "-"]
-MAP_HHBLITS_AATYPE_TO_OUR_AATYPE: Tuple[int, ...] = tuple(
- restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i]) for i in range(len(restypes_with_x_and_gap))
-)
-
-
-def _make_standard_atom_mask() -> np.ndarray:
- """Returns [num_res_types, num_atom_types] mask array."""
- # +1 to account for unknown (all 0s).
- mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
- for restype, restype_letter in enumerate(restypes):
- restype_name = restype_1to3[restype_letter]
- atom_names = residue_atoms[restype_name]
- for atom_name in atom_names:
- atom_type = atom_order[atom_name]
- mask[restype, atom_type] = 1
- return mask
-
-
-STANDARD_ATOM_MASK = _make_standard_atom_mask()
-
-
-# A one hot representation for the first and second atoms defining the axis
-# of rotation for each chi-angle in each residue.
-def chi_angle_atom(atom_index: int) -> np.ndarray:
- """Define chi-angle rigid groups via one-hot representations."""
- chi_angles_index = {}
- one_hots = []
-
- for k, v in chi_angles_atoms.items():
- indices = [atom_types.index(s[atom_index]) for s in v]
- indices.extend([-1] * (4 - len(indices)))
- chi_angles_index[k] = indices
-
- for r in restypes:
- res3 = restype_1to3[r]
- one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
- one_hots.append(one_hot)
-
- one_hots.append(np.zeros([4, atom_type_num])) # Add zeros for residue `X`.
- one_hot = np.stack(one_hots, axis=0)
- one_hot = np.transpose(one_hot, [0, 2, 1])
-
- return one_hot
-
-
-chi_atom_1_one_hot = chi_angle_atom(1)
-chi_atom_2_one_hot = chi_angle_atom(2)
-
-# An array like chi_angles_atoms but using indices rather than names.
-chi_angles_atom_indices_list: List[List[List[str]]] = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
-chi_angles_atom_indices_ours: list = map_structure_with_atom_order(chi_angles_atom_indices_list)
-chi_angles_atom_indices = np.array(
- [chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms))) for chi_atoms in chi_angles_atom_indices_list]
-)
-
-# Mapping from (res_name, atom_name) pairs to the atom's chi group index
-# and atom index within that group.
-chi_groups_for_atom: Dict[Tuple[str, str], List[Tuple[int, int]]] = collections.defaultdict(list)
-for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
- for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
- for atom_i, atom in enumerate(chi_group):
- chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
-chi_groups_for_atom = dict(chi_groups_for_atom)
-
-
-def _make_rigid_transformation_4x4(ex: np.ndarray, ey: np.ndarray, translation: np.ndarray) -> np.ndarray:
- """Create a rigid 4x4 transformation matrix from two axes and transl."""
- # Normalize ex.
- ex_normalized = ex / np.linalg.norm(ex)
-
- # make ey perpendicular to ex
- ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
- ey_normalized /= np.linalg.norm(ey_normalized)
-
- # compute ez as cross product
- eznorm = np.cross(ex_normalized, ey_normalized)
- m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
- m = np.concatenate([m, [[0.0, 0.0, 0.0, 1.0]]], axis=0)
- return m
-
-
-# create an array with (restype, atomtype) --> rigid_group_idx
-# and an array with (restype, atomtype, coord) for the atom positions
-# and compute affine transformation matrices (4,4) from one rigid group to the
-# previous group
-restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=int)
-restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
-restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
-restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=int)
-restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
-restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
-restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)
-
-
-def _make_rigid_group_constants() -> None:
- """Fill the arrays above."""
- for restype, restype_letter in enumerate(restypes):
- resname = restype_1to3[restype_letter]
- for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
- atomtype = atom_order[atomname]
- restype_atom37_to_rigid_group[restype, atomtype] = group_idx
- restype_atom37_mask[restype, atomtype] = 1
- restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position
-
- atom14idx = restype_name_to_atom14_names[resname].index(atomname)
- restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
- restype_atom14_mask[restype, atom14idx] = 1
- restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position
-
- for restype, restype_letter in enumerate(restypes):
- resname = restype_1to3[restype_letter]
- atom_positions: Dict[str, np.ndarray] = {
- name: np.array(pos) for name, _, pos in rigid_group_atom_positions[resname]
- }
-
- # backbone to backbone is the identity transform
- restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)
-
- # pre-omega-frame to backbone (currently dummy identity matrix)
- restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)
-
- # phi-frame to backbone
- mat = _make_rigid_transformation_4x4(
- ex=atom_positions["N"] - atom_positions["CA"],
- ey=np.array([1.0, 0.0, 0.0]),
- translation=atom_positions["N"],
- )
- restype_rigid_group_default_frame[restype, 2, :, :] = mat
-
- # psi-frame to backbone
- mat = _make_rigid_transformation_4x4(
- ex=atom_positions["C"] - atom_positions["CA"],
- ey=atom_positions["CA"] - atom_positions["N"],
- translation=atom_positions["C"],
- )
- restype_rigid_group_default_frame[restype, 3, :, :] = mat
-
- # chi1-frame to backbone
- if chi_angles_mask[restype][0]:
- base_atom_names = chi_angles_atoms[resname][0]
- base_atom_positions = [atom_positions[name] for name in base_atom_names]
- mat = _make_rigid_transformation_4x4(
- ex=base_atom_positions[2] - base_atom_positions[1],
- ey=base_atom_positions[0] - base_atom_positions[1],
- translation=base_atom_positions[2],
- )
- restype_rigid_group_default_frame[restype, 4, :, :] = mat
-
- # chi2-frame to chi1-frame
- # chi3-frame to chi2-frame
- # chi4-frame to chi3-frame
- # luckily all rotation axes for the next frame start at (0,0,0) of the
- # previous frame
- for chi_idx in range(1, 4):
- if chi_angles_mask[restype][chi_idx]:
- axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
- axis_end_atom_position = atom_positions[axis_end_atom_name]
- mat = _make_rigid_transformation_4x4(
- ex=axis_end_atom_position,
- ey=np.array([-1.0, 0.0, 0.0]),
- translation=axis_end_atom_position,
- )
- restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat
-
-
-_make_rigid_group_constants()
-
-
-def make_atom14_dists_bounds(
- overlap_tolerance: float = 1.5,
- bond_length_tolerance_factor: int = 15,
-) -> Dict[str, np.ndarray]:
- """compute upper and lower bounds for bonds to assess violations."""
- restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
- restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
- restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
- residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
- for restype, restype_letter in enumerate(restypes):
- resname = restype_1to3[restype_letter]
- atom_list = restype_name_to_atom14_names[resname]
-
- # create lower and upper bounds for clashes
- for atom1_idx, atom1_name in enumerate(atom_list):
- if not atom1_name:
- continue
- atom1_radius = van_der_waals_radius[atom1_name[0]]
- for atom2_idx, atom2_name in enumerate(atom_list):
- if (not atom2_name) or atom1_idx == atom2_idx:
- continue
- atom2_radius = van_der_waals_radius[atom2_name[0]]
- lower = atom1_radius + atom2_radius - overlap_tolerance
- upper = 1e10
- restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
- restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
- restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
- restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
-
- # overwrite lower and upper bounds for bonds and angles
- for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
- atom1_idx = atom_list.index(b.atom1_name)
- atom2_idx = atom_list.index(b.atom2_name)
- lower = b.length - bond_length_tolerance_factor * b.stddev
- upper = b.length + bond_length_tolerance_factor * b.stddev
- restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
- restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
- restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
- restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
- restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
- restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
- return {
- "lower_bound": restype_atom14_bond_lower_bound, # shape (21,14,14)
- "upper_bound": restype_atom14_bond_upper_bound, # shape (21,14,14)
- "stddev": restype_atom14_bond_stddev, # shape (21,14,14)
- }
-
-
-restype_atom14_ambiguous_atoms = np.zeros((21, 14), dtype=np.float32)
-restype_atom14_ambiguous_atoms_swap_idx: np.ndarray = np.tile(np.arange(14, dtype=int), (21, 1))
-
-
-def _make_atom14_ambiguity_feats() -> None:
- for res, pairs in residue_atom_renaming_swaps.items():
- res_idx = restype_order[restype_3to1[res]]
- for atom1, atom2 in pairs.items():
- atom1_idx = restype_name_to_atom14_names[res].index(atom1)
- atom2_idx = restype_name_to_atom14_names[res].index(atom2)
- restype_atom14_ambiguous_atoms[res_idx, atom1_idx] = 1
- restype_atom14_ambiguous_atoms[res_idx, atom2_idx] = 1
- restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom1_idx] = atom2_idx
- restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom2_idx] = atom1_idx
-
-
-_make_atom14_ambiguity_feats()
-
-
-def aatype_to_str_sequence(aatype: Sequence[int]) -> str:
- return "".join([restypes_with_x[aatype[i]] for i in range(len(aatype))])
diff --git a/transformers/models/esm/openfold_utils/rigid_utils.py b/transformers/models/esm/openfold_utils/rigid_utils.py
deleted file mode 100644
index 2bc2fe5f5c4ebff888e2d66eae3647073be89b4f..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/rigid_utils.py
+++ /dev/null
@@ -1,1242 +0,0 @@
-# Copyright 2021 AlQuraishi Laboratory
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from functools import lru_cache
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
-
-import numpy as np
-import torch
-
-
-def rot_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
- """
- Performs matrix multiplication of two rotation matrix tensors. Written out by hand to avoid AMP downcasting.
-
- Args:
- a: [*, 3, 3] left multiplicand
- b: [*, 3, 3] right multiplicand
- Returns:
- The product ab
- """
-
- def row_mul(i: int) -> torch.Tensor:
- return torch.stack(
- [
- a[..., i, 0] * b[..., 0, 0] + a[..., i, 1] * b[..., 1, 0] + a[..., i, 2] * b[..., 2, 0],
- a[..., i, 0] * b[..., 0, 1] + a[..., i, 1] * b[..., 1, 1] + a[..., i, 2] * b[..., 2, 1],
- a[..., i, 0] * b[..., 0, 2] + a[..., i, 1] * b[..., 1, 2] + a[..., i, 2] * b[..., 2, 2],
- ],
- dim=-1,
- )
-
- return torch.stack(
- [
- row_mul(0),
- row_mul(1),
- row_mul(2),
- ],
- dim=-2,
- )
-
-
-def rot_vec_mul(r: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
- """
- Applies a rotation to a vector. Written out by hand to avoid transfer to avoid AMP downcasting.
-
- Args:
- r: [*, 3, 3] rotation matrices
- t: [*, 3] coordinate tensors
- Returns:
- [*, 3] rotated coordinates
- """
- x, y, z = torch.unbind(t, dim=-1)
- return torch.stack(
- [
- r[..., 0, 0] * x + r[..., 0, 1] * y + r[..., 0, 2] * z,
- r[..., 1, 0] * x + r[..., 1, 1] * y + r[..., 1, 2] * z,
- r[..., 2, 0] * x + r[..., 2, 1] * y + r[..., 2, 2] * z,
- ],
- dim=-1,
- )
-
-
-@lru_cache(maxsize=None)
-def identity_rot_mats(
- batch_dims: Tuple[int, ...],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- requires_grad: bool = True,
-) -> torch.Tensor:
- rots = torch.eye(3, dtype=dtype, device=device, requires_grad=requires_grad)
- rots = rots.view(*((1,) * len(batch_dims)), 3, 3)
- rots = rots.expand(*batch_dims, -1, -1)
- rots = rots.contiguous()
-
- return rots
-
-
-@lru_cache(maxsize=None)
-def identity_trans(
- batch_dims: Tuple[int, ...],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- requires_grad: bool = True,
-) -> torch.Tensor:
- trans = torch.zeros((*batch_dims, 3), dtype=dtype, device=device, requires_grad=requires_grad)
- return trans
-
-
-@lru_cache(maxsize=None)
-def identity_quats(
- batch_dims: Tuple[int, ...],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- requires_grad: bool = True,
-) -> torch.Tensor:
- quat = torch.zeros((*batch_dims, 4), dtype=dtype, device=device, requires_grad=requires_grad)
-
- with torch.no_grad():
- quat[..., 0] = 1
-
- return quat
-
-
-_quat_elements: List[str] = ["a", "b", "c", "d"]
-_qtr_keys: List[str] = [l1 + l2 for l1 in _quat_elements for l2 in _quat_elements]
-_qtr_ind_dict: Dict[str, int] = {key: ind for ind, key in enumerate(_qtr_keys)}
-
-
-def _to_mat(pairs: List[Tuple[str, int]]) -> np.ndarray:
- mat = np.zeros((4, 4))
- for key, value in pairs:
- ind = _qtr_ind_dict[key]
- mat[ind // 4][ind % 4] = value
-
- return mat
-
-
-_QTR_MAT = np.zeros((4, 4, 3, 3))
-_QTR_MAT[..., 0, 0] = _to_mat([("aa", 1), ("bb", 1), ("cc", -1), ("dd", -1)])
-_QTR_MAT[..., 0, 1] = _to_mat([("bc", 2), ("ad", -2)])
-_QTR_MAT[..., 0, 2] = _to_mat([("bd", 2), ("ac", 2)])
-_QTR_MAT[..., 1, 0] = _to_mat([("bc", 2), ("ad", 2)])
-_QTR_MAT[..., 1, 1] = _to_mat([("aa", 1), ("bb", -1), ("cc", 1), ("dd", -1)])
-_QTR_MAT[..., 1, 2] = _to_mat([("cd", 2), ("ab", -2)])
-_QTR_MAT[..., 2, 0] = _to_mat([("bd", 2), ("ac", -2)])
-_QTR_MAT[..., 2, 1] = _to_mat([("cd", 2), ("ab", 2)])
-_QTR_MAT[..., 2, 2] = _to_mat([("aa", 1), ("bb", -1), ("cc", -1), ("dd", 1)])
-
-
-def quat_to_rot(quat: torch.Tensor) -> torch.Tensor:
- """
- Converts a quaternion to a rotation matrix.
-
- Args:
- quat: [*, 4] quaternions
- Returns:
- [*, 3, 3] rotation matrices
- """
- # [*, 4, 4]
- quat = quat[..., None] * quat[..., None, :]
-
- # [4, 4, 3, 3]
- mat = _get_quat("_QTR_MAT", dtype=quat.dtype, device=quat.device)
-
- # [*, 4, 4, 3, 3]
- shaped_qtr_mat = mat.view((1,) * len(quat.shape[:-2]) + mat.shape)
- quat = quat[..., None, None] * shaped_qtr_mat
-
- # [*, 3, 3]
- return torch.sum(quat, dim=(-3, -4))
-
-
-def rot_to_quat(rot: torch.Tensor) -> torch.Tensor:
- if rot.shape[-2:] != (3, 3):
- raise ValueError("Input rotation is incorrectly shaped")
-
- [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = [[rot[..., i, j] for j in range(3)] for i in range(3)]
-
- k = [
- [
- xx + yy + zz,
- zy - yz,
- xz - zx,
- yx - xy,
- ],
- [
- zy - yz,
- xx - yy - zz,
- xy + yx,
- xz + zx,
- ],
- [
- xz - zx,
- xy + yx,
- yy - xx - zz,
- yz + zy,
- ],
- [
- yx - xy,
- xz + zx,
- yz + zy,
- zz - xx - yy,
- ],
- ]
-
- _, vectors = torch.linalg.eigh((1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2))
- return vectors[..., -1]
-
-
-_QUAT_MULTIPLY = np.zeros((4, 4, 4))
-_QUAT_MULTIPLY[:, :, 0] = [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, -1]]
-
-_QUAT_MULTIPLY[:, :, 1] = [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, -1, 0]]
-
-_QUAT_MULTIPLY[:, :, 2] = [[0, 0, 1, 0], [0, 0, 0, -1], [1, 0, 0, 0], [0, 1, 0, 0]]
-
-_QUAT_MULTIPLY[:, :, 3] = [[0, 0, 0, 1], [0, 0, 1, 0], [0, -1, 0, 0], [1, 0, 0, 0]]
-
-_QUAT_MULTIPLY_BY_VEC = _QUAT_MULTIPLY[:, 1:, :]
-
-_CACHED_QUATS: Dict[str, np.ndarray] = {
- "_QTR_MAT": _QTR_MAT,
- "_QUAT_MULTIPLY": _QUAT_MULTIPLY,
- "_QUAT_MULTIPLY_BY_VEC": _QUAT_MULTIPLY_BY_VEC,
-}
-
-
-@lru_cache(maxsize=None)
-def _get_quat(quat_key: str, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
- return torch.tensor(_CACHED_QUATS[quat_key], dtype=dtype, device=device)
-
-
-def quat_multiply(quat1: torch.Tensor, quat2: torch.Tensor) -> torch.Tensor:
- """Multiply a quaternion by another quaternion."""
- mat = _get_quat("_QUAT_MULTIPLY", dtype=quat1.dtype, device=quat1.device)
- reshaped_mat = mat.view((1,) * len(quat1.shape[:-1]) + mat.shape)
- return torch.sum(reshaped_mat * quat1[..., :, None, None] * quat2[..., None, :, None], dim=(-3, -2))
-
-
-def quat_multiply_by_vec(quat: torch.Tensor, vec: torch.Tensor) -> torch.Tensor:
- """Multiply a quaternion by a pure-vector quaternion."""
- mat = _get_quat("_QUAT_MULTIPLY_BY_VEC", dtype=quat.dtype, device=quat.device)
- reshaped_mat = mat.view((1,) * len(quat.shape[:-1]) + mat.shape)
- return torch.sum(reshaped_mat * quat[..., :, None, None] * vec[..., None, :, None], dim=(-3, -2))
-
-
-def invert_rot_mat(rot_mat: torch.Tensor) -> torch.Tensor:
- return rot_mat.transpose(-1, -2)
-
-
-def invert_quat(quat: torch.Tensor) -> torch.Tensor:
- quat_prime = quat.clone()
- quat_prime[..., 1:] *= -1
- inv = quat_prime / torch.sum(quat**2, dim=-1, keepdim=True)
- return inv
-
-
-class Rotation:
- """
- A 3D rotation. Depending on how the object is initialized, the rotation is represented by either a rotation matrix
- or a quaternion, though both formats are made available by helper functions. To simplify gradient computation, the
- underlying format of the rotation cannot be changed in-place. Like Rigid, the class is designed to mimic the
- behavior of a torch Tensor, almost as if each Rotation object were a tensor of rotations, in one format or another.
- """
-
- def __init__(
- self,
- rot_mats: Optional[torch.Tensor] = None,
- quats: Optional[torch.Tensor] = None,
- normalize_quats: bool = True,
- ):
- """
- Args:
- rot_mats:
- A [*, 3, 3] rotation matrix tensor. Mutually exclusive with quats
- quats:
- A [*, 4] quaternion. Mutually exclusive with rot_mats. If normalize_quats is not True, must be a unit
- quaternion
- normalize_quats:
- If quats is specified, whether to normalize quats
- """
- if (rot_mats is None and quats is None) or (rot_mats is not None and quats is not None):
- raise ValueError("Exactly one input argument must be specified")
-
- if (rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or (quats is not None and quats.shape[-1] != 4):
- raise ValueError("Incorrectly shaped rotation matrix or quaternion")
-
- # Force full-precision
- if quats is not None:
- quats = quats.to(dtype=torch.float32)
- if rot_mats is not None:
- rot_mats = rot_mats.to(dtype=torch.float32)
-
- if quats is not None and normalize_quats:
- quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)
-
- self._rot_mats = rot_mats
- self._quats = quats
-
- @staticmethod
- def identity(
- shape,
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- requires_grad: bool = True,
- fmt: str = "quat",
- ) -> Rotation:
- """
- Returns an identity Rotation.
-
- Args:
- shape:
- The "shape" of the resulting Rotation object. See documentation for the shape property
- dtype:
- The torch dtype for the rotation
- device:
- The torch device for the new rotation
- requires_grad:
- Whether the underlying tensors in the new rotation object should require gradient computation
- fmt:
- One of "quat" or "rot_mat". Determines the underlying format of the new object's rotation
- Returns:
- A new identity rotation
- """
- if fmt == "rot_mat":
- rot_mats = identity_rot_mats(
- shape,
- dtype,
- device,
- requires_grad,
- )
- return Rotation(rot_mats=rot_mats, quats=None)
- elif fmt == "quat":
- quats = identity_quats(shape, dtype, device, requires_grad)
- return Rotation(rot_mats=None, quats=quats, normalize_quats=False)
- else:
- raise ValueError(f"Invalid format: f{fmt}")
-
- # Magic methods
-
- def __getitem__(self, index: Any) -> Rotation:
- """
- Allows torch-style indexing over the virtual shape of the rotation object. See documentation for the shape
- property.
-
- Args:
- index:
- A torch index. E.g. (1, 3, 2), or (slice(None,))
- Returns:
- The indexed rotation
- """
- if type(index) != tuple:
- index = (index,)
-
- if self._rot_mats is not None:
- rot_mats = self._rot_mats[index + (slice(None), slice(None))]
- return Rotation(rot_mats=rot_mats)
- elif self._quats is not None:
- quats = self._quats[index + (slice(None),)]
- return Rotation(quats=quats, normalize_quats=False)
- else:
- raise ValueError("Both rotations are None")
-
- def __mul__(self, right: torch.Tensor) -> Rotation:
- """
- Pointwise left multiplication of the rotation with a tensor. Can be used to e.g. mask the Rotation.
-
- Args:
- right:
- The tensor multiplicand
- Returns:
- The product
- """
- if not (isinstance(right, torch.Tensor)):
- raise TypeError("The other multiplicand must be a Tensor")
-
- if self._rot_mats is not None:
- rot_mats = self._rot_mats * right[..., None, None]
- return Rotation(rot_mats=rot_mats, quats=None)
- elif self._quats is not None:
- quats = self._quats * right[..., None]
- return Rotation(rot_mats=None, quats=quats, normalize_quats=False)
- else:
- raise ValueError("Both rotations are None")
-
- def __rmul__(self, left: torch.Tensor) -> Rotation:
- """
- Reverse pointwise multiplication of the rotation with a tensor.
-
- Args:
- left:
- The left multiplicand
- Returns:
- The product
- """
- return self.__mul__(left)
-
- # Properties
-
- @property
- def shape(self) -> torch.Size:
- """
- Returns the virtual shape of the rotation object. This shape is defined as the batch dimensions of the
- underlying rotation matrix or quaternion. If the Rotation was initialized with a [10, 3, 3] rotation matrix
- tensor, for example, the resulting shape would be [10].
-
- Returns:
- The virtual shape of the rotation object
- """
- if self._rot_mats is not None:
- return self._rot_mats.shape[:-2]
- elif self._quats is not None:
- return self._quats.shape[:-1]
- else:
- raise ValueError("Both rotations are None")
-
- @property
- def dtype(self) -> torch.dtype:
- """
- Returns the dtype of the underlying rotation.
-
- Returns:
- The dtype of the underlying rotation
- """
- if self._rot_mats is not None:
- return self._rot_mats.dtype
- elif self._quats is not None:
- return self._quats.dtype
- else:
- raise ValueError("Both rotations are None")
-
- @property
- def device(self) -> torch.device:
- """
- The device of the underlying rotation
-
- Returns:
- The device of the underlying rotation
- """
- if self._rot_mats is not None:
- return self._rot_mats.device
- elif self._quats is not None:
- return self._quats.device
- else:
- raise ValueError("Both rotations are None")
-
- @property
- def requires_grad(self) -> bool:
- """
- Returns the requires_grad property of the underlying rotation
-
- Returns:
- The requires_grad property of the underlying tensor
- """
- if self._rot_mats is not None:
- return self._rot_mats.requires_grad
- elif self._quats is not None:
- return self._quats.requires_grad
- else:
- raise ValueError("Both rotations are None")
-
- def get_rot_mats(self) -> torch.Tensor:
- """
- Returns the underlying rotation as a rotation matrix tensor.
-
- Returns:
- The rotation as a rotation matrix tensor
- """
- if self._rot_mats is not None:
- return self._rot_mats
- elif self._quats is not None:
- return quat_to_rot(self._quats)
- else:
- raise ValueError("Both rotations are None")
-
- def get_quats(self) -> torch.Tensor:
- """
- Returns the underlying rotation as a quaternion tensor.
-
- Depending on whether the Rotation was initialized with a quaternion, this function may call torch.linalg.eigh.
-
- Returns:
- The rotation as a quaternion tensor.
- """
- if self._rot_mats is not None:
- return rot_to_quat(self._rot_mats)
- elif self._quats is not None:
- return self._quats
- else:
- raise ValueError("Both rotations are None")
-
- def get_cur_rot(self) -> torch.Tensor:
- """
- Return the underlying rotation in its current form
-
- Returns:
- The stored rotation
- """
- if self._rot_mats is not None:
- return self._rot_mats
- elif self._quats is not None:
- return self._quats
- else:
- raise ValueError("Both rotations are None")
-
- # Rotation functions
-
- def compose_q_update_vec(self, q_update_vec: torch.Tensor, normalize_quats: bool = True) -> Rotation:
- """
- Returns a new quaternion Rotation after updating the current object's underlying rotation with a quaternion
- update, formatted as a [*, 3] tensor whose final three columns represent x, y, z such that (1, x, y, z) is the
- desired (not necessarily unit) quaternion update.
-
- Args:
- q_update_vec:
- A [*, 3] quaternion update tensor
- normalize_quats:
- Whether to normalize the output quaternion
- Returns:
- An updated Rotation
- """
- quats = self.get_quats()
- new_quats = quats + quat_multiply_by_vec(quats, q_update_vec)
- return Rotation(
- rot_mats=None,
- quats=new_quats,
- normalize_quats=normalize_quats,
- )
-
- def compose_r(self, r: Rotation) -> Rotation:
- """
- Compose the rotation matrices of the current Rotation object with those of another.
-
- Args:
- r:
- An update rotation object
- Returns:
- An updated rotation object
- """
- r1 = self.get_rot_mats()
- r2 = r.get_rot_mats()
- new_rot_mats = rot_matmul(r1, r2)
- return Rotation(rot_mats=new_rot_mats, quats=None)
-
- def compose_q(self, r: Rotation, normalize_quats: bool = True) -> Rotation:
- """
- Compose the quaternions of the current Rotation object with those of another.
-
- Depending on whether either Rotation was initialized with quaternions, this function may call
- torch.linalg.eigh.
-
- Args:
- r:
- An update rotation object
- Returns:
- An updated rotation object
- """
- q1 = self.get_quats()
- q2 = r.get_quats()
- new_quats = quat_multiply(q1, q2)
- return Rotation(rot_mats=None, quats=new_quats, normalize_quats=normalize_quats)
-
- def apply(self, pts: torch.Tensor) -> torch.Tensor:
- """
- Apply the current Rotation as a rotation matrix to a set of 3D coordinates.
-
- Args:
- pts:
- A [*, 3] set of points
- Returns:
- [*, 3] rotated points
- """
- rot_mats = self.get_rot_mats()
- return rot_vec_mul(rot_mats, pts)
-
- def invert_apply(self, pts: torch.Tensor) -> torch.Tensor:
- """
- The inverse of the apply() method.
-
- Args:
- pts:
- A [*, 3] set of points
- Returns:
- [*, 3] inverse-rotated points
- """
- rot_mats = self.get_rot_mats()
- inv_rot_mats = invert_rot_mat(rot_mats)
- return rot_vec_mul(inv_rot_mats, pts)
-
- def invert(self) -> Rotation:
- """
- Returns the inverse of the current Rotation.
-
- Returns:
- The inverse of the current Rotation
- """
- if self._rot_mats is not None:
- return Rotation(rot_mats=invert_rot_mat(self._rot_mats), quats=None)
- elif self._quats is not None:
- return Rotation(
- rot_mats=None,
- quats=invert_quat(self._quats),
- normalize_quats=False,
- )
- else:
- raise ValueError("Both rotations are None")
-
- # "Tensor" stuff
-
- def unsqueeze(self, dim: int) -> Rotation:
- """
- Analogous to torch.unsqueeze. The dimension is relative to the shape of the Rotation object.
-
- Args:
- dim: A positive or negative dimension index.
- Returns:
- The unsqueezed Rotation.
- """
- if dim >= len(self.shape):
- raise ValueError("Invalid dimension")
-
- if self._rot_mats is not None:
- rot_mats = self._rot_mats.unsqueeze(dim if dim >= 0 else dim - 2)
- return Rotation(rot_mats=rot_mats, quats=None)
- elif self._quats is not None:
- quats = self._quats.unsqueeze(dim if dim >= 0 else dim - 1)
- return Rotation(rot_mats=None, quats=quats, normalize_quats=False)
- else:
- raise ValueError("Both rotations are None")
-
- @staticmethod
- def cat(rs: Sequence[Rotation], dim: int) -> Rotation:
- """
- Concatenates rotations along one of the batch dimensions. Analogous to torch.cat().
-
- Note that the output of this operation is always a rotation matrix, regardless of the format of input
- rotations.
-
- Args:
- rs:
- A list of rotation objects
- dim:
- The dimension along which the rotations should be concatenated
- Returns:
- A concatenated Rotation object in rotation matrix format
- """
- rot_mats = torch.cat(
- [r.get_rot_mats() for r in rs],
- dim=dim if dim >= 0 else dim - 2,
- )
-
- return Rotation(rot_mats=rot_mats, quats=None)
-
- def map_tensor_fn(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> Rotation:
- """
- Apply a Tensor -> Tensor function to underlying rotation tensors, mapping over the rotation dimension(s). Can
- be used e.g. to sum out a one-hot batch dimension.
-
- Args:
- fn:
- A Tensor -> Tensor function to be mapped over the Rotation
- Returns:
- The transformed Rotation object
- """
- if self._rot_mats is not None:
- rot_mats = self._rot_mats.view(self._rot_mats.shape[:-2] + (9,))
- rot_mats = torch.stack(list(map(fn, torch.unbind(rot_mats, dim=-1))), dim=-1)
- rot_mats = rot_mats.view(rot_mats.shape[:-1] + (3, 3))
- return Rotation(rot_mats=rot_mats, quats=None)
- elif self._quats is not None:
- quats = torch.stack(list(map(fn, torch.unbind(self._quats, dim=-1))), dim=-1)
- return Rotation(rot_mats=None, quats=quats, normalize_quats=False)
- else:
- raise ValueError("Both rotations are None")
-
- def cuda(self) -> Rotation:
- """
- Analogous to the cuda() method of torch Tensors
-
- Returns:
- A copy of the Rotation in CUDA memory
- """
- if self._rot_mats is not None:
- return Rotation(rot_mats=self._rot_mats.cuda(), quats=None)
- elif self._quats is not None:
- return Rotation(rot_mats=None, quats=self._quats.cuda(), normalize_quats=False)
- else:
- raise ValueError("Both rotations are None")
-
- def to(self, device: Optional[torch.device], dtype: Optional[torch.dtype]) -> Rotation:
- """
- Analogous to the to() method of torch Tensors
-
- Args:
- device:
- A torch device
- dtype:
- A torch dtype
- Returns:
- A copy of the Rotation using the new device and dtype
- """
- if self._rot_mats is not None:
- return Rotation(
- rot_mats=self._rot_mats.to(device=device, dtype=dtype),
- quats=None,
- )
- elif self._quats is not None:
- return Rotation(
- rot_mats=None,
- quats=self._quats.to(device=device, dtype=dtype),
- normalize_quats=False,
- )
- else:
- raise ValueError("Both rotations are None")
-
- def detach(self) -> Rotation:
- """
- Returns a copy of the Rotation whose underlying Tensor has been detached from its torch graph.
-
- Returns:
- A copy of the Rotation whose underlying Tensor has been detached from its torch graph
- """
- if self._rot_mats is not None:
- return Rotation(rot_mats=self._rot_mats.detach(), quats=None)
- elif self._quats is not None:
- return Rotation(
- rot_mats=None,
- quats=self._quats.detach(),
- normalize_quats=False,
- )
- else:
- raise ValueError("Both rotations are None")
-
-
-class Rigid:
- """
- A class representing a rigid transformation. Little more than a wrapper around two objects: a Rotation object and a
- [*, 3] translation Designed to behave approximately like a single torch tensor with the shape of the shared batch
- dimensions of its component parts.
- """
-
- def __init__(self, rots: Optional[Rotation], trans: Optional[torch.Tensor]):
- """
- Args:
- rots: A [*, 3, 3] rotation tensor
- trans: A corresponding [*, 3] translation tensor
- """
- # (we need device, dtype, etc. from at least one input)
-
- batch_dims, dtype, device, requires_grad = None, None, None, None
- if trans is not None:
- batch_dims = trans.shape[:-1]
- dtype = trans.dtype
- device = trans.device
- requires_grad = trans.requires_grad
- elif rots is not None:
- batch_dims = rots.shape
- dtype = rots.dtype
- device = rots.device
- requires_grad = rots.requires_grad
- else:
- raise ValueError("At least one input argument must be specified")
-
- if rots is None:
- rots = Rotation.identity(
- batch_dims,
- dtype,
- device,
- requires_grad,
- )
- elif trans is None:
- trans = identity_trans(
- batch_dims,
- dtype,
- device,
- requires_grad,
- )
-
- assert rots is not None
- assert trans is not None
-
- if (rots.shape != trans.shape[:-1]) or (rots.device != trans.device):
- raise ValueError("Rots and trans incompatible")
-
- # Force full precision. Happens to the rotations automatically.
- trans = trans.to(dtype=torch.float32)
-
- self._rots = rots
- self._trans = trans
-
- @staticmethod
- def identity(
- shape: Tuple[int, ...],
- dtype: Optional[torch.dtype] = None,
- device: Optional[torch.device] = None,
- requires_grad: bool = True,
- fmt: str = "quat",
- ) -> Rigid:
- """
- Constructs an identity transformation.
-
- Args:
- shape:
- The desired shape
- dtype:
- The dtype of both internal tensors
- device:
- The device of both internal tensors
- requires_grad:
- Whether grad should be enabled for the internal tensors
- Returns:
- The identity transformation
- """
- return Rigid(
- Rotation.identity(shape, dtype, device, requires_grad, fmt=fmt),
- identity_trans(shape, dtype, device, requires_grad),
- )
-
- def __getitem__(self, index: Any) -> Rigid:
- """
- Indexes the affine transformation with PyTorch-style indices. The index is applied to the shared dimensions of
- both the rotation and the translation.
-
- E.g.::
-
- r = Rotation(rot_mats=torch.rand(10, 10, 3, 3), quats=None) t = Rigid(r, torch.rand(10, 10, 3)) indexed =
- t[3, 4:6] assert(indexed.shape == (2,)) assert(indexed.get_rots().shape == (2,))
- assert(indexed.get_trans().shape == (2, 3))
-
- Args:
- index: A standard torch tensor index. E.g. 8, (10, None, 3),
- or (3, slice(0, 1, None))
- Returns:
- The indexed tensor
- """
- if type(index) != tuple:
- index = (index,)
-
- return Rigid(
- self._rots[index],
- self._trans[index + (slice(None),)],
- )
-
- def __mul__(self, right: torch.Tensor) -> Rigid:
- """
- Pointwise left multiplication of the transformation with a tensor. Can be used to e.g. mask the Rigid.
-
- Args:
- right:
- The tensor multiplicand
- Returns:
- The product
- """
- if not (isinstance(right, torch.Tensor)):
- raise TypeError("The other multiplicand must be a Tensor")
-
- new_rots = self._rots * right
- new_trans = self._trans * right[..., None]
-
- return Rigid(new_rots, new_trans)
-
- def __rmul__(self, left: torch.Tensor) -> Rigid:
- """
- Reverse pointwise multiplication of the transformation with a tensor.
-
- Args:
- left:
- The left multiplicand
- Returns:
- The product
- """
- return self.__mul__(left)
-
- @property
- def shape(self) -> torch.Size:
- """
- Returns the shape of the shared dimensions of the rotation and the translation.
-
- Returns:
- The shape of the transformation
- """
- return self._trans.shape[:-1]
-
- @property
- def device(self) -> torch.device:
- """
- Returns the device on which the Rigid's tensors are located.
-
- Returns:
- The device on which the Rigid's tensors are located
- """
- return self._trans.device
-
- def get_rots(self) -> Rotation:
- """
- Getter for the rotation.
-
- Returns:
- The rotation object
- """
- return self._rots
-
- def get_trans(self) -> torch.Tensor:
- """
- Getter for the translation.
-
- Returns:
- The stored translation
- """
- return self._trans
-
- def compose_q_update_vec(self, q_update_vec: torch.Tensor) -> Rigid:
- """
- Composes the transformation with a quaternion update vector of shape [*, 6], where the final 6 columns
- represent the x, y, and z values of a quaternion of form (1, x, y, z) followed by a 3D translation.
-
- Args:
- q_vec: The quaternion update vector.
- Returns:
- The composed transformation.
- """
- q_vec, t_vec = q_update_vec[..., :3], q_update_vec[..., 3:]
- new_rots = self._rots.compose_q_update_vec(q_vec)
-
- trans_update = self._rots.apply(t_vec)
- new_translation = self._trans + trans_update
-
- return Rigid(new_rots, new_translation)
-
- def compose(self, r: Rigid) -> Rigid:
- """
- Composes the current rigid object with another.
-
- Args:
- r:
- Another Rigid object
- Returns:
- The composition of the two transformations
- """
- new_rot = self._rots.compose_r(r._rots)
- new_trans = self._rots.apply(r._trans) + self._trans
- return Rigid(new_rot, new_trans)
-
- def apply(self, pts: torch.Tensor) -> torch.Tensor:
- """
- Applies the transformation to a coordinate tensor.
-
- Args:
- pts: A [*, 3] coordinate tensor.
- Returns:
- The transformed points.
- """
- rotated = self._rots.apply(pts)
- return rotated + self._trans
-
- def invert_apply(self, pts: torch.Tensor) -> torch.Tensor:
- """
- Applies the inverse of the transformation to a coordinate tensor.
-
- Args:
- pts: A [*, 3] coordinate tensor
- Returns:
- The transformed points.
- """
- pts = pts - self._trans
- return self._rots.invert_apply(pts)
-
- def invert(self) -> Rigid:
- """
- Inverts the transformation.
-
- Returns:
- The inverse transformation.
- """
- rot_inv = self._rots.invert()
- trn_inv = rot_inv.apply(self._trans)
-
- return Rigid(rot_inv, -1 * trn_inv)
-
- def map_tensor_fn(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> Rigid:
- """
- Apply a Tensor -> Tensor function to underlying translation and rotation tensors, mapping over the
- translation/rotation dimensions respectively.
-
- Args:
- fn:
- A Tensor -> Tensor function to be mapped over the Rigid
- Returns:
- The transformed Rigid object
- """
- new_rots = self._rots.map_tensor_fn(fn)
- new_trans = torch.stack(list(map(fn, torch.unbind(self._trans, dim=-1))), dim=-1)
-
- return Rigid(new_rots, new_trans)
-
- def to_tensor_4x4(self) -> torch.Tensor:
- """
- Converts a transformation to a homogenous transformation tensor.
-
- Returns:
- A [*, 4, 4] homogenous transformation tensor
- """
- tensor = self._trans.new_zeros((*self.shape, 4, 4))
- tensor[..., :3, :3] = self._rots.get_rot_mats()
- tensor[..., :3, 3] = self._trans
- tensor[..., 3, 3] = 1
- return tensor
-
- @staticmethod
- def from_tensor_4x4(t: torch.Tensor) -> Rigid:
- """
- Constructs a transformation from a homogenous transformation tensor.
-
- Args:
- t: [*, 4, 4] homogenous transformation tensor
- Returns:
- T object with shape [*]
- """
- if t.shape[-2:] != (4, 4):
- raise ValueError("Incorrectly shaped input tensor")
-
- rots = Rotation(rot_mats=t[..., :3, :3], quats=None)
- trans = t[..., :3, 3]
-
- return Rigid(rots, trans)
-
- def to_tensor_7(self) -> torch.Tensor:
- """
- Converts a transformation to a tensor with 7 final columns, four for the quaternion followed by three for the
- translation.
-
- Returns:
- A [*, 7] tensor representation of the transformation
- """
- tensor = self._trans.new_zeros((*self.shape, 7))
- tensor[..., :4] = self._rots.get_quats()
- tensor[..., 4:] = self._trans
-
- return tensor
-
- @staticmethod
- def from_tensor_7(t: torch.Tensor, normalize_quats: bool = False) -> Rigid:
- if t.shape[-1] != 7:
- raise ValueError("Incorrectly shaped input tensor")
-
- quats, trans = t[..., :4], t[..., 4:]
-
- rots = Rotation(rot_mats=None, quats=quats, normalize_quats=normalize_quats)
-
- return Rigid(rots, trans)
-
- @staticmethod
- def from_3_points(
- p_neg_x_axis: torch.Tensor, origin: torch.Tensor, p_xy_plane: torch.Tensor, eps: float = 1e-8
- ) -> Rigid:
- """
- Implements algorithm 21. Constructs transformations from sets of 3 points using the Gram-Schmidt algorithm.
-
- Args:
- p_neg_x_axis: [*, 3] coordinates
- origin: [*, 3] coordinates used as frame origins
- p_xy_plane: [*, 3] coordinates
- eps: Small epsilon value
- Returns:
- A transformation object of shape [*]
- """
- p_neg_x_axis_unbound = torch.unbind(p_neg_x_axis, dim=-1)
- origin_unbound = torch.unbind(origin, dim=-1)
- p_xy_plane_unbound = torch.unbind(p_xy_plane, dim=-1)
-
- e0 = [c1 - c2 for c1, c2 in zip(origin_unbound, p_neg_x_axis_unbound)]
- e1 = [c1 - c2 for c1, c2 in zip(p_xy_plane_unbound, origin_unbound)]
-
- denom = torch.sqrt(sum(c * c for c in e0) + eps * torch.ones_like(e0[0]))
- e0 = [c / denom for c in e0]
- dot = sum((c1 * c2 for c1, c2 in zip(e0, e1)))
- e1 = [c2 - c1 * dot for c1, c2 in zip(e0, e1)]
- denom = torch.sqrt(sum((c * c for c in e1)) + eps * torch.ones_like(e1[0]))
- e1 = [c / denom for c in e1]
- e2 = [
- e0[1] * e1[2] - e0[2] * e1[1],
- e0[2] * e1[0] - e0[0] * e1[2],
- e0[0] * e1[1] - e0[1] * e1[0],
- ]
-
- rots = torch.stack([c for tup in zip(e0, e1, e2) for c in tup], dim=-1)
- rots = rots.reshape(rots.shape[:-1] + (3, 3))
-
- rot_obj = Rotation(rot_mats=rots, quats=None)
-
- return Rigid(rot_obj, torch.stack(origin_unbound, dim=-1))
-
- def unsqueeze(self, dim: int) -> Rigid:
- """
- Analogous to torch.unsqueeze. The dimension is relative to the shared dimensions of the rotation/translation.
-
- Args:
- dim: A positive or negative dimension index.
- Returns:
- The unsqueezed transformation.
- """
- if dim >= len(self.shape):
- raise ValueError("Invalid dimension")
- rots = self._rots.unsqueeze(dim)
- trans = self._trans.unsqueeze(dim if dim >= 0 else dim - 1)
-
- return Rigid(rots, trans)
-
- @staticmethod
- def cat(ts: Sequence[Rigid], dim: int) -> Rigid:
- """
- Concatenates transformations along a new dimension.
-
- Args:
- ts:
- A list of T objects
- dim:
- The dimension along which the transformations should be concatenated
- Returns:
- A concatenated transformation object
- """
- rots = Rotation.cat([t._rots for t in ts], dim)
- trans = torch.cat([t._trans for t in ts], dim=dim if dim >= 0 else dim - 1)
-
- return Rigid(rots, trans)
-
- def apply_rot_fn(self, fn: Callable[[Rotation], Rotation]) -> Rigid:
- """
- Applies a Rotation -> Rotation function to the stored rotation object.
-
- Args:
- fn: A function of type Rotation -> Rotation
- Returns:
- A transformation object with a transformed rotation.
- """
- return Rigid(fn(self._rots), self._trans)
-
- def apply_trans_fn(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> Rigid:
- """
- Applies a Tensor -> Tensor function to the stored translation.
-
- Args:
- fn:
- A function of type Tensor -> Tensor to be applied to the translation
- Returns:
- A transformation object with a transformed translation.
- """
- return Rigid(self._rots, fn(self._trans))
-
- def scale_translation(self, trans_scale_factor: float) -> Rigid:
- """
- Scales the translation by a constant factor.
-
- Args:
- trans_scale_factor:
- The constant factor
- Returns:
- A transformation object with a scaled translation.
- """
- return self.apply_trans_fn(lambda t: t * trans_scale_factor)
-
- def stop_rot_gradient(self) -> Rigid:
- """
- Detaches the underlying rotation object
-
- Returns:
- A transformation object with detached rotations
- """
- return self.apply_rot_fn(lambda r: r.detach())
-
- @staticmethod
- def make_transform_from_reference(
- n_xyz: torch.Tensor, ca_xyz: torch.Tensor, c_xyz: torch.Tensor, eps: float = 1e-20
- ) -> Rigid:
- """
- Returns a transformation object from reference coordinates.
-
- Note that this method does not take care of symmetries. If you provide the atom positions in the non-standard
- way, the N atom will end up not at [-0.527250, 1.359329, 0.0] but instead at [-0.527250, -1.359329, 0.0]. You
- need to take care of such cases in your code.
-
- Args:
- n_xyz: A [*, 3] tensor of nitrogen xyz coordinates.
- ca_xyz: A [*, 3] tensor of carbon alpha xyz coordinates.
- c_xyz: A [*, 3] tensor of carbon xyz coordinates.
- Returns:
- A transformation object. After applying the translation and rotation to the reference backbone, the
- coordinates will approximately equal to the input coordinates.
- """
- translation = -1 * ca_xyz
- n_xyz = n_xyz + translation
- c_xyz = c_xyz + translation
-
- c_x, c_y, c_z = [c_xyz[..., i] for i in range(3)]
- norm = torch.sqrt(eps + c_x**2 + c_y**2)
- sin_c1 = -c_y / norm
- cos_c1 = c_x / norm
-
- c1_rots = sin_c1.new_zeros((*sin_c1.shape, 3, 3))
- c1_rots[..., 0, 0] = cos_c1
- c1_rots[..., 0, 1] = -1 * sin_c1
- c1_rots[..., 1, 0] = sin_c1
- c1_rots[..., 1, 1] = cos_c1
- c1_rots[..., 2, 2] = 1
-
- norm = torch.sqrt(eps + c_x**2 + c_y**2 + c_z**2)
- sin_c2 = c_z / norm
- cos_c2 = torch.sqrt(c_x**2 + c_y**2) / norm
-
- c2_rots = sin_c2.new_zeros((*sin_c2.shape, 3, 3))
- c2_rots[..., 0, 0] = cos_c2
- c2_rots[..., 0, 2] = sin_c2
- c2_rots[..., 1, 1] = 1
- c2_rots[..., 2, 0] = -1 * sin_c2
- c2_rots[..., 2, 2] = cos_c2
-
- c_rots = rot_matmul(c2_rots, c1_rots)
- n_xyz = rot_vec_mul(c_rots, n_xyz)
-
- _, n_y, n_z = [n_xyz[..., i] for i in range(3)]
- norm = torch.sqrt(eps + n_y**2 + n_z**2)
- sin_n = -n_z / norm
- cos_n = n_y / norm
-
- n_rots = sin_c2.new_zeros((*sin_c2.shape, 3, 3))
- n_rots[..., 0, 0] = 1
- n_rots[..., 1, 1] = cos_n
- n_rots[..., 1, 2] = -1 * sin_n
- n_rots[..., 2, 1] = sin_n
- n_rots[..., 2, 2] = cos_n
-
- rots = rot_matmul(n_rots, c_rots)
-
- rots = rots.transpose(-1, -2)
- translation = -1 * translation
-
- rot_obj = Rotation(rot_mats=rots, quats=None)
-
- return Rigid(rot_obj, translation)
-
- def cuda(self) -> Rigid:
- """
- Moves the transformation object to GPU memory
-
- Returns:
- A version of the transformation on GPU
- """
- return Rigid(self._rots.cuda(), self._trans.cuda())
diff --git a/transformers/models/esm/openfold_utils/tensor_utils.py b/transformers/models/esm/openfold_utils/tensor_utils.py
deleted file mode 100644
index 99dd6dbe47b68247794e51810fd274c6352e5b4f..0000000000000000000000000000000000000000
--- a/transformers/models/esm/openfold_utils/tensor_utils.py
+++ /dev/null
@@ -1,144 +0,0 @@
-# Copyright 2021 AlQuraishi Laboratory
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from functools import partial
-from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload
-
-import torch
-import torch.nn as nn
-import torch.types
-
-
-def add(m1: torch.Tensor, m2: torch.Tensor, inplace: bool) -> torch.Tensor:
- # The first operation in a checkpoint can't be in-place, but it's
- # nice to have in-place addition during inference. Thus...
- if not inplace:
- m1 = m1 + m2
- else:
- m1 += m2
-
- return m1
-
-
-def permute_final_dims(tensor: torch.Tensor, inds: List[int]) -> torch.Tensor:
- zero_index = -1 * len(inds)
- first_inds = list(range(len(tensor.shape[:zero_index])))
- return tensor.permute(first_inds + [zero_index + i for i in inds])
-
-
-def flatten_final_dims(t: torch.Tensor, no_dims: int) -> torch.Tensor:
- return t.reshape(t.shape[:-no_dims] + (-1,))
-
-
-def masked_mean(mask: torch.Tensor, value: torch.Tensor, dim: int, eps: float = 1e-4) -> torch.Tensor:
- mask = mask.expand(*value.shape)
- return torch.sum(mask * value, dim=dim) / (eps + torch.sum(mask, dim=dim))
-
-
-def pts_to_distogram(
- pts: torch.Tensor, min_bin: torch.types.Number = 2.3125, max_bin: torch.types.Number = 21.6875, no_bins: int = 64
-) -> torch.Tensor:
- boundaries = torch.linspace(min_bin, max_bin, no_bins - 1, device=pts.device)
- dists = torch.sqrt(torch.sum((pts.unsqueeze(-2) - pts.unsqueeze(-3)) ** 2, dim=-1))
- return torch.bucketize(dists, boundaries)
-
-
-def dict_multimap(fn: Callable[[list], Any], dicts: List[dict]) -> dict:
- first = dicts[0]
- new_dict = {}
- for k, v in first.items():
- all_v = [d[k] for d in dicts]
- if isinstance(v, dict):
- new_dict[k] = dict_multimap(fn, all_v)
- else:
- new_dict[k] = fn(all_v)
-
- return new_dict
-
-
-def one_hot(x: torch.Tensor, v_bins: torch.Tensor) -> torch.Tensor:
- reshaped_bins = v_bins.view(((1,) * len(x.shape)) + (len(v_bins),))
- diffs = x[..., None] - reshaped_bins
- am = torch.argmin(torch.abs(diffs), dim=-1)
- return nn.functional.one_hot(am, num_classes=len(v_bins)).float()
-
-
-def batched_gather(data: torch.Tensor, inds: torch.Tensor, dim: int = 0, no_batch_dims: int = 0) -> torch.Tensor:
- ranges: List[Union[slice, torch.Tensor]] = []
- for i, s in enumerate(data.shape[:no_batch_dims]):
- r = torch.arange(s)
- r = r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape) - i - 1))))
- ranges.append(r)
-
- remaining_dims: List[Union[slice, torch.Tensor]] = [slice(None) for _ in range(len(data.shape) - no_batch_dims)]
- remaining_dims[dim - no_batch_dims if dim >= 0 else dim] = inds
- ranges.extend(remaining_dims)
- # Matt note: Editing this to get around the behaviour of using a list as an array index changing
- # in recent Numpy versions
- return data[tuple(ranges)]
-
-
-T = TypeVar("T")
-
-
-# With tree_map, a poor man's JAX tree_map
-def dict_map(
- fn: Callable[[T], Any], dic: Dict[Any, Union[dict, list, tuple, T]], leaf_type: Type[T]
-) -> Dict[Any, Union[dict, list, tuple, Any]]:
- new_dict: Dict[Any, Union[dict, list, tuple, Any]] = {}
- for k, v in dic.items():
- if isinstance(v, dict):
- new_dict[k] = dict_map(fn, v, leaf_type)
- else:
- new_dict[k] = tree_map(fn, v, leaf_type)
-
- return new_dict
-
-
-@overload
-def tree_map(fn: Callable[[T], Any], tree: T, leaf_type: Type[T]) -> Any:
- ...
-
-
-@overload
-def tree_map(fn: Callable[[T], Any], tree: dict, leaf_type: Type[T]) -> dict:
- ...
-
-
-@overload
-def tree_map(fn: Callable[[T], Any], tree: list, leaf_type: Type[T]) -> list:
- ...
-
-
-@overload
-def tree_map(fn: Callable[[T], Any], tree: tuple, leaf_type: Type[T]) -> tuple:
- ...
-
-
-def tree_map(fn, tree, leaf_type):
- if isinstance(tree, dict):
- return dict_map(fn, tree, leaf_type)
- elif isinstance(tree, list):
- return [tree_map(fn, x, leaf_type) for x in tree]
- elif isinstance(tree, tuple):
- return tuple(tree_map(fn, x, leaf_type) for x in tree)
- elif isinstance(tree, leaf_type):
- return fn(tree)
- else:
- print(type(tree))
- raise ValueError("Not supported")
-
-
-tensor_tree_map = partial(tree_map, leaf_type=torch.Tensor)
diff --git a/transformers/models/esm/tokenization_esm.py b/transformers/models/esm/tokenization_esm.py
deleted file mode 100644
index 27a889c87ea0b42397ed1553608aa2e5db2f85bc..0000000000000000000000000000000000000000
--- a/transformers/models/esm/tokenization_esm.py
+++ /dev/null
@@ -1,143 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for ESM."""
-import os
-from typing import List, Optional
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
-
-
-def load_vocab_file(vocab_file):
- with open(vocab_file, "r") as f:
- lines = f.read().splitlines()
- return [l.strip() for l in lines]
-
-
-class EsmTokenizer(PreTrainedTokenizer):
- """
- Constructs an ESM tokenizer.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- unk_token="",
- cls_token="",
- pad_token="",
- mask_token="",
- eos_token="",
- **kwargs,
- ):
- self.all_tokens = load_vocab_file(vocab_file)
- self._id_to_token = dict(enumerate(self.all_tokens))
- self._token_to_id = {tok: ind for ind, tok in enumerate(self.all_tokens)}
- super().__init__(
- unk_token=unk_token,
- cls_token=cls_token,
- pad_token=pad_token,
- mask_token=mask_token,
- eos_token=eos_token,
- **kwargs,
- )
-
- # TODO, all the tokens are added? But they are also part of the vocab... bit strange.
- # none of them are special, but they all need special splitting.
-
- self.unique_no_split_tokens = self.all_tokens
- self._update_trie(self.unique_no_split_tokens)
-
- def _convert_id_to_token(self, index: int) -> str:
- return self._id_to_token.get(index, self.unk_token)
-
- def _convert_token_to_id(self, token: str) -> int:
- return self._token_to_id.get(token, self._token_to_id.get(self.unk_token))
-
- def _tokenize(self, text, **kwargs):
- return text.split()
-
- def get_vocab(self):
- base_vocab = self._token_to_id.copy()
- base_vocab.update(self.added_tokens_encoder)
- return base_vocab
-
- def token_to_id(self, token: str) -> int:
- return self._token_to_id.get(token, self._token_to_id.get(self.unk_token))
-
- def id_to_token(self, index: int) -> str:
- return self._id_to_token.get(index, self.unk_token)
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- cls = [self.cls_token_id]
- sep = [self.eos_token_id] # No sep token in ESM vocabulary
- if token_ids_1 is None:
- if self.eos_token_id is None:
- return cls + token_ids_0
- else:
- return cls + token_ids_0 + sep
- elif self.eos_token_id is None:
- raise ValueError("Cannot tokenize multiple sequences when EOS token is not set!")
- return cls + token_ids_0 + sep + token_ids_1 + sep # Multiple inputs always have an EOS token
-
- def get_special_tokens_mask(
- self, token_ids_0: List, token_ids_1: Optional[List] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
-
- Args:
- token_ids_0 (`List[int]`):
- List of ids of the first sequence.
- token_ids_1 (`List[int]`, *optional*):
- List of ids of the second sequence.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
- if already_has_special_tokens:
- if token_ids_1 is not None:
- raise ValueError(
- "You should not supply a second sequence if the provided sequence of "
- "ids is already formatted with special tokens for the model."
- )
-
- return [1 if token in self.all_special_ids else 0 for token in token_ids_0]
- mask = [1] + ([0] * len(token_ids_0)) + [1]
- if token_ids_1 is not None:
- mask += [0] * len(token_ids_1) + [1]
- return mask
-
- def save_vocabulary(self, save_directory, filename_prefix):
- vocab_file = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.txt")
- with open(vocab_file, "w") as f:
- f.write("\n".join(self.all_tokens))
- return (vocab_file,)
-
- @property
- def vocab_size(self) -> int:
- return len(self.all_tokens)
diff --git a/transformers/models/falcon/__init__.py b/transformers/models/falcon/__init__.py
deleted file mode 100644
index 070e0cc033fbf6c364d2405bbf6367312e79a18d..0000000000000000000000000000000000000000
--- a/transformers/models/falcon/__init__.py
+++ /dev/null
@@ -1,68 +0,0 @@
-# coding=utf-8
-# Copyright 2023 the Falcon authors and HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_falcon": ["FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP", "FalconConfig"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_falcon"] = [
- "FALCON_PRETRAINED_MODEL_ARCHIVE_LIST",
- "FalconForCausalLM",
- "FalconModel",
- "FalconPreTrainedModel",
- "FalconForSequenceClassification",
- "FalconForTokenClassification",
- "FalconForQuestionAnswering",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_falcon import FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP, FalconConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_falcon import (
- FALCON_PRETRAINED_MODEL_ARCHIVE_LIST,
- FalconForCausalLM,
- FalconForQuestionAnswering,
- FalconForSequenceClassification,
- FalconForTokenClassification,
- FalconModel,
- FalconPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/falcon/__pycache__/__init__.cpython-310.pyc b/transformers/models/falcon/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 75760c96503b8029e11c575aa65168c0954d8869..0000000000000000000000000000000000000000
Binary files a/transformers/models/falcon/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/falcon/__pycache__/configuration_falcon.cpython-310.pyc b/transformers/models/falcon/__pycache__/configuration_falcon.cpython-310.pyc
deleted file mode 100644
index 379e314f7e4622ebe08243060d02464ddc522b79..0000000000000000000000000000000000000000
Binary files a/transformers/models/falcon/__pycache__/configuration_falcon.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/falcon/__pycache__/convert_custom_code_checkpoint.cpython-310.pyc b/transformers/models/falcon/__pycache__/convert_custom_code_checkpoint.cpython-310.pyc
deleted file mode 100644
index a8c936df95426c1ec5e516055632ba6a623f561c..0000000000000000000000000000000000000000
Binary files a/transformers/models/falcon/__pycache__/convert_custom_code_checkpoint.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/falcon/__pycache__/modeling_falcon.cpython-310.pyc b/transformers/models/falcon/__pycache__/modeling_falcon.cpython-310.pyc
deleted file mode 100644
index 328f6df052431d3f623efa74e1fca54d7759e1d1..0000000000000000000000000000000000000000
Binary files a/transformers/models/falcon/__pycache__/modeling_falcon.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/falcon/configuration_falcon.py b/transformers/models/falcon/configuration_falcon.py
deleted file mode 100644
index 61d202b0960829e6ea953d72d581b5858553e2bf..0000000000000000000000000000000000000000
--- a/transformers/models/falcon/configuration_falcon.py
+++ /dev/null
@@ -1,201 +0,0 @@
-# coding=utf-8
-# Copyright 2023 the Falcon authors and HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Falcon configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class FalconConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FalconModel`]. It is used to instantiate a Falcon
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the
- [tiiuae/falcon-7b](https://huggingface.co/tiiuae/falcon-7b) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 65024):
- Vocabulary size of the Falcon model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`FalconModel`]
- hidden_size (`int`, *optional*, defaults to 4544):
- Dimension of the hidden representations.
- num_hidden_layers (`int`, *optional*, defaults to 32):
- Number of hidden layers in the Transformer decoder.
- num_attention_heads (`int`, *optional*, defaults to 71):
- Number of attention heads for each attention layer in the Transformer encoder.
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the layer normalization layers.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether the model should return the last key/values attentions (not used by all models). Only relevant if
- `config.is_decoder=True`.
- hidden_dropout (`float`, *optional*, defaults to 0.0):
- The dropout probability for MLP layers.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout probability for attention layers.
- num_kv_heads (`int`, *optional*):
- Number of key-value heads to use per attention layer. If unset, defaults to the same value as
- `num_attention_heads`.
- alibi (`bool`, *optional*, defaults to `False`):
- Whether to use ALiBi positional biases during self-attention.
- new_decoder_architecture (`bool`, *optional*, defaults to `False`):
- Whether to use the new (Falcon-40B) decoder architecture. If `True`, the `multi_query` and `parallel_attn`
- arguments are ignored, as the new decoder always uses parallel attention.
- multi_query (`bool`, *optional*, defaults to `True`):
- Whether to use multi-query attention in the decoder. Ignored when `new_decoder_architecture` is `True`.
- parallel_attn (`bool`, *optional*, defaults to `True`):
- Whether to compute attention in parallel with the feedforward layer. If False, they are consecutive
- instead, as in the original Transformer architecture. Ignored when `new_decoder_architecture` is `True`.
- bias (`bool`, *optional*, defaults to `False`):
- Whether to use bias on Linear layers.
- max_position_embeddings (`int`, *optional*, defaults to 2048):
- The maximum sequence length that this model might ever be used with, when `alibi` is `False`. Pretrained
- Falcon models with RoPE support up to 2048 tokens.
- rope_theta (`float`, *optional*, defaults to 10000.0):
- The base period of the RoPE embeddings.
- rope_scaling (`Dict`, *optional*):
- Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
- strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
- `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
- `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
- these scaling strategies behave:
- https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
- experimental feature, subject to breaking API changes in future versions.
- bos_token_id (`int`, *optional*, defaults to 11):
- The id of the "beginning-of-sequence" token.
- eos_token_id (`int`, *optional*, defaults to 11):
- The id of the "end-of-sequence" token.
- ffn_hidden_size (`int`, *optional*):
- The hidden size of the feedforward layer in the Transformer decoder.
- defaults to 4x hidden dim
- activation (`str`, *optional*, defaults to `"gelu"`):
- The activation function used in the feedforward layer.
-
- Example:
-
- ```python
- >>> from transformers import FalconModel, FalconConfig
-
- >>> # Initializing a small (2-layer) Falcon configuration
- >>> configuration = FalconConfig(num_hidden_layers=2)
-
- >>> # Initializing a model from the small configuration
- >>> model = FalconModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "falcon"
- keys_to_ignore_at_inference = ["past_key_values"]
-
- def __init__(
- self,
- vocab_size=65024,
- hidden_size=4544,
- num_hidden_layers=32,
- num_attention_heads=71,
- layer_norm_epsilon=1e-5,
- initializer_range=0.02,
- use_cache=True,
- hidden_dropout=0.0,
- attention_dropout=0.0,
- num_kv_heads=None,
- alibi=False,
- new_decoder_architecture=False,
- multi_query=True,
- parallel_attn=True,
- bias=False,
- max_position_embeddings=2048,
- rope_theta=10000.0,
- rope_scaling=None,
- bos_token_id=11,
- eos_token_id=11,
- ffn_hidden_size=None,
- activation="gelu",
- **kwargs,
- ):
- self.vocab_size = vocab_size
- # Backward compatibility with n_embed kwarg
- n_embed = kwargs.pop("n_embed", None)
- self.hidden_size = hidden_size if n_embed is None else n_embed
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.layer_norm_epsilon = layer_norm_epsilon
- self.initializer_range = initializer_range
- self.use_cache = use_cache
- self.hidden_dropout = hidden_dropout
- self.attention_dropout = attention_dropout
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
- self.num_kv_heads = num_attention_heads if num_kv_heads is None else num_kv_heads
- self.alibi = alibi
- self.new_decoder_architecture = new_decoder_architecture
- self.multi_query = multi_query # Ignored when new_decoder_architecture is True
- self.parallel_attn = parallel_attn
- self.bias = bias
- self.max_position_embeddings = max_position_embeddings
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
- self.activation = activation
- if ffn_hidden_size is None:
- self.ffn_hidden_size = hidden_size * 4
- else:
- self.ffn_hidden_size = ffn_hidden_size
- self._rope_scaling_validation()
-
- super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
-
- @property
- def head_dim(self):
- return self.hidden_size // self.num_attention_heads
-
- @property
- def rotary(self):
- return not self.alibi
-
- def _rope_scaling_validation(self):
- """
- Validate the `rope_scaling` configuration.
- """
- if self.rope_scaling is None:
- return
-
- if self.alibi:
- raise ValueError("`rope_scaling` is not supported when `alibi` is `True`.")
-
- if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
- raise ValueError(
- "`rope_scaling` must be a dictionary with two fields, `type` and `factor`, " f"got {self.rope_scaling}"
- )
- rope_scaling_type = self.rope_scaling.get("type", None)
- rope_scaling_factor = self.rope_scaling.get("factor", None)
- if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
- raise ValueError(
- f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
- )
- if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
- raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
diff --git a/transformers/models/falcon/convert_custom_code_checkpoint.py b/transformers/models/falcon/convert_custom_code_checkpoint.py
deleted file mode 100644
index 0da817c3ffa73907c0215be12377f08fb5729a85..0000000000000000000000000000000000000000
--- a/transformers/models/falcon/convert_custom_code_checkpoint.py
+++ /dev/null
@@ -1,74 +0,0 @@
-import json
-from argparse import ArgumentParser
-from pathlib import Path
-
-
-"""
-This script converts Falcon custom code checkpoints to modern Falcon checkpoints that use code in the Transformers
-library. After conversion, performance (especially for generation) should improve and the checkpoint can be loaded
-without needing trust_remote_code=True.
-"""
-
-if __name__ == "__main__":
- parser = ArgumentParser()
- parser.add_argument(
- "--checkpoint_dir",
- type=Path,
- required=True,
- help="Directory containing a custom code checkpoint to convert to a modern Falcon checkpoint.",
- )
- args = parser.parse_args()
-
- if not args.checkpoint_dir.is_dir():
- raise ValueError("--checkpoint_dir argument should be a directory!")
-
- if (
- not (args.checkpoint_dir / "configuration_RW.py").is_file()
- or not (args.checkpoint_dir / "modelling_RW.py").is_file()
- ):
- raise ValueError(
- "The model directory should contain configuration_RW.py and modelling_RW.py files! Are you sure this is a custom code checkpoint?"
- )
- (args.checkpoint_dir / "configuration_RW.py").unlink()
- (args.checkpoint_dir / "modelling_RW.py").unlink()
-
- config = args.checkpoint_dir / "config.json"
- text = config.read_text()
- text = text.replace("RWForCausalLM", "FalconForCausalLM")
- text = text.replace("RefinedWebModel", "falcon")
- text = text.replace("RefinedWeb", "falcon")
- json_config = json.loads(text)
- del json_config["auto_map"]
-
- if "n_head" in json_config:
- json_config["num_attention_heads"] = json_config.pop("n_head")
- if "n_layer" in json_config:
- json_config["num_hidden_layers"] = json_config.pop("n_layer")
- if "n_head_kv" in json_config:
- json_config["num_kv_heads"] = json_config.pop("n_head_kv")
- json_config["new_decoder_architecture"] = True
- else:
- json_config["new_decoder_architecture"] = False
- bos_token_id = json_config.get("bos_token_id", 1)
- eos_token_id = json_config.get("eos_token_id", 2)
- config.unlink()
- config.write_text(json.dumps(json_config, indent=2, sort_keys=True))
-
- tokenizer_config = args.checkpoint_dir / "tokenizer_config.json"
- if tokenizer_config.is_file():
- text = tokenizer_config.read_text()
- json_config = json.loads(text)
- if json_config["tokenizer_class"] == "PreTrainedTokenizerFast":
- json_config["model_input_names"] = ["input_ids", "attention_mask"]
- tokenizer_config.unlink()
- tokenizer_config.write_text(json.dumps(json_config, indent=2, sort_keys=True))
-
- generation_config_path = args.checkpoint_dir / "generation_config.json"
- generation_dict = {
- "_from_model_config": True,
- "bos_token_id": bos_token_id,
- "eos_token_id": eos_token_id,
- "transformers_version": "4.33.0.dev0",
- }
- generation_config_path.write_text(json.dumps(generation_dict, indent=2, sort_keys=True))
- print("Done! Please double-check that the new checkpoint works as expected.")
diff --git a/transformers/models/falcon/modeling_falcon.py b/transformers/models/falcon/modeling_falcon.py
deleted file mode 100644
index d9254bec0a7342d10388f5954c54a3b16d747aaa..0000000000000000000000000000000000000000
--- a/transformers/models/falcon/modeling_falcon.py
+++ /dev/null
@@ -1,1640 +0,0 @@
-# coding=utf-8
-# Copyright 2023 the Falcon authors and HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch Falcon model."""
-
-import math
-import warnings
-from typing import TYPE_CHECKING, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
-from torch.nn import functional as F
-
-from ...activations import get_activation
-from ...modeling_attn_mask_utils import (
- AttentionMaskConverter,
- _prepare_4d_causal_attention_mask,
- _prepare_4d_causal_attention_mask_for_sdpa,
-)
-from ...modeling_outputs import (
- BaseModelOutputWithPastAndCrossAttentions,
- CausalLMOutputWithCrossAttentions,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutputWithPast,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import is_torch_greater_or_equal_than_2_0
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- logging,
-)
-from .configuration_falcon import FalconConfig
-
-
-if TYPE_CHECKING:
- from ...configuration_utils import PretrainedConfig
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-logger = logging.get_logger(__name__)
-
-from ..deprecated._archive_maps import FALCON_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-_CHECKPOINT_FOR_DOC = "Rocketknight1/falcon-rw-1b"
-_CONFIG_FOR_DOC = "FalconConfig"
-
-
-# NOTE(Hesslow): Unfortunately we did not fuse matmul and bias during training, this means that there's one additional quantization to bfloat16 between the operations.
-# In order not to degrade the quality of our HF-port, we keep these characteristics in the final model.
-class FalconLinear(nn.Linear):
- def forward(self, input: torch.Tensor) -> torch.Tensor:
- hidden_states = input @ self.weight.T
- if self.bias is None:
- return hidden_states
- return hidden_states + self.bias
-
-
-# Copied from transformers.models.llama.modeling_llama.rotate_half
-def rotate_half(x):
- """Rotates half the hidden dims of the input."""
- x1 = x[..., : x.shape[-1] // 2]
- x2 = x[..., x.shape[-1] // 2 :]
- return torch.cat((-x2, x1), dim=-1)
-
-
-# Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
- """Applies Rotary Position Embedding to the query and key tensors.
-
- Args:
- q (`torch.Tensor`): The query tensor.
- k (`torch.Tensor`): The key tensor.
- cos (`torch.Tensor`): The cosine part of the rotary embedding.
- sin (`torch.Tensor`): The sine part of the rotary embedding.
- position_ids (`torch.Tensor`):
- The position indices of the tokens corresponding to the query and key tensors. For example, this can be
- used to pass offsetted position ids when working with a KV-cache.
- unsqueeze_dim (`int`, *optional*, defaults to 1):
- The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
- sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
- that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
- k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
- cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
- the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
- Returns:
- `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
- """
- cos = cos[position_ids].unsqueeze(unsqueeze_dim)
- sin = sin[position_ids].unsqueeze(unsqueeze_dim)
- q_embed = (q * cos) + (rotate_half(q) * sin)
- k_embed = (k * cos) + (rotate_half(k) * sin)
- return q_embed, k_embed
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-# Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->Falcon
-class FalconRotaryEmbedding(nn.Module):
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
- super().__init__()
-
- self.dim = dim
- self.max_position_embeddings = max_position_embeddings
- self.base = base
- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
-
- # Build here to make `torch.jit.trace` work.
- self._set_cos_sin_cache(
- seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
- )
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
-
- def forward(self, x, seq_len=None):
- # x: [bs, num_attention_heads, seq_len, head_size]
- if seq_len > self.max_seq_len_cached:
- self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-
- return (
- self.cos_cached[:seq_len].to(dtype=x.dtype),
- self.sin_cached[:seq_len].to(dtype=x.dtype),
- )
-
-
-# copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->Falcon
-# TODO @joao no longer copied from LLama after static cache, fix me (copied -> Copied)
-class FalconLinearScalingRotaryEmbedding(FalconRotaryEmbedding):
- """FalconRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
-
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
- self.scaling_factor = scaling_factor
- super().__init__(dim, max_position_embeddings, base, device)
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
- t = t / self.scaling_factor
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
-
-
-# copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->Falcon
-# TODO @joao no longer copied from LLama after static cache, fix me (copied -> Copied)
-class FalconDynamicNTKScalingRotaryEmbedding(FalconRotaryEmbedding):
- """FalconRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
-
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
- self.scaling_factor = scaling_factor
- super().__init__(dim, max_position_embeddings, base, device)
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
-
- if seq_len > self.max_position_embeddings:
- base = self.base * (
- (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
- ) ** (self.dim / (self.dim - 2))
- inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
-
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
- self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
-
-
-def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
- batch_size, seq_length = attention_mask.shape
- closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
- base = torch.tensor(
- 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
- )
- powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)
- slopes = torch.pow(base, powers)
-
- if closest_power_of_2 != num_heads:
- extra_base = torch.tensor(
- 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
- )
- num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
- extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)
- slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
-
- # Note: alibi will added to the attention bias that will be applied to the query, key product of attention
- # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length)
- # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length)
- # => the query_length dimension will then be broadcasted correctly
- # This is more or less identical to T5's relative position bias:
- # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527
- arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :]
- alibi = slopes[..., None].bfloat16() * arange_tensor
- return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype)
-
-
-# Copied from transformers.models.bloom.modeling_bloom.dropout_add
-def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:
- """
- Dropout add function
-
- Args:
- x (`torch.tensor`, *required*):
- input tensor
- residual (`torch.tensor`, *required*):
- residual tensor
- prob (`float`, *required*):
- dropout probability
- training (`bool`, *required*):
- training mode
- """
- out = F.dropout(x, p=prob, training=training)
- out = residual + out
- return out
-
-
-class FalconAttention(nn.Module):
- def __init__(self, config: FalconConfig):
- super().__init__()
-
- self.config = config
- self.hidden_size = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.hidden_size // self.num_heads
- self.split_size = self.hidden_size
- self.hidden_dropout = config.hidden_dropout
- self.max_position_embeddings = config.max_position_embeddings
- self.rope_theta = config.rope_theta
- self.is_causal = True
- self._use_sdpa = config._attn_implementation == "sdpa"
-
- if self.head_dim * self.num_heads != self.hidden_size:
- raise ValueError(
- f"`hidden_size` must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"
- f" {self.num_heads})."
- )
-
- if config.rotary:
- self._init_rope()
-
- # Layer-wise attention scaling
- self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
- self.beta = self.inv_norm_factor
- if config.new_decoder_architecture:
- qkv_out_dim = (config.num_kv_heads * 2 + config.num_attention_heads) * self.head_dim
- elif config.multi_query:
- qkv_out_dim = self.hidden_size + 2 * self.head_dim
- else:
- qkv_out_dim = 3 * self.hidden_size
- self.query_key_value = FalconLinear(self.hidden_size, qkv_out_dim, bias=config.bias)
- self.new_decoder_architecture = config.new_decoder_architecture
- self.multi_query = config.multi_query
- self.dense = FalconLinear(self.hidden_size, self.hidden_size, bias=config.bias)
- self.attention_dropout = nn.Dropout(config.attention_dropout)
- self.num_kv_heads = config.num_kv_heads if (self.new_decoder_architecture or not self.multi_query) else 1
-
- # Copied from transformers.models.llama.modeling_llama.LlamaAttention._init_rope with Llama->Falcon
- def _init_rope(self):
- if self.config.rope_scaling is None:
- self.rotary_emb = FalconRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- base=self.rope_theta,
- )
- else:
- scaling_type = self.config.rope_scaling["type"]
- scaling_factor = self.config.rope_scaling["factor"]
- if scaling_type == "linear":
- self.rotary_emb = FalconLinearScalingRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- scaling_factor=scaling_factor,
- base=self.rope_theta,
- )
- elif scaling_type == "dynamic":
- self.rotary_emb = FalconDynamicNTKScalingRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- scaling_factor=scaling_factor,
- base=self.rope_theta,
- )
- else:
- raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
-
- def _split_heads(self, fused_qkv: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- """
- Split the last dimension into (num_heads, head_dim), results share same memory storage as `fused_qkv`
-
- Args:
- fused_qkv (`torch.tensor`, *required*): [batch_size, seq_length, num_heads * 3 * head_dim]
-
- Returns:
- query: [batch_size, seq_length, num_heads, head_dim] key: [batch_size, seq_length, num_heads, head_dim]
- value: [batch_size, seq_length, num_heads, head_dim]
- """
- if self.new_decoder_architecture:
- batch, seq_len, _ = fused_qkv.shape
- qkv = fused_qkv.view(batch, seq_len, -1, self.num_heads // self.num_kv_heads + 2, self.head_dim)
- query = qkv[:, :, :, :-2]
- key = qkv[:, :, :, [-2]]
- value = qkv[:, :, :, [-1]]
- key = torch.broadcast_to(key, query.shape)
- value = torch.broadcast_to(value, query.shape)
-
- query, key, value = [x.flatten(2, 3) for x in (query, key, value)]
- return query, key, value
- elif not self.multi_query:
- batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
- fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads, 3, self.head_dim)
- return fused_qkv[..., 0, :], fused_qkv[..., 1, :], fused_qkv[..., 2, :]
- else:
- batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
- fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads + 2, self.head_dim)
- return fused_qkv[..., :-2, :], fused_qkv[..., [-2], :], fused_qkv[..., [-1], :]
-
- # Copied from transformers.models.bloom.modeling_bloom.BloomAttention._merge_heads
- def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
- """
- Merge heads together over the last dimension
-
- Args:
- x (`torch.tensor`, *required*): [batch_size * num_heads, seq_length, head_dim]
-
- Returns:
- torch.tensor: [batch_size, seq_length, num_heads * head_dim]
- """
- # What we want to achieve is:
- # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim
- batch_size_and_num_heads, seq_length, _ = x.shape
- batch_size = batch_size_and_num_heads // self.num_heads
-
- # First view to decompose the batch size
- # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim
- x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)
-
- # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim
- x = x.permute(0, 2, 1, 3)
-
- # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim
- return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- alibi: Optional[torch.Tensor],
- attention_mask: torch.Tensor,
- position_ids: Optional[torch.LongTensor] = None,
- layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- use_cache: bool = False,
- output_attentions: bool = False,
- **kwargs,
- ):
- if "padding_mask" in kwargs:
- warnings.warn(
- "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
- )
-
- fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]
- num_kv_heads = self.num_heads if self.new_decoder_architecture else self.num_kv_heads
- # 3 x [batch_size, seq_length, num_heads, head_dim]
- (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)
-
- batch_size, query_length, _, _ = query_layer.shape
-
- query_layer = query_layer.transpose(1, 2).reshape(batch_size, self.num_heads, query_length, self.head_dim)
- key_layer = key_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
- value_layer = value_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
-
- kv_seq_len = key_layer.shape[-2]
- if layer_past is not None:
- kv_seq_len += layer_past[0].shape[-2]
- if alibi is None:
- cos, sin = self.rotary_emb(value_layer, seq_len=kv_seq_len)
- query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin, position_ids)
-
- if layer_past is not None:
- past_key, past_value = layer_past
- # concatenate along seq_length dimension:
- # - key: [batch_size, self.num_heads, kv_length, head_dim]
- # - value: [batch_size, self.num_heads, kv_length, head_dim]
- key_layer = torch.cat((past_key, key_layer), dim=-2)
- value_layer = torch.cat((past_value, value_layer), dim=-2)
-
- kv_length = key_layer.shape[-2]
- if use_cache:
- present = (key_layer, value_layer)
- else:
- present = None
-
- if self._use_sdpa and query_layer.device.type == "cuda" and attention_mask is not None:
- # For torch<=2.1.2, SDPA with memory-efficient backend is bugged with non-contiguous inputs with custom attn_mask,
- # Reference: https://github.com/pytorch/pytorch/issues/112577.
- query_layer = query_layer.contiguous()
- key_layer = key_layer.contiguous()
- value_layer = value_layer.contiguous()
-
- if alibi is None:
- if self._use_sdpa and not output_attentions:
- attn_output = F.scaled_dot_product_attention(
- query_layer,
- key_layer,
- value_layer,
- attention_mask,
- 0.0,
- # The query_length > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case query_length == 1.
- is_causal=self.is_causal and attention_mask is None and query_length > 1,
- )
-
- attention_scores = None
- else:
- attention_scores = query_layer @ key_layer.transpose(-1, -2)
- attention_scores /= math.sqrt(self.head_dim)
-
- attention_scores = F.softmax(attention_scores + attention_mask, dim=-1, dtype=hidden_states.dtype)
- # It is unclear why neither dropout nor head_mask is applied here (while it is with alibi).
- attn_output = attention_scores @ value_layer
-
- attn_output = attn_output.view(batch_size, self.num_heads, query_length, self.head_dim)
- attn_output = attn_output.permute(0, 2, 1, 3)
- attn_output = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
-
- attn_output = self.dense(attn_output)
-
- if output_attentions:
- return attn_output, present, attention_scores
- else:
- return attn_output, present
-
- else:
- if self._use_sdpa and not output_attentions and head_mask is None:
- attn_output = F.scaled_dot_product_attention(
- query_layer,
- key_layer,
- value_layer,
- attn_mask=attention_mask,
- dropout_p=self.attention_dropout.p if self.training else 0.0,
- is_causal=self.is_causal and attention_mask is None and query_length > 1,
- )
- attn_output = attn_output.transpose(1, 2)
- attn_output = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
-
- attn_output = self.dense(attn_output)
- else:
- matmul_result = query_layer @ key_layer.transpose(-1, -2)
-
- # change view to [batch_size, num_heads, q_length, kv_length]
- attention_scores = matmul_result.view(batch_size, self.num_heads, query_length, kv_length)
-
- # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length]
- input_dtype = attention_scores.dtype
- # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38`
- if input_dtype == torch.float16 or input_dtype == torch.bfloat16:
- attention_scores = attention_scores.to(torch.float32)
-
- attention_logits = attention_scores + alibi.view(batch_size, self.num_heads, 1, -1)
- attention_logits *= self.inv_norm_factor
- attention_probs = F.softmax(attention_logits + attention_mask, dim=-1, dtype=hidden_states.dtype)
- # [batch_size, num_heads, q_length, kv_length]
- attention_probs = self.attention_dropout(attention_probs)
-
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- # change view [batch_size, num_heads, q_length, kv_length]
- attention_probs_reshaped = attention_probs.view(batch_size, self.num_heads, query_length, kv_length)
-
- # matmul: [batch_size * num_heads, q_length, head_dim]
- attn_output = (attention_probs_reshaped @ value_layer).flatten(0, 1)
-
- # change view [batch_size, q_length, num_heads * head_dim]
- attn_output = self._merge_heads(attn_output)
-
- attn_output = self.dense(attn_output)
-
- if output_attentions:
- return attn_output, present, attention_probs
- else:
- return attn_output, present
-
-
-class FalconFlashAttention2(FalconAttention):
- """
- Falcon flash attention module. This module inherits from `FalconAttention` as the weights of the module stays
- untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
- flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- alibi: Optional[torch.Tensor],
- attention_mask: torch.Tensor,
- position_ids: Optional[torch.LongTensor] = None,
- layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- use_cache: bool = False,
- output_attentions: bool = False,
- **kwargs,
- ):
- if "padding_mask" in kwargs:
- warnings.warn(
- "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
- )
-
- # overwrite attention_mask with padding_mask
- attention_mask = kwargs.pop("padding_mask")
-
- fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]
- num_kv_heads = self.num_heads if self.new_decoder_architecture else self.num_kv_heads
- # 3 x [batch_size, seq_length, num_heads, head_dim]
- (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)
-
- batch_size, query_length, _, _ = query_layer.shape
-
- query_layer = query_layer.transpose(1, 2).reshape(batch_size, self.num_heads, query_length, self.head_dim)
- key_layer = key_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
- value_layer = value_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
-
- kv_seq_len = key_layer.shape[-2]
- if layer_past is not None:
- kv_seq_len += layer_past[0].shape[-2]
- if alibi is None:
- cos, sin = self.rotary_emb(value_layer, seq_len=kv_seq_len)
- query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin, position_ids)
-
- if layer_past is not None and use_cache:
- past_key, past_value = layer_past
- # concatenate along seq_length dimension:
- # - key: [batch_size, self.num_heads, kv_length, head_dim]
- # - value: [batch_size, self.num_heads, kv_length, head_dim]
- key_layer = torch.cat((past_key, key_layer), dim=-2)
- value_layer = torch.cat((past_value, value_layer), dim=-2)
-
- past_key_value = (key_layer, value_layer) if use_cache else None
-
- # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
- # to be able to avoid many of these transpose/reshape/view.
- query_layer = query_layer.transpose(1, 2)
- key_layer = key_layer.transpose(1, 2)
- value_layer = value_layer.transpose(1, 2)
-
- if alibi is not None:
- raise ValueError("`alibi` is not supported when `use_flash_attn` is True")
-
- attn_dropout = self.config.attention_dropout if self.training else 0.0
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in float16 just to be sure everything works as expected.
- input_dtype = query_layer.dtype
- if input_dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.query_key_value.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
-
- query_layer = query_layer.to(target_dtype)
- key_layer = key_layer.to(target_dtype)
- value_layer = value_layer.to(target_dtype)
-
- attn_output = self._flash_attention_forward(
- query_layer, key_layer, value_layer, attention_mask, query_length, dropout=attn_dropout
- )
-
- attn_weights = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
- attn_output = self.dense(attn_weights)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, past_key_value, attn_weights
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-class FalconMLP(nn.Module):
- def __init__(self, config: FalconConfig):
- super().__init__()
- hidden_size = config.hidden_size
-
- self.dense_h_to_4h = FalconLinear(hidden_size, config.ffn_hidden_size, bias=config.bias)
- self.act = get_activation(config.activation)
- self.dense_4h_to_h = FalconLinear(config.ffn_hidden_size, hidden_size, bias=config.bias)
- self.hidden_dropout = config.hidden_dropout
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- x = self.act(self.dense_h_to_4h(x))
- x = self.dense_4h_to_h(x)
- return x
-
-
-FALCON_ATTENTION_CLASSES = {
- "eager": FalconAttention,
- "sdpa": FalconAttention, # FalconAttention originally implemented both a forward with & without SDPA
- "flash_attention_2": FalconFlashAttention2,
-}
-
-
-class FalconDecoderLayer(nn.Module):
- def __init__(self, config: FalconConfig):
- super().__init__()
- hidden_size = config.hidden_size
- self.num_heads = config.num_attention_heads
-
- self.self_attention = FALCON_ATTENTION_CLASSES[config._attn_implementation](config)
- self.mlp = FalconMLP(config)
- self.hidden_dropout = config.hidden_dropout
- self.config = config
-
- if config.new_decoder_architecture:
- # The layer norm before self-attention
- self.ln_attn = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
- # The layer norm before the MLP
- self.ln_mlp = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
- else:
- self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
- if not config.parallel_attn:
- self.post_attention_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- alibi: Optional[torch.Tensor],
- attention_mask: torch.Tensor,
- position_ids: Optional[torch.LongTensor] = None,
- layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- use_cache: bool = False,
- output_attentions: bool = False,
- **kwargs,
- ):
- if "padding_mask" in kwargs:
- warnings.warn(
- "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
- )
-
- residual = hidden_states
-
- if self.config.new_decoder_architecture:
- attention_layernorm_out = self.ln_attn(hidden_states)
- mlp_layernorm_out = self.ln_mlp(hidden_states)
- else:
- attention_layernorm_out = self.input_layernorm(hidden_states)
-
- # Self attention.
- attn_outputs = self.self_attention(
- attention_layernorm_out,
- layer_past=layer_past,
- attention_mask=attention_mask,
- position_ids=position_ids,
- alibi=alibi,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- **kwargs,
- )
-
- attention_output = attn_outputs[0]
-
- if not self.config.new_decoder_architecture:
- if self.config.parallel_attn:
- mlp_layernorm_out = attention_layernorm_out
- else:
- residual = dropout_add(
- attention_output, residual, self.config.attention_dropout, training=self.training
- )
- mlp_layernorm_out = self.post_attention_layernorm(residual)
-
- outputs = attn_outputs[1:]
-
- # MLP.
- mlp_output = self.mlp(mlp_layernorm_out)
-
- if self.config.new_decoder_architecture or self.config.parallel_attn:
- mlp_output += attention_output
-
- output = dropout_add(mlp_output, residual, self.config.hidden_dropout, training=self.training)
-
- if use_cache:
- outputs = (output,) + outputs
- else:
- outputs = (output,) + outputs[1:]
-
- return outputs # hidden_states, present, attentions
-
-
-FALCON_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`FalconConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FALCON_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0][0].shape[2]`
- (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
-
- If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
- `input_ids`.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.num_hidden_layers`):
- Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
- their past given to this model should not be passed as `input_ids` as they have already been computed.
-
- Each element of `past_key_values` is a tuple (past_key, past_value):
- - past_key: [batch_size * num_heads, head_dim, kv_length]
- - past_value: [batch_size * num_heads, kv_length, head_dim]
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
-
- If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see
- `past_key_values`).
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class FalconPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = FalconConfig
- base_model_prefix = "transformer"
- supports_gradient_checkpointing = True
- _no_split_modules = ["FalconDecoderLayer"]
- _supports_flash_attn_2 = True
- _supports_sdpa = True
-
- def __init__(self, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- def _init_weights(self, module: nn.Module):
- """Initialize the weights."""
- if isinstance(module, nn.Linear) or isinstance(module, FalconLinear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
- # Adapted from transformers.modeling_utils.PreTrainedModel._check_and_enable_sdpa
- @classmethod
- def _check_and_enable_sdpa(cls, config, hard_check_only: bool = False) -> "PretrainedConfig":
- # NOTE: Falcon supported SDPA from PyTorch 2.0. We keep it like that for backward compatibility (automatically use SDPA for torch>=2.0).
- if hard_check_only:
- if not is_torch_greater_or_equal_than_2_0:
- raise ImportError("PyTorch SDPA requirements in Transformers are not met. Please install torch>=2.0.")
-
- if not is_torch_greater_or_equal_than_2_0:
- return config
-
- _is_bettertransformer = getattr(cls, "use_bettertransformer", False)
- if _is_bettertransformer:
- return config
-
- if not hard_check_only:
- config._attn_implementation = "sdpa"
- return config
-
-
-@add_start_docstrings(
- "The bare Falcon Model transformer outputting raw hidden-states without any specific head on top.",
- FALCON_START_DOCSTRING,
-)
-class FalconModel(FalconPreTrainedModel):
- def __init__(self, config: FalconConfig):
- super().__init__(config)
-
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.use_alibi = config.alibi
-
- # Embedding + LN Embedding
- self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)
-
- # Transformer blocks
- self.h = nn.ModuleList([FalconDecoderLayer(config) for _ in range(config.num_hidden_layers)])
- self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
- self._use_sdpa = config._attn_implementation == "sdpa"
-
- # Final Layer Norm
- self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
-
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.word_embeddings
-
- def set_input_embeddings(self, new_embeddings: torch.Tensor):
- self.word_embeddings = new_embeddings
-
- @add_start_docstrings_to_model_forward(FALCON_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPastAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- batch_size, seq_length = input_ids.shape
- elif inputs_embeds is not None:
- batch_size, seq_length, _ = inputs_embeds.shape
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if past_key_values is None:
- past_key_values = tuple([None] * len(self.h))
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
-
- hidden_states = inputs_embeds
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
- presents = () if use_cache else None
- all_self_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
-
- # Compute alibi tensor: check build_alibi_tensor documentation
- past_key_values_length = 0
- if past_key_values[0] is not None:
- past_key_values_length = past_key_values[0][0].shape[-2]
-
- if self.use_alibi:
- mask = (
- torch.ones(
- (batch_size, seq_length + past_key_values_length), device=inputs_embeds.device, dtype=torch.long
- )
- if attention_mask is None
- else attention_mask
- )
- alibi = build_alibi_tensor(mask, self.num_heads, dtype=hidden_states.dtype)
- else:
- alibi = None
- if position_ids is None:
- device = input_ids.device if input_ids is not None else inputs_embeds.device
- position_ids = torch.arange(
- past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
- )
- position_ids = position_ids.unsqueeze(0)
-
- if self._use_flash_attention_2:
- # 2d mask is passed through the layers
- attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
- elif self._use_sdpa and not output_attentions:
- # output_attentions=True can not be supported when using SDPA, and we fall back on
- # the manual implementation that requires a 4D causal mask in all cases.
- if alibi is None:
- attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
- attention_mask,
- (batch_size, seq_length),
- inputs_embeds,
- past_key_values_length,
- )
- elif head_mask is None:
- alibi = alibi.reshape(batch_size, -1, *alibi.shape[1:])
-
- # We don't call _prepare_4d_causal_attention_mask_for_sdpa as we need to mask alibi using the 4D attention_mask untouched.
- attention_mask = _prepare_4d_causal_attention_mask(
- attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
- )
-
- # We take care to integrate alibi bias in the attention_mask here.
- min_dtype = torch.finfo(alibi.dtype).min
- attention_mask = torch.masked_fill(
- alibi / math.sqrt(self.config.hidden_size // self.num_heads),
- attention_mask < -1,
- min_dtype,
- )
-
- # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
- # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
- if seq_length > 1 and attention_mask.device.type == "cuda":
- attention_mask = AttentionMaskConverter._unmask_unattended(attention_mask, min_dtype=min_dtype)
- else:
- # PyTorch SDPA does not support head_mask, we fall back on the eager implementation in this case.
- attention_mask = _prepare_4d_causal_attention_mask(
- attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
- )
- else:
- # 4d mask is passed through the layers
- attention_mask = _prepare_4d_causal_attention_mask(
- attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
- )
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape batch_size x num_heads x N x N
- # head_mask has shape n_layer x batch x num_heads x N x N
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- alibi,
- attention_mask,
- position_ids,
- head_mask[i],
- layer_past,
- use_cache,
- output_attentions,
- )
- else:
- outputs = block(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask[i],
- use_cache=use_cache,
- output_attentions=output_attentions,
- alibi=alibi,
- )
-
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
-
- # Add last hidden state
- hidden_states = self.ln_f(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
-
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-@add_start_docstrings(
- "The Falcon Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings).",
- FALCON_START_DOCSTRING,
-)
-class FalconForCausalLM(FalconPreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config: FalconConfig):
- super().__init__(config)
- self.transformer = FalconModel(config)
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings: torch.Tensor):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(
- self,
- input_ids: torch.LongTensor,
- past_key_values: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- **kwargs,
- ) -> dict:
- if past_key_values is not None:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
-
- # Note: versions of Falcon with alibi do not use position_ids. It is used with RoPE.
- if not self.transformer.use_alibi and attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- return {
- "input_ids": input_ids,
- "position_ids": position_ids,
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "attention_mask": attention_mask,
- }
-
- @add_start_docstrings_to_model_forward(FALCON_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutputWithCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
-
- lm_logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- batch_size, seq_length, vocab_size = shift_logits.shape
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(
- shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length)
- )
-
- if not return_dict:
- output = (lm_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithCrossAttentions(
- loss=loss,
- logits=lm_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def _reorder_cache(
- self, past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
- ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
- """
- This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
- [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
- beam_idx at every generation step.
-
- Output shares the same memory storage as `past`.
- """
-
- # Get a copy of `beam_idx` on all the devices where we need those indices.
- device_to_beam_idx = {
- past_state.device: beam_idx.to(past_state.device) for layer_past in past for past_state in layer_past
- }
- reordered_past = tuple(
- (
- layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]),
- layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]),
- )
- for layer_past in past
- )
- return reordered_past
-
-
-@add_start_docstrings(
- """
- The Falcon Model transformer with a sequence classification head on top (linear layer).
-
- [`FalconForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT-1) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- FALCON_START_DOCSTRING,
-)
-class FalconForSequenceClassification(FalconPreTrainedModel):
- def __init__(self, config: FalconConfig):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = FalconModel(config)
- self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FALCON_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size = input_ids.shape[0]
- else:
- batch_size = inputs_embeds.shape[0]
-
- if self.config.pad_token_id is None and batch_size != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
-
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Falcon Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- FALCON_START_DOCSTRING,
-)
-class FalconForTokenClassification(FalconPreTrainedModel):
- def __init__(self, config: FalconConfig):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.transformer = FalconModel(config)
- if getattr(config, "classifier_dropout", None) is not None:
- classifier_dropout = config.classifier_dropout
- elif getattr(config, "hidden_dropout", None) is not None:
- classifier_dropout = config.hidden_dropout
- else:
- classifier_dropout = 0.1
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FALCON_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = transformer_outputs[0]
- hidden_states = self.dropout(hidden_states)
- logits = self.classifier(hidden_states)
-
- loss = None
- if labels is not None:
- batch_size, seq_length = labels.shape
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(
- logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length)
- )
-
- if not return_dict:
- output = (logits,) + transformer_outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- The Falcon Model transformer with a span classification head on top for extractive question-answering tasks like
- SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- FALCON_START_DOCSTRING,
-)
-class FalconForQuestionAnswering(FalconPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.transformer = FalconModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, 2)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FALCON_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- start_positions: Optional[torch.LongTensor] = None,
- end_positions: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/fastspeech2_conformer/__init__.py b/transformers/models/fastspeech2_conformer/__init__.py
deleted file mode 100644
index 1fd5cbf1dc272ec4d9747e2dd1616ed37d5acdbf..0000000000000000000000000000000000000000
--- a/transformers/models/fastspeech2_conformer/__init__.py
+++ /dev/null
@@ -1,77 +0,0 @@
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_fastspeech2_conformer": [
- "FASTSPEECH2_CONFORMER_HIFIGAN_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "FASTSPEECH2_CONFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "FASTSPEECH2_CONFORMER_WITH_HIFIGAN_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "FastSpeech2ConformerConfig",
- "FastSpeech2ConformerHifiGanConfig",
- "FastSpeech2ConformerWithHifiGanConfig",
- ],
- "tokenization_fastspeech2_conformer": ["FastSpeech2ConformerTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_fastspeech2_conformer"] = [
- "FASTSPEECH2_CONFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
- "FastSpeech2ConformerWithHifiGan",
- "FastSpeech2ConformerHifiGan",
- "FastSpeech2ConformerModel",
- "FastSpeech2ConformerPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_fastspeech2_conformer import (
- FASTSPEECH2_CONFORMER_HIFIGAN_PRETRAINED_CONFIG_ARCHIVE_MAP,
- FASTSPEECH2_CONFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
- FASTSPEECH2_CONFORMER_WITH_HIFIGAN_PRETRAINED_CONFIG_ARCHIVE_MAP,
- FastSpeech2ConformerConfig,
- FastSpeech2ConformerHifiGanConfig,
- FastSpeech2ConformerWithHifiGanConfig,
- )
- from .tokenization_fastspeech2_conformer import FastSpeech2ConformerTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_fastspeech2_conformer import (
- FASTSPEECH2_CONFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
- FastSpeech2ConformerHifiGan,
- FastSpeech2ConformerModel,
- FastSpeech2ConformerPreTrainedModel,
- FastSpeech2ConformerWithHifiGan,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/fastspeech2_conformer/__pycache__/__init__.cpython-310.pyc b/transformers/models/fastspeech2_conformer/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 939314f650286e799303b7787cb3961ad9956a23..0000000000000000000000000000000000000000
Binary files a/transformers/models/fastspeech2_conformer/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fastspeech2_conformer/__pycache__/configuration_fastspeech2_conformer.cpython-310.pyc b/transformers/models/fastspeech2_conformer/__pycache__/configuration_fastspeech2_conformer.cpython-310.pyc
deleted file mode 100644
index 3b5280dd3e53778bea7b50e535def4204ac1a826..0000000000000000000000000000000000000000
Binary files a/transformers/models/fastspeech2_conformer/__pycache__/configuration_fastspeech2_conformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fastspeech2_conformer/__pycache__/convert_fastspeech2_conformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/fastspeech2_conformer/__pycache__/convert_fastspeech2_conformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 6f2e12f21ce95e75ee1a263ed84f559cb0765ede..0000000000000000000000000000000000000000
Binary files a/transformers/models/fastspeech2_conformer/__pycache__/convert_fastspeech2_conformer_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fastspeech2_conformer/__pycache__/convert_hifigan.cpython-310.pyc b/transformers/models/fastspeech2_conformer/__pycache__/convert_hifigan.cpython-310.pyc
deleted file mode 100644
index 2ebd63e00944b1554b00beed9518fd3752696ba7..0000000000000000000000000000000000000000
Binary files a/transformers/models/fastspeech2_conformer/__pycache__/convert_hifigan.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fastspeech2_conformer/__pycache__/convert_model_with_hifigan.cpython-310.pyc b/transformers/models/fastspeech2_conformer/__pycache__/convert_model_with_hifigan.cpython-310.pyc
deleted file mode 100644
index 2778dd8fbaad39af4bca6bbcc43afbbd615ac1be..0000000000000000000000000000000000000000
Binary files a/transformers/models/fastspeech2_conformer/__pycache__/convert_model_with_hifigan.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fastspeech2_conformer/__pycache__/modeling_fastspeech2_conformer.cpython-310.pyc b/transformers/models/fastspeech2_conformer/__pycache__/modeling_fastspeech2_conformer.cpython-310.pyc
deleted file mode 100644
index f854cc16e35172d8b21414ab6988fb6c66965a55..0000000000000000000000000000000000000000
Binary files a/transformers/models/fastspeech2_conformer/__pycache__/modeling_fastspeech2_conformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fastspeech2_conformer/__pycache__/tokenization_fastspeech2_conformer.cpython-310.pyc b/transformers/models/fastspeech2_conformer/__pycache__/tokenization_fastspeech2_conformer.cpython-310.pyc
deleted file mode 100644
index 72cfb6c4510656dd63925ba112cd2501b521e82a..0000000000000000000000000000000000000000
Binary files a/transformers/models/fastspeech2_conformer/__pycache__/tokenization_fastspeech2_conformer.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fastspeech2_conformer/configuration_fastspeech2_conformer.py b/transformers/models/fastspeech2_conformer/configuration_fastspeech2_conformer.py
deleted file mode 100644
index adb038ad1b2a0b2b7072f76f2a9fc714b5937a0d..0000000000000000000000000000000000000000
--- a/transformers/models/fastspeech2_conformer/configuration_fastspeech2_conformer.py
+++ /dev/null
@@ -1,482 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" FastSpeech2Conformer model configuration"""
-
-from typing import Dict
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import ( # noqa: F401, E402
- FASTSPEECH2_CONFORMER_HIFIGAN_PRETRAINED_CONFIG_ARCHIVE_MAP, # noqa: F401, E402
- FASTSPEECH2_CONFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, # noqa: F401, E402
- FASTSPEECH2_CONFORMER_WITH_HIFIGAN_PRETRAINED_CONFIG_ARCHIVE_MAP, # noqa: F401, E402
-)
-
-
-class FastSpeech2ConformerConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FastSpeech2ConformerModel`]. It is used to
- instantiate a FastSpeech2Conformer model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the
- FastSpeech2Conformer [espnet/fastspeech2_conformer](https://huggingface.co/espnet/fastspeech2_conformer)
- architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- hidden_size (`int`, *optional*, defaults to 384):
- The dimensionality of the hidden layers.
- vocab_size (`int`, *optional*, defaults to 78):
- The size of the vocabulary.
- num_mel_bins (`int`, *optional*, defaults to 80):
- The number of mel filters used in the filter bank.
- encoder_num_attention_heads (`int`, *optional*, defaults to 2):
- The number of attention heads in the encoder.
- encoder_layers (`int`, *optional*, defaults to 4):
- The number of layers in the encoder.
- encoder_linear_units (`int`, *optional*, defaults to 1536):
- The number of units in the linear layer of the encoder.
- decoder_layers (`int`, *optional*, defaults to 4):
- The number of layers in the decoder.
- decoder_num_attention_heads (`int`, *optional*, defaults to 2):
- The number of attention heads in the decoder.
- decoder_linear_units (`int`, *optional*, defaults to 1536):
- The number of units in the linear layer of the decoder.
- speech_decoder_postnet_layers (`int`, *optional*, defaults to 5):
- The number of layers in the post-net of the speech decoder.
- speech_decoder_postnet_units (`int`, *optional*, defaults to 256):
- The number of units in the post-net layers of the speech decoder.
- speech_decoder_postnet_kernel (`int`, *optional*, defaults to 5):
- The kernel size in the post-net of the speech decoder.
- positionwise_conv_kernel_size (`int`, *optional*, defaults to 3):
- The size of the convolution kernel used in the position-wise layer.
- encoder_normalize_before (`bool`, *optional*, defaults to `False`):
- Specifies whether to normalize before encoder layers.
- decoder_normalize_before (`bool`, *optional*, defaults to `False`):
- Specifies whether to normalize before decoder layers.
- encoder_concat_after (`bool`, *optional*, defaults to `False`):
- Specifies whether to concatenate after encoder layers.
- decoder_concat_after (`bool`, *optional*, defaults to `False`):
- Specifies whether to concatenate after decoder layers.
- reduction_factor (`int`, *optional*, defaults to 1):
- The factor by which the speech frame rate is reduced.
- speaking_speed (`float`, *optional*, defaults to 1.0):
- The speed of the speech produced.
- use_macaron_style_in_conformer (`bool`, *optional*, defaults to `True`):
- Specifies whether to use macaron style in the conformer.
- use_cnn_in_conformer (`bool`, *optional*, defaults to `True`):
- Specifies whether to use convolutional neural networks in the conformer.
- encoder_kernel_size (`int`, *optional*, defaults to 7):
- The kernel size used in the encoder.
- decoder_kernel_size (`int`, *optional*, defaults to 31):
- The kernel size used in the decoder.
- duration_predictor_layers (`int`, *optional*, defaults to 2):
- The number of layers in the duration predictor.
- duration_predictor_channels (`int`, *optional*, defaults to 256):
- The number of channels in the duration predictor.
- duration_predictor_kernel_size (`int`, *optional*, defaults to 3):
- The kernel size used in the duration predictor.
- energy_predictor_layers (`int`, *optional*, defaults to 2):
- The number of layers in the energy predictor.
- energy_predictor_channels (`int`, *optional*, defaults to 256):
- The number of channels in the energy predictor.
- energy_predictor_kernel_size (`int`, *optional*, defaults to 3):
- The kernel size used in the energy predictor.
- energy_predictor_dropout (`float`, *optional*, defaults to 0.5):
- The dropout rate in the energy predictor.
- energy_embed_kernel_size (`int`, *optional*, defaults to 1):
- The kernel size used in the energy embed layer.
- energy_embed_dropout (`float`, *optional*, defaults to 0.0):
- The dropout rate in the energy embed layer.
- stop_gradient_from_energy_predictor (`bool`, *optional*, defaults to `False`):
- Specifies whether to stop gradients from the energy predictor.
- pitch_predictor_layers (`int`, *optional*, defaults to 5):
- The number of layers in the pitch predictor.
- pitch_predictor_channels (`int`, *optional*, defaults to 256):
- The number of channels in the pitch predictor.
- pitch_predictor_kernel_size (`int`, *optional*, defaults to 5):
- The kernel size used in the pitch predictor.
- pitch_predictor_dropout (`float`, *optional*, defaults to 0.5):
- The dropout rate in the pitch predictor.
- pitch_embed_kernel_size (`int`, *optional*, defaults to 1):
- The kernel size used in the pitch embed layer.
- pitch_embed_dropout (`float`, *optional*, defaults to 0.0):
- The dropout rate in the pitch embed layer.
- stop_gradient_from_pitch_predictor (`bool`, *optional*, defaults to `True`):
- Specifies whether to stop gradients from the pitch predictor.
- encoder_dropout_rate (`float`, *optional*, defaults to 0.2):
- The dropout rate in the encoder.
- encoder_positional_dropout_rate (`float`, *optional*, defaults to 0.2):
- The positional dropout rate in the encoder.
- encoder_attention_dropout_rate (`float`, *optional*, defaults to 0.2):
- The attention dropout rate in the encoder.
- decoder_dropout_rate (`float`, *optional*, defaults to 0.2):
- The dropout rate in the decoder.
- decoder_positional_dropout_rate (`float`, *optional*, defaults to 0.2):
- The positional dropout rate in the decoder.
- decoder_attention_dropout_rate (`float`, *optional*, defaults to 0.2):
- The attention dropout rate in the decoder.
- duration_predictor_dropout_rate (`float`, *optional*, defaults to 0.2):
- The dropout rate in the duration predictor.
- speech_decoder_postnet_dropout (`float`, *optional*, defaults to 0.5):
- The dropout rate in the speech decoder postnet.
- max_source_positions (`int`, *optional*, defaults to 5000):
- if `"relative"` position embeddings are used, defines the maximum source input positions.
- use_masking (`bool`, *optional*, defaults to `True`):
- Specifies whether to use masking in the model.
- use_weighted_masking (`bool`, *optional*, defaults to `False`):
- Specifies whether to use weighted masking in the model.
- num_speakers (`int`, *optional*):
- Number of speakers. If set to > 1, assume that the speaker ids will be provided as the input and use
- speaker id embedding layer.
- num_languages (`int`, *optional*):
- Number of languages. If set to > 1, assume that the language ids will be provided as the input and use the
- languge id embedding layer.
- speaker_embed_dim (`int`, *optional*):
- Speaker embedding dimension. If set to > 0, assume that speaker_embedding will be provided as the input.
- is_encoder_decoder (`bool`, *optional*, defaults to `True`):
- Specifies whether the model is an encoder-decoder.
-
- Example:
-
- ```python
- >>> from transformers import FastSpeech2ConformerModel, FastSpeech2ConformerConfig
-
- >>> # Initializing a FastSpeech2Conformer style configuration
- >>> configuration = FastSpeech2ConformerConfig()
-
- >>> # Initializing a model from the FastSpeech2Conformer style configuration
- >>> model = FastSpeech2ConformerModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "fastspeech2_conformer"
- attribute_map = {"num_hidden_layers": "encoder_layers", "num_attention_heads": "encoder_num_attention_heads"}
-
- def __init__(
- self,
- hidden_size=384,
- vocab_size=78,
- num_mel_bins=80,
- encoder_num_attention_heads=2,
- encoder_layers=4,
- encoder_linear_units=1536,
- decoder_layers=4,
- decoder_num_attention_heads=2,
- decoder_linear_units=1536,
- speech_decoder_postnet_layers=5,
- speech_decoder_postnet_units=256,
- speech_decoder_postnet_kernel=5,
- positionwise_conv_kernel_size=3,
- encoder_normalize_before=False,
- decoder_normalize_before=False,
- encoder_concat_after=False,
- decoder_concat_after=False,
- reduction_factor=1,
- speaking_speed=1.0,
- use_macaron_style_in_conformer=True,
- use_cnn_in_conformer=True,
- encoder_kernel_size=7,
- decoder_kernel_size=31,
- duration_predictor_layers=2,
- duration_predictor_channels=256,
- duration_predictor_kernel_size=3,
- energy_predictor_layers=2,
- energy_predictor_channels=256,
- energy_predictor_kernel_size=3,
- energy_predictor_dropout=0.5,
- energy_embed_kernel_size=1,
- energy_embed_dropout=0.0,
- stop_gradient_from_energy_predictor=False,
- pitch_predictor_layers=5,
- pitch_predictor_channels=256,
- pitch_predictor_kernel_size=5,
- pitch_predictor_dropout=0.5,
- pitch_embed_kernel_size=1,
- pitch_embed_dropout=0.0,
- stop_gradient_from_pitch_predictor=True,
- encoder_dropout_rate=0.2,
- encoder_positional_dropout_rate=0.2,
- encoder_attention_dropout_rate=0.2,
- decoder_dropout_rate=0.2,
- decoder_positional_dropout_rate=0.2,
- decoder_attention_dropout_rate=0.2,
- duration_predictor_dropout_rate=0.2,
- speech_decoder_postnet_dropout=0.5,
- max_source_positions=5000,
- use_masking=True,
- use_weighted_masking=False,
- num_speakers=None,
- num_languages=None,
- speaker_embed_dim=None,
- is_encoder_decoder=True,
- **kwargs,
- ):
- if positionwise_conv_kernel_size % 2 == 0:
- raise ValueError(
- f"positionwise_conv_kernel_size must be odd, but got {positionwise_conv_kernel_size} instead."
- )
- if encoder_kernel_size % 2 == 0:
- raise ValueError(f"encoder_kernel_size must be odd, but got {encoder_kernel_size} instead.")
- if decoder_kernel_size % 2 == 0:
- raise ValueError(f"decoder_kernel_size must be odd, but got {decoder_kernel_size} instead.")
- if duration_predictor_kernel_size % 2 == 0:
- raise ValueError(
- f"duration_predictor_kernel_size must be odd, but got {duration_predictor_kernel_size} instead."
- )
- if energy_predictor_kernel_size % 2 == 0:
- raise ValueError(
- f"energy_predictor_kernel_size must be odd, but got {energy_predictor_kernel_size} instead."
- )
- if energy_embed_kernel_size % 2 == 0:
- raise ValueError(f"energy_embed_kernel_size must be odd, but got {energy_embed_kernel_size} instead.")
- if pitch_predictor_kernel_size % 2 == 0:
- raise ValueError(
- f"pitch_predictor_kernel_size must be odd, but got {pitch_predictor_kernel_size} instead."
- )
- if pitch_embed_kernel_size % 2 == 0:
- raise ValueError(f"pitch_embed_kernel_size must be odd, but got {pitch_embed_kernel_size} instead.")
- if hidden_size % encoder_num_attention_heads != 0:
- raise ValueError("The hidden_size must be evenly divisible by encoder_num_attention_heads.")
- if hidden_size % decoder_num_attention_heads != 0:
- raise ValueError("The hidden_size must be evenly divisible by decoder_num_attention_heads.")
- if use_masking and use_weighted_masking:
- raise ValueError("Either use_masking or use_weighted_masking can be True, but not both.")
-
- self.hidden_size = hidden_size
- self.vocab_size = vocab_size
- self.num_mel_bins = num_mel_bins
- self.encoder_config = {
- "num_attention_heads": encoder_num_attention_heads,
- "layers": encoder_layers,
- "kernel_size": encoder_kernel_size,
- "attention_dropout_rate": encoder_attention_dropout_rate,
- "dropout_rate": encoder_dropout_rate,
- "positional_dropout_rate": encoder_positional_dropout_rate,
- "linear_units": encoder_linear_units,
- "normalize_before": encoder_normalize_before,
- "concat_after": encoder_concat_after,
- }
- self.decoder_config = {
- "num_attention_heads": decoder_num_attention_heads,
- "layers": decoder_layers,
- "kernel_size": decoder_kernel_size,
- "attention_dropout_rate": decoder_attention_dropout_rate,
- "dropout_rate": decoder_dropout_rate,
- "positional_dropout_rate": decoder_positional_dropout_rate,
- "linear_units": decoder_linear_units,
- "normalize_before": decoder_normalize_before,
- "concat_after": decoder_concat_after,
- }
- self.encoder_num_attention_heads = encoder_num_attention_heads
- self.encoder_layers = encoder_layers
- self.duration_predictor_channels = duration_predictor_channels
- self.duration_predictor_kernel_size = duration_predictor_kernel_size
- self.duration_predictor_layers = duration_predictor_layers
- self.energy_embed_dropout = energy_embed_dropout
- self.energy_embed_kernel_size = energy_embed_kernel_size
- self.energy_predictor_channels = energy_predictor_channels
- self.energy_predictor_dropout = energy_predictor_dropout
- self.energy_predictor_kernel_size = energy_predictor_kernel_size
- self.energy_predictor_layers = energy_predictor_layers
- self.pitch_embed_dropout = pitch_embed_dropout
- self.pitch_embed_kernel_size = pitch_embed_kernel_size
- self.pitch_predictor_channels = pitch_predictor_channels
- self.pitch_predictor_dropout = pitch_predictor_dropout
- self.pitch_predictor_kernel_size = pitch_predictor_kernel_size
- self.pitch_predictor_layers = pitch_predictor_layers
- self.positionwise_conv_kernel_size = positionwise_conv_kernel_size
- self.speech_decoder_postnet_units = speech_decoder_postnet_units
- self.speech_decoder_postnet_dropout = speech_decoder_postnet_dropout
- self.speech_decoder_postnet_kernel = speech_decoder_postnet_kernel
- self.speech_decoder_postnet_layers = speech_decoder_postnet_layers
- self.reduction_factor = reduction_factor
- self.speaking_speed = speaking_speed
- self.stop_gradient_from_energy_predictor = stop_gradient_from_energy_predictor
- self.stop_gradient_from_pitch_predictor = stop_gradient_from_pitch_predictor
- self.max_source_positions = max_source_positions
- self.use_cnn_in_conformer = use_cnn_in_conformer
- self.use_macaron_style_in_conformer = use_macaron_style_in_conformer
- self.use_masking = use_masking
- self.use_weighted_masking = use_weighted_masking
- self.num_speakers = num_speakers
- self.num_languages = num_languages
- self.speaker_embed_dim = speaker_embed_dim
- self.duration_predictor_dropout_rate = duration_predictor_dropout_rate
- self.is_encoder_decoder = is_encoder_decoder
-
- super().__init__(
- is_encoder_decoder=is_encoder_decoder,
- **kwargs,
- )
-
-
-class FastSpeech2ConformerHifiGanConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FastSpeech2ConformerHifiGanModel`]. It is used to
- instantiate a FastSpeech2Conformer HiFi-GAN vocoder model according to the specified arguments, defining the model
- architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the
- FastSpeech2Conformer
- [espnet/fastspeech2_conformer_hifigan](https://huggingface.co/espnet/fastspeech2_conformer_hifigan) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- model_in_dim (`int`, *optional*, defaults to 80):
- The number of frequency bins in the input log-mel spectrogram.
- upsample_initial_channel (`int`, *optional*, defaults to 512):
- The number of input channels into the upsampling network.
- upsample_rates (`Tuple[int]` or `List[int]`, *optional*, defaults to `[8, 8, 2, 2]`):
- A tuple of integers defining the stride of each 1D convolutional layer in the upsampling network. The
- length of *upsample_rates* defines the number of convolutional layers and has to match the length of
- *upsample_kernel_sizes*.
- upsample_kernel_sizes (`Tuple[int]` or `List[int]`, *optional*, defaults to `[16, 16, 4, 4]`):
- A tuple of integers defining the kernel size of each 1D convolutional layer in the upsampling network. The
- length of *upsample_kernel_sizes* defines the number of convolutional layers and has to match the length of
- *upsample_rates*.
- resblock_kernel_sizes (`Tuple[int]` or `List[int]`, *optional*, defaults to `[3, 7, 11]`):
- A tuple of integers defining the kernel sizes of the 1D convolutional layers in the multi-receptive field
- fusion (MRF) module.
- resblock_dilation_sizes (`Tuple[Tuple[int]]` or `List[List[int]]`, *optional*, defaults to `[[1, 3, 5], [1, 3, 5], [1, 3, 5]]`):
- A nested tuple of integers defining the dilation rates of the dilated 1D convolutional layers in the
- multi-receptive field fusion (MRF) module.
- initializer_range (`float`, *optional*, defaults to 0.01):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- leaky_relu_slope (`float`, *optional*, defaults to 0.1):
- The angle of the negative slope used by the leaky ReLU activation.
- normalize_before (`bool`, *optional*, defaults to `True`):
- Whether or not to normalize the spectrogram before vocoding using the vocoder's learned mean and variance.
-
- Example:
-
- ```python
- >>> from transformers import FastSpeech2ConformerHifiGan, FastSpeech2ConformerHifiGanConfig
-
- >>> # Initializing a FastSpeech2ConformerHifiGan configuration
- >>> configuration = FastSpeech2ConformerHifiGanConfig()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = FastSpeech2ConformerHifiGan(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "hifigan"
-
- def __init__(
- self,
- model_in_dim=80,
- upsample_initial_channel=512,
- upsample_rates=[8, 8, 2, 2],
- upsample_kernel_sizes=[16, 16, 4, 4],
- resblock_kernel_sizes=[3, 7, 11],
- resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
- initializer_range=0.01,
- leaky_relu_slope=0.1,
- normalize_before=True,
- **kwargs,
- ):
- self.model_in_dim = model_in_dim
- self.upsample_initial_channel = upsample_initial_channel
- self.upsample_rates = upsample_rates
- self.upsample_kernel_sizes = upsample_kernel_sizes
- self.resblock_kernel_sizes = resblock_kernel_sizes
- self.resblock_dilation_sizes = resblock_dilation_sizes
- self.initializer_range = initializer_range
- self.leaky_relu_slope = leaky_relu_slope
- self.normalize_before = normalize_before
- super().__init__(**kwargs)
-
-
-class FastSpeech2ConformerWithHifiGanConfig(PretrainedConfig):
- """
- This is the configuration class to store the configuration of a [`FastSpeech2ConformerWithHifiGan`]. It is used to
- instantiate a `FastSpeech2ConformerWithHifiGanModel` model according to the specified sub-models configurations,
- defining the model architecture.
-
- Instantiating a configuration with the defaults will yield a similar configuration to that of the
- FastSpeech2ConformerModel [espnet/fastspeech2_conformer](https://huggingface.co/espnet/fastspeech2_conformer) and
- FastSpeech2ConformerHifiGan
- [espnet/fastspeech2_conformer_hifigan](https://huggingface.co/espnet/fastspeech2_conformer_hifigan) architectures.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- model_config (`typing.Dict`, *optional*):
- Configuration of the text-to-speech model.
- vocoder_config (`typing.Dict`, *optional*):
- Configuration of the vocoder model.
- model_config ([`FastSpeech2ConformerConfig`], *optional*):
- Configuration of the text-to-speech model.
- vocoder_config ([`FastSpeech2ConformerHiFiGanConfig`], *optional*):
- Configuration of the vocoder model.
-
- Example:
-
- ```python
- >>> from transformers import (
- ... FastSpeech2ConformerConfig,
- ... FastSpeech2ConformerHifiGanConfig,
- ... FastSpeech2ConformerWithHifiGanConfig,
- ... FastSpeech2ConformerWithHifiGan,
- ... )
-
- >>> # Initializing FastSpeech2ConformerWithHifiGan sub-modules configurations.
- >>> model_config = FastSpeech2ConformerConfig()
- >>> vocoder_config = FastSpeech2ConformerHifiGanConfig()
-
- >>> # Initializing a FastSpeech2ConformerWithHifiGan module style configuration
- >>> configuration = FastSpeech2ConformerWithHifiGanConfig(model_config.to_dict(), vocoder_config.to_dict())
-
- >>> # Initializing a model (with random weights)
- >>> model = FastSpeech2ConformerWithHifiGan(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```
- """
-
- model_type = "fastspeech2_conformer_with_hifigan"
- is_composition = True
-
- def __init__(
- self,
- model_config: Dict = None,
- vocoder_config: Dict = None,
- **kwargs,
- ):
- if model_config is None:
- model_config = {}
- logger.info("model_config is None. initializing the model with default values.")
-
- if vocoder_config is None:
- vocoder_config = {}
- logger.info("vocoder_config is None. initializing the coarse model with default values.")
-
- self.model_config = FastSpeech2ConformerConfig(**model_config)
- self.vocoder_config = FastSpeech2ConformerHifiGanConfig(**vocoder_config)
-
- super().__init__(**kwargs)
diff --git a/transformers/models/fastspeech2_conformer/convert_fastspeech2_conformer_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/fastspeech2_conformer/convert_fastspeech2_conformer_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index bb9c432f82292f0a22c276821130f65e30f45e6a..0000000000000000000000000000000000000000
--- a/transformers/models/fastspeech2_conformer/convert_fastspeech2_conformer_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,210 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert FastSpeech2Conformer checkpoint."""
-
-import argparse
-import json
-import re
-from pathlib import Path
-from tempfile import TemporaryDirectory
-
-import torch
-import yaml
-
-from transformers import (
- FastSpeech2ConformerConfig,
- FastSpeech2ConformerModel,
- FastSpeech2ConformerTokenizer,
- logging,
-)
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger("transformers.models.FastSpeech2Conformer")
-
-CONFIG_MAPPING = {
- "adim": "hidden_size",
- "aheads": "num_attention_heads",
- "conformer_dec_kernel_size": "decoder_kernel_size",
- "conformer_enc_kernel_size": "encoder_kernel_size",
- "decoder_normalize_before": "decoder_normalize_before",
- "dlayers": "decoder_layers",
- "dunits": "decoder_linear_units",
- "duration_predictor_chans": "duration_predictor_channels",
- "duration_predictor_kernel_size": "duration_predictor_kernel_size",
- "duration_predictor_layers": "duration_predictor_layers",
- "elayers": "encoder_layers",
- "encoder_normalize_before": "encoder_normalize_before",
- "energy_embed_dropout": "energy_embed_dropout",
- "energy_embed_kernel_size": "energy_embed_kernel_size",
- "energy_predictor_chans": "energy_predictor_channels",
- "energy_predictor_dropout": "energy_predictor_dropout",
- "energy_predictor_kernel_size": "energy_predictor_kernel_size",
- "energy_predictor_layers": "energy_predictor_layers",
- "eunits": "encoder_linear_units",
- "pitch_embed_dropout": "pitch_embed_dropout",
- "pitch_embed_kernel_size": "pitch_embed_kernel_size",
- "pitch_predictor_chans": "pitch_predictor_channels",
- "pitch_predictor_dropout": "pitch_predictor_dropout",
- "pitch_predictor_kernel_size": "pitch_predictor_kernel_size",
- "pitch_predictor_layers": "pitch_predictor_layers",
- "positionwise_conv_kernel_size": "positionwise_conv_kernel_size",
- "postnet_chans": "speech_decoder_postnet_units",
- "postnet_filts": "speech_decoder_postnet_kernel",
- "postnet_layers": "speech_decoder_postnet_layers",
- "reduction_factor": "reduction_factor",
- "stop_gradient_from_energy_predictor": "stop_gradient_from_energy_predictor",
- "stop_gradient_from_pitch_predictor": "stop_gradient_from_pitch_predictor",
- "transformer_dec_attn_dropout_rate": "decoder_attention_dropout_rate",
- "transformer_dec_dropout_rate": "decoder_dropout_rate",
- "transformer_dec_positional_dropout_rate": "decoder_positional_dropout_rate",
- "transformer_enc_attn_dropout_rate": "encoder_attention_dropout_rate",
- "transformer_enc_dropout_rate": "encoder_dropout_rate",
- "transformer_enc_positional_dropout_rate": "encoder_positional_dropout_rate",
- "use_cnn_in_conformer": "use_cnn_in_conformer",
- "use_macaron_style_in_conformer": "use_macaron_style_in_conformer",
- "use_masking": "use_masking",
- "use_weighted_masking": "use_weighted_masking",
- "idim": "input_dim",
- "odim": "num_mel_bins",
- "spk_embed_dim": "speaker_embed_dim",
- "langs": "num_languages",
- "spks": "num_speakers",
-}
-
-
-def remap_model_yaml_config(yaml_config_path):
- with Path(yaml_config_path).open("r", encoding="utf-8") as f:
- args = yaml.safe_load(f)
- args = argparse.Namespace(**args)
-
- remapped_config = {}
-
- model_params = args.tts_conf["text2mel_params"]
- # espnet_config_key -> hf_config_key, any keys not included are ignored
- for espnet_config_key, hf_config_key in CONFIG_MAPPING.items():
- if espnet_config_key in model_params:
- remapped_config[hf_config_key] = model_params[espnet_config_key]
-
- return remapped_config, args.g2p, args.token_list
-
-
-def convert_espnet_state_dict_to_hf(state_dict):
- new_state_dict = {}
- for key in state_dict:
- if "tts.generator.text2mel." in key:
- new_key = key.replace("tts.generator.text2mel.", "")
- if "postnet" in key:
- new_key = new_key.replace("postnet.postnet", "speech_decoder_postnet.layers")
- new_key = new_key.replace(".0.weight", ".conv.weight")
- new_key = new_key.replace(".1.weight", ".batch_norm.weight")
- new_key = new_key.replace(".1.bias", ".batch_norm.bias")
- new_key = new_key.replace(".1.running_mean", ".batch_norm.running_mean")
- new_key = new_key.replace(".1.running_var", ".batch_norm.running_var")
- new_key = new_key.replace(".1.num_batches_tracked", ".batch_norm.num_batches_tracked")
- if "feat_out" in key:
- if "weight" in key:
- new_key = "speech_decoder_postnet.feat_out.weight"
- if "bias" in key:
- new_key = "speech_decoder_postnet.feat_out.bias"
- if "encoder.embed.0.weight" in key:
- new_key = new_key.replace("0.", "")
- if "w_1" in key:
- new_key = new_key.replace("w_1", "conv1")
- if "w_2" in key:
- new_key = new_key.replace("w_2", "conv2")
- if "predictor.conv" in key:
- new_key = new_key.replace(".conv", ".conv_layers")
- pattern = r"(\d)\.(\d)"
- replacement = (
- r"\1.conv" if ("2.weight" not in new_key) and ("2.bias" not in new_key) else r"\1.layer_norm"
- )
- new_key = re.sub(pattern, replacement, new_key)
- if "pitch_embed" in key or "energy_embed" in key:
- new_key = new_key.replace("0", "conv")
- if "encoders" in key:
- new_key = new_key.replace("encoders", "conformer_layers")
- new_key = new_key.replace("norm_final", "final_layer_norm")
- new_key = new_key.replace("norm_mha", "self_attn_layer_norm")
- new_key = new_key.replace("norm_ff_macaron", "ff_macaron_layer_norm")
- new_key = new_key.replace("norm_ff", "ff_layer_norm")
- new_key = new_key.replace("norm_conv", "conv_layer_norm")
- if "lid_emb" in key:
- new_key = new_key.replace("lid_emb", "language_id_embedding")
- if "sid_emb" in key:
- new_key = new_key.replace("sid_emb", "speaker_id_embedding")
-
- new_state_dict[new_key] = state_dict[key]
-
- return new_state_dict
-
-
-@torch.no_grad()
-def convert_FastSpeech2ConformerModel_checkpoint(
- checkpoint_path,
- yaml_config_path,
- pytorch_dump_folder_path,
- repo_id=None,
-):
- model_params, tokenizer_name, vocab = remap_model_yaml_config(yaml_config_path)
- config = FastSpeech2ConformerConfig(**model_params)
-
- # Prepare the model
- model = FastSpeech2ConformerModel(config)
-
- espnet_checkpoint = torch.load(checkpoint_path)
- hf_compatible_state_dict = convert_espnet_state_dict_to_hf(espnet_checkpoint)
-
- model.load_state_dict(hf_compatible_state_dict)
-
- model.save_pretrained(pytorch_dump_folder_path)
-
- # Prepare the tokenizer
- with TemporaryDirectory() as tempdir:
- vocab = {token: id for id, token in enumerate(vocab)}
- vocab_file = Path(tempdir) / "vocab.json"
- with open(vocab_file, "w") as f:
- json.dump(vocab, f)
- should_strip_spaces = "no_space" in tokenizer_name
- tokenizer = FastSpeech2ConformerTokenizer(str(vocab_file), should_strip_spaces=should_strip_spaces)
-
- tokenizer.save_pretrained(pytorch_dump_folder_path)
-
- if repo_id:
- print("Pushing to the hub...")
- model.push_to_hub(repo_id)
- tokenizer.push_to_hub(repo_id)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint")
- parser.add_argument(
- "--yaml_config_path", required=True, default=None, type=str, help="Path to config.yaml of model to convert"
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model."
- )
- parser.add_argument(
- "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
- )
-
- args = parser.parse_args()
- convert_FastSpeech2ConformerModel_checkpoint(
- args.checkpoint_path,
- args.yaml_config_path,
- args.pytorch_dump_folder_path,
- args.push_to_hub,
- )
diff --git a/transformers/models/fastspeech2_conformer/convert_hifigan.py b/transformers/models/fastspeech2_conformer/convert_hifigan.py
deleted file mode 100644
index ec9f57ce7142d619259555fb89f4f1366947fe71..0000000000000000000000000000000000000000
--- a/transformers/models/fastspeech2_conformer/convert_hifigan.py
+++ /dev/null
@@ -1,134 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert FastSpeech2Conformer HiFi-GAN checkpoint."""
-
-import argparse
-from pathlib import Path
-
-import torch
-import yaml
-
-from transformers import FastSpeech2ConformerHifiGan, FastSpeech2ConformerHifiGanConfig, logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger("transformers.models.FastSpeech2Conformer")
-
-
-def load_weights(checkpoint, hf_model, config):
- vocoder_key_prefix = "tts.generator.vocoder."
- checkpoint = {k.replace(vocoder_key_prefix, ""): v for k, v in checkpoint.items() if vocoder_key_prefix in k}
-
- hf_model.apply_weight_norm()
-
- hf_model.conv_pre.weight_g.data = checkpoint["input_conv.weight_g"]
- hf_model.conv_pre.weight_v.data = checkpoint["input_conv.weight_v"]
- hf_model.conv_pre.bias.data = checkpoint["input_conv.bias"]
-
- for i in range(len(config.upsample_rates)):
- hf_model.upsampler[i].weight_g.data = checkpoint[f"upsamples.{i}.1.weight_g"]
- hf_model.upsampler[i].weight_v.data = checkpoint[f"upsamples.{i}.1.weight_v"]
- hf_model.upsampler[i].bias.data = checkpoint[f"upsamples.{i}.1.bias"]
-
- for i in range(len(config.upsample_rates) * len(config.resblock_kernel_sizes)):
- for j in range(len(config.resblock_dilation_sizes)):
- hf_model.resblocks[i].convs1[j].weight_g.data = checkpoint[f"blocks.{i}.convs1.{j}.1.weight_g"]
- hf_model.resblocks[i].convs1[j].weight_v.data = checkpoint[f"blocks.{i}.convs1.{j}.1.weight_v"]
- hf_model.resblocks[i].convs1[j].bias.data = checkpoint[f"blocks.{i}.convs1.{j}.1.bias"]
-
- hf_model.resblocks[i].convs2[j].weight_g.data = checkpoint[f"blocks.{i}.convs2.{j}.1.weight_g"]
- hf_model.resblocks[i].convs2[j].weight_v.data = checkpoint[f"blocks.{i}.convs2.{j}.1.weight_v"]
- hf_model.resblocks[i].convs2[j].bias.data = checkpoint[f"blocks.{i}.convs2.{j}.1.bias"]
-
- hf_model.conv_post.weight_g.data = checkpoint["output_conv.1.weight_g"]
- hf_model.conv_post.weight_v.data = checkpoint["output_conv.1.weight_v"]
- hf_model.conv_post.bias.data = checkpoint["output_conv.1.bias"]
-
- hf_model.remove_weight_norm()
-
-
-def remap_hifigan_yaml_config(yaml_config_path):
- with Path(yaml_config_path).open("r", encoding="utf-8") as f:
- args = yaml.safe_load(f)
- args = argparse.Namespace(**args)
-
- vocoder_type = args.tts_conf["vocoder_type"]
- if vocoder_type != "hifigan_generator":
- raise TypeError(f"Vocoder config must be for `hifigan_generator`, but got {vocoder_type}")
-
- remapped_dict = {}
- vocoder_params = args.tts_conf["vocoder_params"]
-
- # espnet_config_key -> hf_config_key
- key_mappings = {
- "channels": "upsample_initial_channel",
- "in_channels": "model_in_dim",
- "resblock_dilations": "resblock_dilation_sizes",
- "resblock_kernel_sizes": "resblock_kernel_sizes",
- "upsample_kernel_sizes": "upsample_kernel_sizes",
- "upsample_scales": "upsample_rates",
- }
- for espnet_config_key, hf_config_key in key_mappings.items():
- remapped_dict[hf_config_key] = vocoder_params[espnet_config_key]
- remapped_dict["sampling_rate"] = args.tts_conf["sampling_rate"]
- remapped_dict["normalize_before"] = False
- remapped_dict["leaky_relu_slope"] = vocoder_params["nonlinear_activation_params"]["negative_slope"]
-
- return remapped_dict
-
-
-@torch.no_grad()
-def convert_hifigan_checkpoint(
- checkpoint_path,
- pytorch_dump_folder_path,
- yaml_config_path=None,
- repo_id=None,
-):
- if yaml_config_path is not None:
- config_kwargs = remap_hifigan_yaml_config(yaml_config_path)
- config = FastSpeech2ConformerHifiGanConfig(**config_kwargs)
- else:
- config = FastSpeech2ConformerHifiGanConfig()
-
- model = FastSpeech2ConformerHifiGan(config)
-
- orig_checkpoint = torch.load(checkpoint_path)
- load_weights(orig_checkpoint, model, config)
-
- model.save_pretrained(pytorch_dump_folder_path)
-
- if repo_id:
- print("Pushing to the hub...")
- model.push_to_hub(repo_id)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint")
- parser.add_argument("--yaml_config_path", default=None, type=str, help="Path to config.yaml of model to convert")
- parser.add_argument(
- "--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model."
- )
- parser.add_argument(
- "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
- )
-
- args = parser.parse_args()
- convert_hifigan_checkpoint(
- args.checkpoint_path,
- args.pytorch_dump_folder_path,
- args.yaml_config_path,
- args.push_to_hub,
- )
diff --git a/transformers/models/fastspeech2_conformer/convert_model_with_hifigan.py b/transformers/models/fastspeech2_conformer/convert_model_with_hifigan.py
deleted file mode 100644
index 2a780d5cf0b8ea8a69a0bfc7f02796fbaa2b8c5b..0000000000000000000000000000000000000000
--- a/transformers/models/fastspeech2_conformer/convert_model_with_hifigan.py
+++ /dev/null
@@ -1,102 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert FastSpeech2Conformer checkpoint."""
-
-import argparse
-
-import torch
-
-from transformers import (
- FastSpeech2ConformerConfig,
- FastSpeech2ConformerHifiGan,
- FastSpeech2ConformerHifiGanConfig,
- FastSpeech2ConformerModel,
- FastSpeech2ConformerWithHifiGan,
- FastSpeech2ConformerWithHifiGanConfig,
- logging,
-)
-
-from .convert_fastspeech2_conformer_original_pytorch_checkpoint_to_pytorch import (
- convert_espnet_state_dict_to_hf,
- remap_model_yaml_config,
-)
-from .convert_hifigan import load_weights, remap_hifigan_yaml_config
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger("transformers.models.FastSpeech2Conformer")
-
-
-def convert_FastSpeech2ConformerWithHifiGan_checkpoint(
- checkpoint_path,
- yaml_config_path,
- pytorch_dump_folder_path,
- repo_id=None,
-):
- # Prepare the model
- model_params, *_ = remap_model_yaml_config(yaml_config_path)
- model_config = FastSpeech2ConformerConfig(**model_params)
-
- model = FastSpeech2ConformerModel(model_config)
-
- espnet_checkpoint = torch.load(checkpoint_path)
- hf_compatible_state_dict = convert_espnet_state_dict_to_hf(espnet_checkpoint)
- model.load_state_dict(hf_compatible_state_dict)
-
- # Prepare the vocoder
- config_kwargs = remap_hifigan_yaml_config(yaml_config_path)
- vocoder_config = FastSpeech2ConformerHifiGanConfig(**config_kwargs)
-
- vocoder = FastSpeech2ConformerHifiGan(vocoder_config)
- load_weights(espnet_checkpoint, vocoder, vocoder_config)
-
- # Prepare the model + vocoder
- config = FastSpeech2ConformerWithHifiGanConfig.from_sub_model_configs(model_config, vocoder_config)
- with_hifigan_model = FastSpeech2ConformerWithHifiGan(config)
- with_hifigan_model.model = model
- with_hifigan_model.vocoder = vocoder
-
- with_hifigan_model.save_pretrained(pytorch_dump_folder_path)
-
- if repo_id:
- print("Pushing to the hub...")
- with_hifigan_model.push_to_hub(repo_id)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint")
- parser.add_argument(
- "--yaml_config_path", required=True, default=None, type=str, help="Path to config.yaml of model to convert"
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- required=True,
- default=None,
- type=str,
- help="Path to the output `FastSpeech2ConformerModel` PyTorch model.",
- )
- parser.add_argument(
- "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
- )
-
- args = parser.parse_args()
-
- convert_FastSpeech2ConformerWithHifiGan_checkpoint(
- args.checkpoint_path,
- args.yaml_config_path,
- args.pytorch_dump_folder_path,
- args.push_to_hub,
- )
diff --git a/transformers/models/fastspeech2_conformer/modeling_fastspeech2_conformer.py b/transformers/models/fastspeech2_conformer/modeling_fastspeech2_conformer.py
deleted file mode 100644
index c46ef2a8365f0cb8c632cf50ba2b9f8b44d48c69..0000000000000000000000000000000000000000
--- a/transformers/models/fastspeech2_conformer/modeling_fastspeech2_conformer.py
+++ /dev/null
@@ -1,1684 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The Espnet authors, IMS Toucan authors, and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch FastSpeech2Conformer model."""
-
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-from torch import nn
-
-from ...modeling_outputs import BaseModelOutput
-from ...modeling_utils import PreTrainedModel
-from ...utils import ModelOutput, add_start_docstrings, logging, replace_return_docstrings
-from .configuration_fastspeech2_conformer import (
- FastSpeech2ConformerConfig,
- FastSpeech2ConformerHifiGanConfig,
- FastSpeech2ConformerWithHifiGanConfig,
-)
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FASTSPEECH2_CONFORMER_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class FastSpeech2ConformerModelOutput(ModelOutput):
- """
- Output type of [`FastSpeech2ConformerModel`].
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
- Spectrogram generation loss.
- spectrogram (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_bins)`):
- The predicted spectrogram.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
- one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the encoder at the output of each layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
- self-attention heads.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
- one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the decoder at the output of each layer plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
- self-attention heads.
- duration_outputs (`torch.LongTensor` of shape `(batch_size, max_text_length + 1)`, *optional*):
- Outputs of the duration predictor.
- pitch_outputs (`torch.FloatTensor` of shape `(batch_size, max_text_length + 1, 1)`, *optional*):
- Outputs of the pitch predictor.
- energy_outputs (`torch.FloatTensor` of shape `(batch_size, max_text_length + 1, 1)`, *optional*):
- Outputs of the energy predictor.
-
- """
-
- loss: Optional[torch.FloatTensor] = None
- spectrogram: torch.FloatTensor = None
- encoder_last_hidden_state: Optional[torch.FloatTensor] = None
- encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
- duration_outputs: torch.LongTensor = None
- pitch_outputs: torch.FloatTensor = None
- energy_outputs: torch.FloatTensor = None
-
-
-@dataclass
-class FastSpeech2ConformerWithHifiGanOutput(FastSpeech2ConformerModelOutput):
- """
- Output type of [`FastSpeech2ConformerWithHifiGan`].
-
- Args:
- waveform (`torch.FloatTensor` of shape `(batch_size, audio_length)`):
- Speech output as a result of passing the predicted mel spectrogram through the vocoder.
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
- Spectrogram generation loss.
- spectrogram (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_bins)`):
- The predicted spectrogram.
- encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder of the model.
- encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
- one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the encoder at the output of each layer plus the initial embedding outputs.
- encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
- self-attention heads.
- decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
- one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the decoder at the output of each layer plus the initial embedding outputs.
- decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
- self-attention heads.
- duration_outputs (`torch.LongTensor` of shape `(batch_size, max_text_length + 1)`, *optional*):
- Outputs of the duration predictor.
- pitch_outputs (`torch.FloatTensor` of shape `(batch_size, max_text_length + 1, 1)`, *optional*):
- Outputs of the pitch predictor.
- energy_outputs (`torch.FloatTensor` of shape `(batch_size, max_text_length + 1, 1)`, *optional*):
- Outputs of the energy predictor.
- """
-
- waveform: torch.FloatTensor = None
-
-
-_CONFIG_FOR_DOC = "FastSpeech2ConformerConfig"
-
-FASTSPEECH2_CONFORMER_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`FastSpeech2ConformerConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-HIFIGAN_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`FastSpeech2ConformerConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FASTSPEECH2_CONFORMER_WITH_HIFIGAN_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`FastSpeech2ConformerWithHifiGanConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-def length_regulator(encoded_embeddings, duration_labels, speaking_speed=1.0):
- """
- Length regulator for feed-forward Transformer.
-
- This is the length regulator module described in `FastSpeech: Fast, Robust and Controllable Text to Speech`
- https://arxiv.org/pdf/1905.09263.pdf. The length regulator expands char or phoneme-level embedding features to
- frame-level by repeating each feature based on the corresponding predicted durations.
-
- Args:
- encoded_embeddings (`torch.Tensor` of shape `(batch_size, max_text_length, embedding_dim)`):
- Batch of sequences of char or phoneme embeddings.
- duration_labels (`torch.LongTensor` of shape `(batch_size, time)`):
- Batch of durations of each frame.
- speaking_speed (`float`, *optional*, defaults to 1.0):
- Value to control speed of speech.
-
- Returns:
- `torch.Tensor`:
- Replicated input tensor based on durations (batch_size, time*, embedding_dim).
- """
-
- if speaking_speed <= 0:
- raise ValueError("`speaking_speed` must be greater than 0.")
- elif speaking_speed != 1.0:
- duration_labels = torch.round(duration_labels.float() * speaking_speed).long()
-
- if duration_labels.sum() == 0:
- duration_labels[duration_labels.sum(dim=1).eq(0)] = 1
-
- # Calculate the maximum length needed
- max_len = torch.sum(duration_labels, dim=1).max()
-
- # Create a padded tensor to hold the results
- hidden_states = torch.zeros(
- (encoded_embeddings.size(0), max_len, encoded_embeddings.size(2)),
- dtype=torch.float,
- device=encoded_embeddings.device,
- )
-
- # Loop through the batch and fill in the data
- for i, (encoded_embedding, target_duration) in enumerate(zip(encoded_embeddings, duration_labels)):
- repeated = torch.repeat_interleave(encoded_embedding, target_duration, dim=0)
- hidden_states[i, : repeated.size(0)] = repeated
-
- return hidden_states
-
-
-class FastSpeech2ConformerDurationPredictor(nn.Module):
- """
- Duration predictor module.
-
- This is a module of duration predictor described in the paper 'FastSpeech: Fast, Robust and Controllable Text to
- Speech' https://arxiv.org/pdf/1905.09263.pdf The duration predictor predicts a duration of each frame in log domain
- from the hidden embeddings of encoder.
-
- Note:
- The calculation domain of outputs is different between in `forward` and in `inference`. In `forward`, the
- outputs are calculated in log domain but in `inference`, those are calculated in linear domain.
-
- """
-
- def __init__(self, config: FastSpeech2ConformerConfig):
- super().__init__()
-
- self.conv_layers = nn.ModuleList()
- self.log_domain_offset = 1.0
-
- for layer_idx in range(config.duration_predictor_layers):
- num_chans = config.duration_predictor_channels
- input_channels = config.hidden_size if layer_idx == 0 else num_chans
- layer = FastSpeech2ConformerPredictorLayer(
- input_channels,
- num_chans,
- config.duration_predictor_kernel_size,
- config.duration_predictor_dropout_rate,
- )
- self.conv_layers.append(layer)
- self.linear = nn.Linear(config.duration_predictor_channels, 1)
-
- def forward(self, encoder_hidden_states):
- """
- Args:
- hidden_states (`torch.Tensor` of shape `(batch_size, max_text_length, input_dim)`):
- Batch of input sequences.
- padding_masks (`torch.ByteTensor` of shape `(batch_size, max_text_length)`, *optional*):
- Batch of masks indicating padded part.
-
- Returns:
- `torch.Tensor`: Batch of predicted durations in log domain `(batch_size, max_text_length)`.
-
- """
- # (batch_size, input_dim, max_text_length)
- hidden_states = encoder_hidden_states.transpose(1, -1)
- for layer in self.conv_layers:
- hidden_states = layer(hidden_states)
-
- # NOTE: calculate in log domain, (batch_size, max_text_length)
- hidden_states = self.linear(hidden_states.transpose(1, -1)).squeeze(-1)
-
- if not self.training:
- # NOTE: calculate in linear domain
- hidden_states = torch.clamp(torch.round(hidden_states.exp() - self.log_domain_offset), min=0).long()
-
- return hidden_states
-
-
-# Copied from transformers.models.speecht5.modeling_speecht5.SpeechT5BatchNormConvLayer
-class FastSpeech2ConformerBatchNormConvLayer(nn.Module):
- def __init__(self, config, layer_id=0):
- super().__init__()
-
- if layer_id == 0:
- in_conv_dim = config.num_mel_bins
- else:
- in_conv_dim = config.speech_decoder_postnet_units
-
- if layer_id == config.speech_decoder_postnet_layers - 1:
- out_conv_dim = config.num_mel_bins
- else:
- out_conv_dim = config.speech_decoder_postnet_units
-
- self.conv = nn.Conv1d(
- in_conv_dim,
- out_conv_dim,
- kernel_size=config.speech_decoder_postnet_kernel,
- stride=1,
- padding=(config.speech_decoder_postnet_kernel - 1) // 2,
- bias=False,
- )
- self.batch_norm = nn.BatchNorm1d(out_conv_dim)
-
- if layer_id < config.speech_decoder_postnet_layers - 1:
- self.activation = nn.Tanh()
- else:
- self.activation = None
-
- self.dropout = nn.Dropout(config.speech_decoder_postnet_dropout)
-
- def forward(self, hidden_states):
- hidden_states = self.conv(hidden_states)
- hidden_states = self.batch_norm(hidden_states)
- if self.activation is not None:
- hidden_states = self.activation(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-class FastSpeech2ConformerSpeechDecoderPostnet(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.feat_out = nn.Linear(config.hidden_size, config.num_mel_bins * config.reduction_factor)
- self.layers = nn.ModuleList(
- [FastSpeech2ConformerBatchNormConvLayer(config, i) for i in range(config.speech_decoder_postnet_layers)]
- )
-
- def forward(self, hidden_states: torch.Tensor):
- outputs_before_postnet = self.feat_out(hidden_states).view(hidden_states.size(0), -1, self.config.num_mel_bins)
- layer_output = outputs_before_postnet.transpose(1, 2)
- for layer in self.layers:
- layer_output = layer(layer_output)
- outputs_after_postnet = outputs_before_postnet + layer_output.transpose(1, 2)
- return outputs_before_postnet, outputs_after_postnet
-
-
-class FastSpeech2ConformerPredictorLayer(nn.Module):
- def __init__(self, input_channels, num_chans, kernel_size, dropout_rate):
- super().__init__()
- self.conv = nn.Conv1d(
- input_channels,
- num_chans,
- kernel_size,
- stride=1,
- padding=(kernel_size - 1) // 2,
- )
- self.activation = nn.ReLU()
- self.layer_norm = nn.LayerNorm(num_chans)
- self.dropout = nn.Dropout(dropout_rate)
-
- def forward(self, hidden_states):
- hidden_states = self.conv(hidden_states)
- hidden_states = self.activation(hidden_states)
-
- # Perform layer norm on dimension 1
- hidden_states = hidden_states.transpose(1, -1)
- hidden_states = self.layer_norm(hidden_states)
- hidden_states = hidden_states.transpose(1, -1)
-
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-class FastSpeech2ConformerVariancePredictor(nn.Module):
- def __init__(
- self,
- config: FastSpeech2ConformerConfig,
- num_layers=2,
- num_chans=384,
- kernel_size=3,
- dropout_rate=0.5,
- ):
- """
- Initilize variance predictor module.
-
- Args:
- input_dim (`int`): Input dimension.
- num_layers (`int`, *optional*, defaults to 2): Number of convolutional layers.
- num_chans (`int`, *optional*, defaults to 384): Number of channels of convolutional layers.
- kernel_size (`int`, *optional*, defaults to 3): Kernel size of convolutional layers.
- dropout_rate (`float`, *optional*, defaults to 0.5): Dropout rate.
- """
- super().__init__()
- self.conv_layers = nn.ModuleList()
- for idx in range(num_layers):
- input_channels = config.hidden_size if idx == 0 else num_chans
- layer = FastSpeech2ConformerPredictorLayer(input_channels, num_chans, kernel_size, dropout_rate)
- self.conv_layers.append(layer)
- self.linear = nn.Linear(num_chans, 1)
-
- def forward(self, encoder_hidden_states, padding_masks=None):
- """
- Calculate forward propagation.
-
- Args:
- encoder_hidden_states (`torch.Tensor` of shape `(batch_size, max_text_length, input_dim)`):
- Batch of input sequences.
- padding_masks (`torch.ByteTensor` of shape `(batch_size, max_text_length)`, *optional*):
- Batch of masks indicating padded part.
-
- Returns:
- Tensor: Batch of predicted sequences `(batch_size, max_text_length, 1)`.
- """
- # (batch_size, input_dim, max_text_length)
- hidden_states = encoder_hidden_states.transpose(1, -1)
- for layer in self.conv_layers:
- hidden_states = layer(hidden_states)
-
- hidden_states = self.linear(hidden_states.transpose(1, 2))
-
- if padding_masks is not None:
- hidden_states = hidden_states.masked_fill(padding_masks, 0.0)
-
- return hidden_states
-
-
-class FastSpeech2ConformerVarianceEmbedding(nn.Module):
- def __init__(
- self,
- in_channels=1,
- out_channels=384,
- kernel_size=1,
- padding=0,
- dropout_rate=0.0,
- ):
- super().__init__()
- self.conv = nn.Conv1d(
- in_channels=in_channels,
- out_channels=out_channels,
- kernel_size=kernel_size,
- padding=padding,
- )
- self.dropout = nn.Dropout(dropout_rate)
-
- def forward(self, hidden_states):
- hidden_states = hidden_states.transpose(1, 2)
- hidden_states = self.conv(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = hidden_states.transpose(1, 2)
- return hidden_states
-
-
-class FastSpeech2ConformerAttention(nn.Module):
- """
- Multi-Head attention layer with relative position encoding. Details can be found in
- https://github.com/espnet/espnet/pull/2816. Paper: https://arxiv.org/abs/1901.02860.
- """
-
- def __init__(self, config: FastSpeech2ConformerConfig, module_config):
- """Construct an FastSpeech2ConformerAttention object."""
- super().__init__()
- # We assume d_v always equals dim_key
- self.num_heads = module_config["num_attention_heads"]
- self.hidden_size = config.hidden_size
- self.dim_key = self.hidden_size // self.num_heads
- self.head_dim = self.hidden_size // self.num_heads
- self.linear_q = nn.Linear(self.hidden_size, self.hidden_size)
- self.linear_k = nn.Linear(self.hidden_size, self.hidden_size)
- self.linear_v = nn.Linear(self.hidden_size, self.hidden_size)
- self.linear_out = nn.Linear(self.hidden_size, self.hidden_size)
- self.dropout = nn.Dropout(p=module_config["attention_dropout_rate"])
-
- # linear transformation for positional encoding
- self.linear_pos = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
- # these two learnable bias are used in matrix c and matrix d
- # as described in https://arxiv.org/abs/1901.02860 Section 3.3
- self.pos_bias_u = nn.Parameter(torch.Tensor(self.num_heads, self.head_dim))
- self.pos_bias_v = nn.Parameter(torch.Tensor(self.num_heads, self.head_dim))
-
- def shift_relative_position_tensor(self, pos_tensor):
- """
- Args:
- pos_tensor (torch.Tensor of shape (batch_size, head, time1, 2*time1-1)): Input tensor.
- """
- zero_pad = torch.zeros((*pos_tensor.size()[:3], 1), device=pos_tensor.device, dtype=pos_tensor.dtype)
- pos_tensor_padded = torch.cat([zero_pad, pos_tensor], dim=-1)
-
- pos_tensor_padded = pos_tensor_padded.view(*pos_tensor.size()[:2], pos_tensor.size(3) + 1, pos_tensor.size(2))
- # only keep the positions from 0 to time2
- pos_tensor = pos_tensor_padded[:, :, 1:].view_as(pos_tensor)[:, :, :, : pos_tensor.size(-1) // 2 + 1]
-
- return pos_tensor
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- pos_emb: Optional[torch.Tensor] = None,
- output_attentions: Optional[torch.Tensor] = False,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
- """
- Compute 'Scaled Dot Product Attention' with rel. positional encoding.
-
- Args:
- hidden_states (`torch.Tensor` of shape `(batch, time2, size)`): Values of the hidden states
- attention_mask (`torch.Tensor` of shape `(batch, time1, time2)`): Mask tensor.
- pos_emb (`torch.Tensor` of shape `(batch, 2*time1-1, size)`): Positional embedding tensor.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- Returns:
- `torch.Tensor`: Output tensor of shape `(batch, time1, d_model)`.
- """
- bsz, q_len, _ = hidden_states.size()
- query_states = self.linear_q(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
- key_states = self.linear_k(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
- value_states = self.linear_v(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
-
- bsz_pos = pos_emb.size(0)
- pos_encoding = self.linear_pos(pos_emb).view(bsz_pos, -1, self.num_heads, self.head_dim)
-
- # (batch_size, head, time1, dim_key)
- query_with_bias_u = (query_states + self.pos_bias_u).transpose(1, 2)
- # (batch_size, head, time1, dim_key)
- query_with_bias_v = (query_states + self.pos_bias_v).transpose(1, 2)
-
- # compute attention score
- # first compute matrix a and matrix c
- # as described in https://arxiv.org/abs/1901.02860 Section 3.3
- # (batch_size, head, time1, time2)
- matrix_ac = torch.matmul(query_with_bias_u, key_states.permute(0, 2, 3, 1))
-
- # compute matrix b and matrix d
- # (batch_size, head, time1, 2*time1-1)
- matrix_bd = torch.matmul(query_with_bias_v, pos_encoding.permute(0, 2, 3, 1))
- matrix_bd = self.shift_relative_position_tensor(matrix_bd)
-
- # (batch_size, head, time1, time2)
- scores = (matrix_ac + matrix_bd) / math.sqrt(self.dim_key)
-
- # Forward attention
- if attention_mask is not None:
- expected_size = (bsz, 1, q_len)
- if attention_mask.size() != expected_size:
- raise ValueError(f"Attention mask should be of size {expected_size}, but is {attention_mask.size()}")
- attention_mask = attention_mask.unsqueeze(1).eq(0)
- min_value = float(torch.finfo(scores.dtype).min)
- scores = scores.masked_fill(attention_mask, min_value)
- attn_weights = torch.softmax(scores, dim=-1).masked_fill(attention_mask, 0.0)
- else:
- attn_weights = torch.softmax(scores, dim=-1)
-
- attn_weights = self.dropout(attn_weights)
- attn_output = torch.matmul(attn_weights, value_states.transpose(1, 2))
- attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
-
- attn_output = self.linear_out(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights
-
-
-class FastSpeech2ConformerConvolutionModule(nn.Module):
- def __init__(self, config: FastSpeech2ConformerConfig, module_config):
- super().__init__()
- # kernel_size should be an odd number for 'SAME' padding
- channels = config.hidden_size
- kernel_size = module_config["kernel_size"]
- self.pointwise_conv1 = nn.Conv1d(channels, 2 * channels, kernel_size=1, stride=1, padding=0, bias=True)
- self.depthwise_conv = nn.Conv1d(
- channels, channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2, groups=channels, bias=True
- )
- self.norm = nn.BatchNorm1d(channels)
- self.pointwise_conv2 = nn.Conv1d(channels, channels, kernel_size=1, stride=1, padding=0, bias=True)
-
- def forward(self, hidden_states):
- """
- Compute convolution module.
-
- Args:
- hidden_states (`torch.Tensor` of shape `(batch, time, channels)`): Input tensor.
-
- Returns:
- `torch.Tensor`: Output tensor of shape `(batch, time, channels)`.
-
- """
- # exchange the temporal dimension and the feature dimension
- hidden_states = hidden_states.transpose(1, 2)
-
- # GLU mechanism, (batch_size, 2*channel, dim)
- hidden_states = self.pointwise_conv1(hidden_states)
- # (batch_size, channel, dim)
- hidden_states = nn.functional.glu(hidden_states, dim=1)
-
- # 1D Depthwise Conv
- hidden_states = self.depthwise_conv(hidden_states)
- hidden_states = self.norm(hidden_states)
-
- hidden_states = hidden_states * torch.sigmoid(hidden_states)
-
- hidden_states = self.pointwise_conv2(hidden_states)
-
- return hidden_states.transpose(1, 2)
-
-
-class FastSpeech2ConformerEncoderLayer(nn.Module):
- def __init__(self, config: FastSpeech2ConformerConfig, module_config):
- super().__init__()
-
- # self-attention module definition
- self.self_attn = FastSpeech2ConformerAttention(config, module_config)
-
- # feed-forward module definition
- self.feed_forward = FastSpeech2ConformerMultiLayeredConv1d(config, module_config)
-
- self.macaron_style = config.use_macaron_style_in_conformer
- if self.macaron_style:
- self.feed_forward_macaron = FastSpeech2ConformerMultiLayeredConv1d(config, module_config)
- self.ff_macaron_layer_norm = nn.LayerNorm(config.hidden_size)
- self.ff_scale = 0.5
- else:
- self.ff_scale = 1.0
-
- # convolution module definition
- self.use_cnn_module = config.use_cnn_in_conformer
- if self.use_cnn_module:
- self.conv_module = FastSpeech2ConformerConvolutionModule(config, module_config)
- self.conv_layer_norm = nn.LayerNorm(config.hidden_size)
- self.final_layer_norm = nn.LayerNorm(config.hidden_size)
-
- self.ff_layer_norm = nn.LayerNorm(config.hidden_size)
-
- self.self_attn_layer_norm = nn.LayerNorm(config.hidden_size)
-
- self.dropout = nn.Dropout(module_config["dropout_rate"])
- self.size = config.hidden_size
- self.normalize_before = module_config["normalize_before"]
- self.concat_after = module_config["concat_after"]
- if self.concat_after:
- self.concat_linear = nn.Linear(config.hidden_size + config.hidden_size, config.hidden_size)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- pos_emb: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[torch.Tensor] = False,
- ):
- """
- Compute encoded features.
-
- Args:
- hidden_states (`torch.Tensor` of shape `(batch, time, size)`): Input tensor.
- pos_emb (`torch.Tensor` of shape `(1, time, size)`): Positional embeddings tensor.
- attention_mask (`torch.Tensor` of shape `(batch, time)`): Attention mask tensor for the input.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- Returns:
- `torch.Tensor`: Output tensor of shape `(batch, time, size)`.
-
- """
- # whether to use macaron style
- if self.macaron_style:
- residual = hidden_states
- if self.normalize_before:
- hidden_states = self.ff_macaron_layer_norm(hidden_states)
- hidden_states = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(hidden_states))
- if not self.normalize_before:
- hidden_states = self.ff_macaron_layer_norm(hidden_states)
-
- # multi-headed self-attention module
- residual = hidden_states
- if self.normalize_before:
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- attention_output, attention_scores = self.self_attn(
- hidden_states, attention_mask=attention_mask, pos_emb=pos_emb, output_attentions=output_attentions
- )
-
- if self.concat_after:
- x_concat = torch.cat((hidden_states, attention_output), dim=-1)
- hidden_states = self.concat_linear(x_concat)
- hidden_states = residual + hidden_states
- else:
- hidden_states = self.dropout(attention_output)
- hidden_states = residual + hidden_states
- if not self.normalize_before:
- hidden_states = self.self_attn_layer_norm(hidden_states)
-
- # convolution module
- if self.use_cnn_module:
- residual = hidden_states
- if self.normalize_before:
- hidden_states = self.conv_layer_norm(hidden_states)
- hidden_states = self.conv_module(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = residual + hidden_states
- if not self.normalize_before:
- hidden_states = self.conv_layer_norm(hidden_states)
-
- # feed forward module
- residual = hidden_states
- if self.normalize_before:
- hidden_states = self.ff_layer_norm(hidden_states)
- hidden_states = self.feed_forward(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = residual + self.ff_scale * hidden_states
- if not self.normalize_before:
- hidden_states = self.ff_layer_norm(hidden_states)
-
- if self.conv_module is not None:
- hidden_states = self.final_layer_norm(hidden_states)
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attention_scores,)
-
- return outputs
-
-
-class FastSpeech2ConformerMultiLayeredConv1d(nn.Module):
- """
- Multi-layered conv1d for Transformer block.
-
- This is a module of multi-layered conv1d designed to replace positionwise feed-forward network in Transformer
- block, which is introduced in 'FastSpeech: Fast, Robust and Controllable Text to Speech'
- https://arxiv.org/pdf/1905.09263.pdf
- """
-
- def __init__(self, config: FastSpeech2ConformerConfig, module_config):
- """
- Initialize FastSpeech2ConformerMultiLayeredConv1d module.
-
- Args:
- input_channels (`int`): Number of input channels.
- hidden_channels (`int`): Number of hidden channels.
- kernel_size (`int`): Kernel size of conv1d.
- dropout_rate (`float`): Dropout rate.
- """
- super().__init__()
- input_channels = config.hidden_size
- hidden_channels = module_config["linear_units"]
- kernel_size = config.positionwise_conv_kernel_size
- self.conv1 = nn.Conv1d(input_channels, hidden_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)
- self.conv2 = nn.Conv1d(hidden_channels, input_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)
- self.dropout = nn.Dropout(module_config["dropout_rate"])
-
- def forward(self, hidden_states):
- """
- Calculate forward propagation.
-
- Args:
- hidden_states (torch.Tensor): Batch of input tensors (batch_size, time, input_channels).
-
- Returns:
- torch.Tensor: Batch of output tensors (batch_size, time, hidden_channels).
- """
- hidden_states = hidden_states.transpose(-1, 1)
- hidden_states = self.conv1(hidden_states)
- hidden_states = torch.relu(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.conv2(hidden_states)
- hidden_states = hidden_states.transpose(-1, 1)
- return hidden_states
-
-
-class FastSpeech2ConformerRelPositionalEncoding(nn.Module):
- """
- Args:
- Relative positional encoding module (new implementation). Details can be found in
- https://github.com/espnet/espnet/pull/2816. See : Appendix Batch in https://arxiv.org/abs/1901.02860
- config (`FastSpeech2ConformerConfig`):
- FastSpeech2ConformerConfig instance.
- module_config (`dict`):
- Dictionary containing the encoder or decoder module configuration from the `FastSpeech2ConformerConfig`.
- """
-
- def __init__(self, config: FastSpeech2ConformerConfig, module_config):
- """
- Construct an PositionalEncoding object.
- """
- super().__init__()
- self.embed_dim = config.hidden_size
- self.input_scale = math.sqrt(self.embed_dim)
- self.dropout = nn.Dropout(p=module_config["positional_dropout_rate"])
- self.pos_enc = None
- self.max_len = 5000
- self.extend_pos_enc(torch.tensor(0.0).expand(1, self.max_len))
-
- def extend_pos_enc(self, x):
- """Reset the positional encodings."""
- if self.pos_enc is not None:
- # self.pos_enc contains both positive and negative parts
- # the length of self.pos_enc is 2 * input_len - 1
- if self.pos_enc.size(1) >= x.size(1) * 2 - 1:
- if self.pos_enc.dtype != x.dtype or self.pos_enc.device != x.device:
- self.pos_enc = self.pos_enc.to(dtype=x.dtype, device=x.device)
- return
- # Suppose `i` means to the position of query vector and `j` means the
- # position of key vector. We use position relative positions when keys
- # are to the left (i>j) and negative relative positions otherwise (i 1
- if self.multilingual_model:
- self.language_id_embedding = torch.nn.Embedding(config.num_languages, self.hidden_size)
-
- self.multispeaker_model = config.num_speakers is not None and config.num_speakers > 1
- if self.multispeaker_model:
- self.speaker_id_embedding = torch.nn.Embedding(config.num_speakers, config.hidden_size)
-
- self.speaker_embed_dim = config.speaker_embed_dim
- if self.speaker_embed_dim:
- self.projection = nn.Linear(config.hidden_size + self.speaker_embed_dim, config.hidden_size)
-
- self.encoder = FastSpeech2ConformerEncoder(config, config.encoder_config, use_encoder_input_layer=True)
-
- self.duration_predictor = FastSpeech2ConformerDurationPredictor(config)
-
- self.pitch_predictor = FastSpeech2ConformerVariancePredictor(
- config,
- num_layers=config.pitch_predictor_layers,
- num_chans=config.pitch_predictor_channels,
- kernel_size=config.pitch_predictor_kernel_size,
- dropout_rate=config.pitch_predictor_dropout,
- )
- # continuous pitch + FastPitch style avg
- self.pitch_embed = FastSpeech2ConformerVarianceEmbedding(
- out_channels=self.hidden_size,
- kernel_size=config.pitch_embed_kernel_size,
- padding=(config.pitch_embed_kernel_size - 1) // 2,
- dropout_rate=config.pitch_embed_dropout,
- )
-
- self.energy_predictor = FastSpeech2ConformerVariancePredictor(
- config,
- num_layers=config.energy_predictor_layers,
- num_chans=config.energy_predictor_channels,
- kernel_size=config.energy_predictor_kernel_size,
- dropout_rate=config.energy_predictor_dropout,
- )
- # continuous energy + FastPitch style avg
- self.energy_embed = FastSpeech2ConformerVarianceEmbedding(
- out_channels=self.hidden_size,
- kernel_size=config.energy_embed_kernel_size,
- padding=(config.energy_embed_kernel_size - 1) // 2,
- dropout_rate=config.energy_embed_dropout,
- )
-
- # The decoder is an encoder
- self.decoder = FastSpeech2ConformerEncoder(config, config.decoder_config, use_encoder_input_layer=False)
-
- self.speech_decoder_postnet = FastSpeech2ConformerSpeechDecoderPostnet(config)
-
- self.criterion = FastSpeech2ConformerLoss(config)
-
- self.post_init()
-
- @replace_return_docstrings(output_type=FastSpeech2ConformerModelOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: torch.LongTensor,
- attention_mask: Optional[torch.LongTensor] = None,
- spectrogram_labels: Optional[torch.FloatTensor] = None,
- duration_labels: Optional[torch.LongTensor] = None,
- pitch_labels: Optional[torch.FloatTensor] = None,
- energy_labels: Optional[torch.FloatTensor] = None,
- speaker_ids: Optional[torch.LongTensor] = None,
- lang_ids: Optional[torch.LongTensor] = None,
- speaker_embedding: Optional[torch.FloatTensor] = None,
- return_dict: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- ) -> Union[Tuple, FastSpeech2ConformerModelOutput]:
- """
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Input sequence of text vectors.
- attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*, defaults to `None`):
- Mask to avoid performing convolution and attention on padding token indices. Mask values selected in
- `[0, 1]`: 0 for tokens that are **masked**, 1 for tokens that are **not masked**.
- spectrogram_labels (`torch.FloatTensor` of shape `(batch_size, max_spectrogram_length, num_mel_bins)`, *optional*, defaults to `None`):
- Batch of padded target features.
- duration_labels (`torch.LongTensor` of shape `(batch_size, sequence_length + 1)`, *optional*, defaults to `None`):
- Batch of padded durations.
- pitch_labels (`torch.FloatTensor` of shape `(batch_size, sequence_length + 1, 1)`, *optional*, defaults to `None`):
- Batch of padded token-averaged pitch.
- energy_labels (`torch.FloatTensor` of shape `(batch_size, sequence_length + 1, 1)`, *optional*, defaults to `None`):
- Batch of padded token-averaged energy.
- speaker_ids (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*, defaults to `None`):
- Speaker ids used to condition features of speech output by the model.
- lang_ids (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*, defaults to `None`):
- Language ids used to condition features of speech output by the model.
- speaker_embedding (`torch.FloatTensor` of shape `(batch_size, embedding_dim)`, *optional*, defaults to `None`):
- Embedding containing conditioning signals for the features of the speech.
- return_dict (`bool`, *optional*, defaults to `None`):
- Whether or not to return a [`FastSpeech2ConformerModelOutput`] instead of a plain tuple.
- output_attentions (`bool`, *optional*, defaults to `None`):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*, defaults to `None`):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import (
- ... FastSpeech2ConformerTokenizer,
- ... FastSpeech2ConformerModel,
- ... FastSpeech2ConformerHifiGan,
- ... )
-
- >>> tokenizer = FastSpeech2ConformerTokenizer.from_pretrained("espnet/fastspeech2_conformer")
- >>> inputs = tokenizer("some text to convert to speech", return_tensors="pt")
- >>> input_ids = inputs["input_ids"]
-
- >>> model = FastSpeech2ConformerModel.from_pretrained("espnet/fastspeech2_conformer")
- >>> output_dict = model(input_ids, return_dict=True)
- >>> spectrogram = output_dict["spectrogram"]
-
- >>> vocoder = FastSpeech2ConformerHifiGan.from_pretrained("espnet/fastspeech2_conformer_hifigan")
- >>> waveform = vocoder(spectrogram)
- >>> print(waveform.shape)
- torch.Size([1, 49664])
- ```
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- if attention_mask is None:
- attention_mask = torch.ones(input_ids.shape, device=input_ids.device)
-
- has_missing_labels = (
- spectrogram_labels is None or duration_labels is None or pitch_labels is None or energy_labels is None
- )
- if self.training and has_missing_labels:
- raise ValueError("All labels must be provided to run in training mode.")
-
- # forward encoder
- text_masks = attention_mask.unsqueeze(-2)
-
- encoder_outputs = self.encoder(
- input_ids,
- text_masks,
- output_hidden_states=output_hidden_states,
- output_attentions=output_attentions,
- return_dict=return_dict,
- )
- hidden_states = encoder_outputs[0]
-
- # Integrate with language id, speaker id, and speaker embedding
- if self.multispeaker_model and speaker_ids is not None:
- speaker_id_embeddings = self.speaker_id_embedding(speaker_ids.view(-1))
- hidden_states = hidden_states + speaker_id_embeddings.unsqueeze(1)
-
- if self.multilingual_model and lang_ids is not None:
- language_id_embbedings = self.language_id_embedding(lang_ids.view(-1))
- hidden_states = hidden_states + language_id_embbedings.unsqueeze(1)
-
- if self.speaker_embed_dim is not None and speaker_embedding is not None:
- embeddings_expanded = (
- nn.functional.normalize(speaker_embedding).unsqueeze(1).expand(-1, hidden_states.size(1), -1)
- )
- hidden_states = self.projection(torch.cat([hidden_states, embeddings_expanded], dim=-1))
-
- # forward duration predictor and variance predictors
- duration_mask = ~attention_mask.bool()
-
- if self.stop_gradient_from_pitch_predictor:
- pitch_predictions = self.pitch_predictor(hidden_states.detach(), duration_mask.unsqueeze(-1))
- else:
- pitch_predictions = self.pitch_predictor(hidden_states, duration_mask.unsqueeze(-1))
-
- if self.stop_gradient_from_energy_predictor:
- energy_predictions = self.energy_predictor(hidden_states.detach(), duration_mask.unsqueeze(-1))
- else:
- energy_predictions = self.energy_predictor(hidden_states, duration_mask.unsqueeze(-1))
-
- duration_predictions = self.duration_predictor(hidden_states)
- duration_predictions = duration_predictions.masked_fill(duration_mask, 0.0)
-
- if not self.training:
- # use prediction in inference
- embedded_pitch_curve = self.pitch_embed(pitch_predictions)
- embedded_energy_curve = self.energy_embed(energy_predictions)
- hidden_states = hidden_states + embedded_energy_curve + embedded_pitch_curve
- hidden_states = length_regulator(hidden_states, duration_predictions, self.config.speaking_speed)
- else:
- # use groundtruth in training
- embedded_pitch_curve = self.pitch_embed(pitch_labels)
- embedded_energy_curve = self.energy_embed(energy_labels)
- hidden_states = hidden_states + embedded_energy_curve + embedded_pitch_curve
- hidden_states = length_regulator(hidden_states, duration_labels)
-
- # forward decoder
- if not self.training:
- hidden_mask = None
- else:
- spectrogram_mask = (spectrogram_labels != -100).any(dim=-1)
- spectrogram_mask = spectrogram_mask.int()
- if self.reduction_factor > 1:
- length_dim = spectrogram_mask.shape[1] - spectrogram_mask.shape[1] % self.reduction_factor
- spectrogram_mask = spectrogram_mask[:, :, :length_dim]
- hidden_mask = spectrogram_mask.unsqueeze(-2)
-
- decoder_outputs = self.decoder(
- hidden_states,
- hidden_mask,
- output_hidden_states=output_hidden_states,
- output_attentions=output_attentions,
- return_dict=return_dict,
- )
-
- outputs_before_postnet, outputs_after_postnet = self.speech_decoder_postnet(decoder_outputs[0])
-
- loss = None
- if self.training:
- # calculate loss
- loss_duration_mask = ~duration_mask
- loss_spectrogram_mask = spectrogram_mask.unsqueeze(-1).bool()
- loss = self.criterion(
- outputs_after_postnet=outputs_after_postnet,
- outputs_before_postnet=outputs_before_postnet,
- duration_outputs=duration_predictions,
- pitch_outputs=pitch_predictions,
- energy_outputs=energy_predictions,
- spectrogram_labels=spectrogram_labels,
- duration_labels=duration_labels,
- pitch_labels=pitch_labels,
- energy_labels=energy_labels,
- duration_mask=loss_duration_mask,
- spectrogram_mask=loss_spectrogram_mask,
- )
-
- if not return_dict:
- postnet_outputs = (outputs_after_postnet,)
- audio_feature_predictions = (
- duration_predictions,
- pitch_predictions,
- energy_predictions,
- )
- outputs = postnet_outputs + encoder_outputs + decoder_outputs[1:] + audio_feature_predictions
- return ((loss,) + outputs) if loss is not None else outputs
-
- return FastSpeech2ConformerModelOutput(
- loss=loss,
- spectrogram=outputs_after_postnet,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- duration_outputs=duration_predictions,
- pitch_outputs=pitch_predictions,
- energy_outputs=energy_predictions,
- )
-
-
-# Copied from transformers.models.speecht5.modeling_speecht5.HifiGanResidualBlock
-class HifiGanResidualBlock(nn.Module):
- def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), leaky_relu_slope=0.1):
- super().__init__()
- self.leaky_relu_slope = leaky_relu_slope
-
- self.convs1 = nn.ModuleList(
- [
- nn.Conv1d(
- channels,
- channels,
- kernel_size,
- stride=1,
- dilation=dilation[i],
- padding=self.get_padding(kernel_size, dilation[i]),
- )
- for i in range(len(dilation))
- ]
- )
- self.convs2 = nn.ModuleList(
- [
- nn.Conv1d(
- channels,
- channels,
- kernel_size,
- stride=1,
- dilation=1,
- padding=self.get_padding(kernel_size, 1),
- )
- for _ in range(len(dilation))
- ]
- )
-
- def get_padding(self, kernel_size, dilation=1):
- return (kernel_size * dilation - dilation) // 2
-
- def apply_weight_norm(self):
- for layer in self.convs1:
- nn.utils.weight_norm(layer)
- for layer in self.convs2:
- nn.utils.weight_norm(layer)
-
- def remove_weight_norm(self):
- for layer in self.convs1:
- nn.utils.remove_weight_norm(layer)
- for layer in self.convs2:
- nn.utils.remove_weight_norm(layer)
-
- def forward(self, hidden_states):
- for conv1, conv2 in zip(self.convs1, self.convs2):
- residual = hidden_states
- hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
- hidden_states = conv1(hidden_states)
- hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
- hidden_states = conv2(hidden_states)
- hidden_states = hidden_states + residual
- return hidden_states
-
-
-@add_start_docstrings(
- """HiFi-GAN vocoder.""",
- HIFIGAN_START_DOCSTRING,
-)
-# Copied from transformers.models.speecht5.modeling_speecht5.SpeechT5HifiGan with SpeechT5->FastSpeech2Conformer
-class FastSpeech2ConformerHifiGan(PreTrainedModel):
- config_class = FastSpeech2ConformerHifiGanConfig
- main_input_name = "spectrogram"
-
- def __init__(self, config: FastSpeech2ConformerHifiGanConfig):
- super().__init__(config)
- self.num_kernels = len(config.resblock_kernel_sizes)
- self.num_upsamples = len(config.upsample_rates)
- self.conv_pre = nn.Conv1d(
- config.model_in_dim,
- config.upsample_initial_channel,
- kernel_size=7,
- stride=1,
- padding=3,
- )
-
- self.upsampler = nn.ModuleList()
- for i, (upsample_rate, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes)):
- self.upsampler.append(
- nn.ConvTranspose1d(
- config.upsample_initial_channel // (2**i),
- config.upsample_initial_channel // (2 ** (i + 1)),
- kernel_size=kernel_size,
- stride=upsample_rate,
- padding=(kernel_size - upsample_rate) // 2,
- )
- )
-
- self.resblocks = nn.ModuleList()
- for i in range(len(self.upsampler)):
- channels = config.upsample_initial_channel // (2 ** (i + 1))
- for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes):
- self.resblocks.append(HifiGanResidualBlock(channels, kernel_size, dilation, config.leaky_relu_slope))
-
- self.conv_post = nn.Conv1d(channels, 1, kernel_size=7, stride=1, padding=3)
-
- self.register_buffer("mean", torch.zeros(config.model_in_dim))
- self.register_buffer("scale", torch.ones(config.model_in_dim))
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, (nn.Linear, nn.Conv1d)):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
-
- def apply_weight_norm(self):
- nn.utils.weight_norm(self.conv_pre)
- for layer in self.upsampler:
- nn.utils.weight_norm(layer)
- for layer in self.resblocks:
- layer.apply_weight_norm()
- nn.utils.weight_norm(self.conv_post)
-
- def remove_weight_norm(self):
- nn.utils.remove_weight_norm(self.conv_pre)
- for layer in self.upsampler:
- nn.utils.remove_weight_norm(layer)
- for layer in self.resblocks:
- layer.remove_weight_norm()
- nn.utils.remove_weight_norm(self.conv_post)
-
- def forward(self, spectrogram: torch.FloatTensor) -> torch.FloatTensor:
- r"""
- Converts a log-mel spectrogram into a speech waveform. Passing a batch of log-mel spectrograms returns a batch
- of speech waveforms. Passing a single, un-batched log-mel spectrogram returns a single, un-batched speech
- waveform.
-
- Args:
- spectrogram (`torch.FloatTensor`):
- Tensor containing the log-mel spectrograms. Can be batched and of shape `(batch_size, sequence_length,
- config.model_in_dim)`, or un-batched and of shape `(sequence_length, config.model_in_dim)`.
-
- Returns:
- `torch.FloatTensor`: Tensor containing the speech waveform. If the input spectrogram is batched, will be of
- shape `(batch_size, num_frames,)`. If un-batched, will be of shape `(num_frames,)`.
- """
- if self.config.normalize_before:
- spectrogram = (spectrogram - self.mean) / self.scale
-
- is_batched = spectrogram.dim() == 3
- if not is_batched:
- spectrogram = spectrogram.unsqueeze(0)
-
- hidden_states = spectrogram.transpose(2, 1)
-
- hidden_states = self.conv_pre(hidden_states)
- for i in range(self.num_upsamples):
- hidden_states = nn.functional.leaky_relu(hidden_states, self.config.leaky_relu_slope)
- hidden_states = self.upsampler[i](hidden_states)
-
- res_state = self.resblocks[i * self.num_kernels](hidden_states)
- for j in range(1, self.num_kernels):
- res_state += self.resblocks[i * self.num_kernels + j](hidden_states)
- hidden_states = res_state / self.num_kernels
-
- hidden_states = nn.functional.leaky_relu(hidden_states)
- hidden_states = self.conv_post(hidden_states)
- hidden_states = torch.tanh(hidden_states)
-
- if not is_batched:
- # remove batch dim and collapse tensor to 1-d audio waveform
- waveform = hidden_states.squeeze(0).transpose(1, 0).view(-1)
- else:
- # remove seq-len dim since this collapses to 1
- waveform = hidden_states.squeeze(1)
-
- return waveform
-
-
-@add_start_docstrings(
- "The FastSpeech2ConformerModel with a FastSpeech2ConformerHifiGan vocoder head that performs text-to-speech (waveform).",
- FASTSPEECH2_CONFORMER_WITH_HIFIGAN_START_DOCSTRING,
-)
-class FastSpeech2ConformerWithHifiGan(PreTrainedModel):
- config_class = FastSpeech2ConformerWithHifiGanConfig
-
- def __init__(self, config: FastSpeech2ConformerWithHifiGanConfig):
- super().__init__(config)
-
- self.model = FastSpeech2ConformerModel(config.model_config)
- self.vocoder = FastSpeech2ConformerHifiGan(config.vocoder_config)
-
- self.config = config
-
- @replace_return_docstrings(
- output_type=FastSpeech2ConformerWithHifiGanOutput, config_class=FastSpeech2ConformerWithHifiGanConfig
- )
- def forward(
- self,
- input_ids: torch.LongTensor,
- attention_mask: Optional[torch.LongTensor] = None,
- spectrogram_labels: Optional[torch.FloatTensor] = None,
- duration_labels: Optional[torch.LongTensor] = None,
- pitch_labels: Optional[torch.FloatTensor] = None,
- energy_labels: Optional[torch.FloatTensor] = None,
- speaker_ids: Optional[torch.LongTensor] = None,
- lang_ids: Optional[torch.LongTensor] = None,
- speaker_embedding: Optional[torch.FloatTensor] = None,
- return_dict: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- ) -> Union[Tuple, FastSpeech2ConformerModelOutput]:
- """
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Input sequence of text vectors.
- attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*, defaults to `None`):
- Mask to avoid performing convolution and attention on padding token indices. Mask values selected in
- `[0, 1]`: 0 for tokens that are **masked**, 1 for tokens that are **not masked**.
- spectrogram_labels (`torch.FloatTensor` of shape `(batch_size, max_spectrogram_length, num_mel_bins)`, *optional*, defaults to `None`):
- Batch of padded target features.
- duration_labels (`torch.LongTensor` of shape `(batch_size, sequence_length + 1)`, *optional*, defaults to `None`):
- Batch of padded durations.
- pitch_labels (`torch.FloatTensor` of shape `(batch_size, sequence_length + 1, 1)`, *optional*, defaults to `None`):
- Batch of padded token-averaged pitch.
- energy_labels (`torch.FloatTensor` of shape `(batch_size, sequence_length + 1, 1)`, *optional*, defaults to `None`):
- Batch of padded token-averaged energy.
- speaker_ids (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*, defaults to `None`):
- Speaker ids used to condition features of speech output by the model.
- lang_ids (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*, defaults to `None`):
- Language ids used to condition features of speech output by the model.
- speaker_embedding (`torch.FloatTensor` of shape `(batch_size, embedding_dim)`, *optional*, defaults to `None`):
- Embedding containing conditioning signals for the features of the speech.
- return_dict (`bool`, *optional*, defaults to `None`):
- Whether or not to return a [`FastSpeech2ConformerModelOutput`] instead of a plain tuple.
- output_attentions (`bool`, *optional*, defaults to `None`):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*, defaults to `None`):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import (
- ... FastSpeech2ConformerTokenizer,
- ... FastSpeech2ConformerWithHifiGan,
- ... )
-
- >>> tokenizer = FastSpeech2ConformerTokenizer.from_pretrained("espnet/fastspeech2_conformer")
- >>> inputs = tokenizer("some text to convert to speech", return_tensors="pt")
- >>> input_ids = inputs["input_ids"]
-
- >>> model = FastSpeech2ConformerWithHifiGan.from_pretrained("espnet/fastspeech2_conformer_with_hifigan")
- >>> output_dict = model(input_ids, return_dict=True)
- >>> waveform = output_dict["waveform"]
- >>> print(waveform.shape)
- torch.Size([1, 49664])
- ```
- """
- return_dict = return_dict if return_dict is not None else self.config.model_config.use_return_dict
- output_attentions = (
- output_attentions if output_attentions is not None else self.config.model_config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.model_config.output_hidden_states
- )
-
- model_outputs = self.model(
- input_ids,
- attention_mask,
- spectrogram_labels=spectrogram_labels,
- duration_labels=duration_labels,
- pitch_labels=pitch_labels,
- energy_labels=energy_labels,
- speaker_ids=speaker_ids,
- lang_ids=lang_ids,
- speaker_embedding=speaker_embedding,
- return_dict=return_dict,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- )
-
- if not return_dict:
- has_missing_labels = (
- spectrogram_labels is None or duration_labels is None or pitch_labels is None or energy_labels is None
- )
- if has_missing_labels:
- spectrogram = model_outputs[0]
- else:
- spectrogram = model_outputs[1]
- else:
- spectrogram = model_outputs["spectrogram"]
- waveform = self.vocoder(spectrogram)
-
- if not return_dict:
- return model_outputs + (waveform,)
-
- return FastSpeech2ConformerWithHifiGanOutput(waveform=waveform, **model_outputs)
diff --git a/transformers/models/fastspeech2_conformer/tokenization_fastspeech2_conformer.py b/transformers/models/fastspeech2_conformer/tokenization_fastspeech2_conformer.py
deleted file mode 100644
index 5b979c8761c42cfffbb215578430d82596b961fc..0000000000000000000000000000000000000000
--- a/transformers/models/fastspeech2_conformer/tokenization_fastspeech2_conformer.py
+++ /dev/null
@@ -1,184 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Team and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for FastSpeech2Conformer."""
-import json
-import os
-from typing import Optional, Tuple
-
-import regex
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...utils import logging, requires_backends
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.json"}
-
-
-class FastSpeech2ConformerTokenizer(PreTrainedTokenizer):
- """
- Construct a FastSpeech2Conformer tokenizer.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- bos_token (`str`, *optional*, defaults to `""`):
- The begin of sequence token. Note that for FastSpeech2, it is the same as the `eos_token`.
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sequence token. Note that for FastSpeech2, it is the same as the `bos_token`.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- should_strip_spaces (`bool`, *optional*, defaults to `False`):
- Whether or not to strip the spaces from the list of tokens.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- bos_token="",
- eos_token="",
- pad_token="",
- unk_token="",
- should_strip_spaces=False,
- **kwargs,
- ):
- requires_backends(self, "g2p_en")
-
- with open(vocab_file, encoding="utf-8") as vocab_handle:
- self.encoder = json.load(vocab_handle)
-
- import g2p_en
-
- self.g2p = g2p_en.G2p()
-
- self.decoder = {v: k for k, v in self.encoder.items()}
-
- super().__init__(
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- pad_token=pad_token,
- should_strip_spaces=should_strip_spaces,
- **kwargs,
- )
-
- self.should_strip_spaces = should_strip_spaces
-
- @property
- def vocab_size(self):
- return len(self.decoder)
-
- def get_vocab(self):
- "Returns vocab as a dict"
- return dict(self.encoder, **self.added_tokens_encoder)
-
- def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
- # expand symbols
- text = regex.sub(";", ",", text)
- text = regex.sub(":", ",", text)
- text = regex.sub("-", " ", text)
- text = regex.sub("&", "and", text)
-
- # strip unnecessary symbols
- text = regex.sub(r"[\(\)\[\]\<\>\"]+", "", text)
-
- # strip whitespaces
- text = regex.sub(r"\s+", " ", text)
-
- text = text.upper()
-
- return text, kwargs
-
- def _tokenize(self, text):
- """Returns a tokenized string."""
- # phonemize
- tokens = self.g2p(text)
-
- if self.should_strip_spaces:
- tokens = list(filter(lambda s: s != " ", tokens))
-
- tokens.append(self.eos_token)
-
- return tokens
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index, self.unk_token)
-
- # Override since phonemes cannot be converted back to strings
- def decode(self, token_ids, **kwargs):
- logger.warning(
- "Phonemes cannot be reliably converted to a string due to the one-many mapping, converting to tokens instead."
- )
- return self.convert_ids_to_tokens(token_ids)
-
- # Override since phonemes cannot be converted back to strings
- def convert_tokens_to_string(self, tokens, **kwargs):
- logger.warning(
- "Phonemes cannot be reliably converted to a string due to the one-many mapping, returning the tokens."
- )
- return tokens
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- """
- Save the vocabulary and special tokens file to a directory.
-
- Args:
- save_directory (`str`):
- The directory in which to save the vocabulary.
-
- Returns:
- `Tuple(str)`: Paths to the files saved.
- """
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- with open(vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(self.get_vocab(), ensure_ascii=False))
-
- return (vocab_file,)
-
- def __getstate__(self):
- state = self.__dict__.copy()
- state["g2p"] = None
- return state
-
- def __setstate__(self, d):
- self.__dict__ = d
-
- try:
- import g2p_en
-
- self.g2p = g2p_en.G2p()
- except ImportError:
- raise ImportError(
- "You need to install g2p-en to use FastSpeech2ConformerTokenizer. "
- "See https://pypi.org/project/g2p-en/ for installation."
- )
diff --git a/transformers/models/flaubert/__init__.py b/transformers/models/flaubert/__init__.py
deleted file mode 100644
index 210d80b00f9ea2195b41bf5c6f3c0cd885fddae2..0000000000000000000000000000000000000000
--- a/transformers/models/flaubert/__init__.py
+++ /dev/null
@@ -1,103 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available
-
-
-_import_structure = {
- "configuration_flaubert": ["FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "FlaubertConfig", "FlaubertOnnxConfig"],
- "tokenization_flaubert": ["FlaubertTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flaubert"] = [
- "FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "FlaubertForMultipleChoice",
- "FlaubertForQuestionAnswering",
- "FlaubertForQuestionAnsweringSimple",
- "FlaubertForSequenceClassification",
- "FlaubertForTokenClassification",
- "FlaubertModel",
- "FlaubertWithLMHeadModel",
- "FlaubertPreTrainedModel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_flaubert"] = [
- "TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFFlaubertForMultipleChoice",
- "TFFlaubertForQuestionAnsweringSimple",
- "TFFlaubertForSequenceClassification",
- "TFFlaubertForTokenClassification",
- "TFFlaubertModel",
- "TFFlaubertPreTrainedModel",
- "TFFlaubertWithLMHeadModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_flaubert import FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, FlaubertConfig, FlaubertOnnxConfig
- from .tokenization_flaubert import FlaubertTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flaubert import (
- FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
- FlaubertForMultipleChoice,
- FlaubertForQuestionAnswering,
- FlaubertForQuestionAnsweringSimple,
- FlaubertForSequenceClassification,
- FlaubertForTokenClassification,
- FlaubertModel,
- FlaubertPreTrainedModel,
- FlaubertWithLMHeadModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_flaubert import (
- TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFFlaubertForMultipleChoice,
- TFFlaubertForQuestionAnsweringSimple,
- TFFlaubertForSequenceClassification,
- TFFlaubertForTokenClassification,
- TFFlaubertModel,
- TFFlaubertPreTrainedModel,
- TFFlaubertWithLMHeadModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/flaubert/__pycache__/__init__.cpython-310.pyc b/transformers/models/flaubert/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 703cd505ddceb82fb87e3c6fc5ed0a39a6a9b1b9..0000000000000000000000000000000000000000
Binary files a/transformers/models/flaubert/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flaubert/__pycache__/configuration_flaubert.cpython-310.pyc b/transformers/models/flaubert/__pycache__/configuration_flaubert.cpython-310.pyc
deleted file mode 100644
index 9e9ca2c1e954d33a1c66443a3733b30146cd8650..0000000000000000000000000000000000000000
Binary files a/transformers/models/flaubert/__pycache__/configuration_flaubert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flaubert/__pycache__/modeling_flaubert.cpython-310.pyc b/transformers/models/flaubert/__pycache__/modeling_flaubert.cpython-310.pyc
deleted file mode 100644
index 1bbc3272dd654ac1c6d18ca3ec4248aa4266cc8c..0000000000000000000000000000000000000000
Binary files a/transformers/models/flaubert/__pycache__/modeling_flaubert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flaubert/__pycache__/modeling_tf_flaubert.cpython-310.pyc b/transformers/models/flaubert/__pycache__/modeling_tf_flaubert.cpython-310.pyc
deleted file mode 100644
index ade2fb668369e7e729269dc514493a694bdc58e7..0000000000000000000000000000000000000000
Binary files a/transformers/models/flaubert/__pycache__/modeling_tf_flaubert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flaubert/__pycache__/tokenization_flaubert.cpython-310.pyc b/transformers/models/flaubert/__pycache__/tokenization_flaubert.cpython-310.pyc
deleted file mode 100644
index ed4faf957ee2f550d65530fdfb5fcd17de05d653..0000000000000000000000000000000000000000
Binary files a/transformers/models/flaubert/__pycache__/tokenization_flaubert.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flaubert/configuration_flaubert.py b/transformers/models/flaubert/configuration_flaubert.py
deleted file mode 100644
index fb4ef2992cbb8823e603989de42ff125f9c53854..0000000000000000000000000000000000000000
--- a/transformers/models/flaubert/configuration_flaubert.py
+++ /dev/null
@@ -1,234 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present CNRS, Facebook Inc. and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Flaubert configuration"""
-from collections import OrderedDict
-from typing import Mapping
-
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class FlaubertConfig(PretrainedConfig):
- """
- This is the configuration class to store the configuration of a [`FlaubertModel`] or a [`TFFlaubertModel`]. It is
- used to instantiate a FlauBERT model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the FlauBERT
- [flaubert/flaubert_base_uncased](https://huggingface.co/flaubert/flaubert_base_uncased) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- pre_norm (`bool`, *optional*, defaults to `False`):
- Whether to apply the layer normalization before or after the feed forward layer following the attention in
- each layer (Vaswani et al., Tensor2Tensor for Neural Machine Translation. 2018)
- layerdrop (`float`, *optional*, defaults to 0.0):
- Probability to drop layers during training (Fan et al., Reducing Transformer Depth on Demand with
- Structured Dropout. ICLR 2020)
- vocab_size (`int`, *optional*, defaults to 30145):
- Vocabulary size of the FlauBERT model. Defines the number of different tokens that can be represented by
- the `inputs_ids` passed when calling [`FlaubertModel`] or [`TFFlaubertModel`].
- emb_dim (`int`, *optional*, defaults to 2048):
- Dimensionality of the encoder layers and the pooler layer.
- n_layer (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- n_head (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for the attention mechanism
- gelu_activation (`bool`, *optional*, defaults to `True`):
- Whether or not to use a *gelu* activation instead of *relu*.
- sinusoidal_embeddings (`bool`, *optional*, defaults to `False`):
- Whether or not to use sinusoidal positional embeddings instead of absolute positional embeddings.
- causal (`bool`, *optional*, defaults to `False`):
- Whether or not the model should behave in a causal manner. Causal models use a triangular attention mask in
- order to only attend to the left-side context instead if a bidirectional context.
- asm (`bool`, *optional*, defaults to `False`):
- Whether or not to use an adaptive log softmax projection layer instead of a linear layer for the prediction
- layer.
- n_langs (`int`, *optional*, defaults to 1):
- The number of languages the model handles. Set to 1 for monolingual models.
- use_lang_emb (`bool`, *optional*, defaults to `True`)
- Whether to use language embeddings. Some models use additional language embeddings, see [the multilingual
- models page](http://huggingface.co/transformers/multilingual.html#xlm-language-embeddings) for information
- on how to use them.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- embed_init_std (`float`, *optional*, defaults to 2048^-0.5):
- The standard deviation of the truncated_normal_initializer for initializing the embedding matrices.
- init_std (`int`, *optional*, defaults to 50257):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices except the
- embedding matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- bos_index (`int`, *optional*, defaults to 0):
- The index of the beginning of sentence token in the vocabulary.
- eos_index (`int`, *optional*, defaults to 1):
- The index of the end of sentence token in the vocabulary.
- pad_index (`int`, *optional*, defaults to 2):
- The index of the padding token in the vocabulary.
- unk_index (`int`, *optional*, defaults to 3):
- The index of the unknown token in the vocabulary.
- mask_index (`int`, *optional*, defaults to 5):
- The index of the masking token in the vocabulary.
- is_encoder(`bool`, *optional*, defaults to `True`):
- Whether or not the initialized model should be a transformer encoder or decoder as seen in Vaswani et al.
- summary_type (`string`, *optional*, defaults to "first"):
- Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
-
- Has to be one of the following options:
-
- - `"last"`: Take the last token hidden state (like XLNet).
- - `"first"`: Take the first token hidden state (like BERT).
- - `"mean"`: Take the mean of all tokens hidden states.
- - `"cls_index"`: Supply a Tensor of classification token position (like GPT/GPT-2).
- - `"attn"`: Not implemented now, use multi-head attention.
- summary_use_proj (`bool`, *optional*, defaults to `True`):
- Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
-
- Whether or not to add a projection after the vector extraction.
- summary_activation (`str`, *optional*):
- Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
-
- Pass `"tanh"` for a tanh activation to the output, any other value will result in no activation.
- summary_proj_to_labels (`bool`, *optional*, defaults to `True`):
- Used in the sequence classification and multiple choice models.
-
- Whether the projection outputs should have `config.num_labels` or `config.hidden_size` classes.
- summary_first_dropout (`float`, *optional*, defaults to 0.1):
- Used in the sequence classification and multiple choice models.
-
- The dropout ratio to be used after the projection and activation.
- start_n_top (`int`, *optional*, defaults to 5):
- Used in the SQuAD evaluation script.
- end_n_top (`int`, *optional*, defaults to 5):
- Used in the SQuAD evaluation script.
- mask_token_id (`int`, *optional*, defaults to 0):
- Model agnostic parameter to identify masked tokens when generating text in an MLM context.
- lang_id (`int`, *optional*, defaults to 1):
- The ID of the language used by the model. This parameter is used when generating text in a given language.
- """
-
- model_type = "flaubert"
- attribute_map = {
- "hidden_size": "emb_dim",
- "num_attention_heads": "n_heads",
- "num_hidden_layers": "n_layers",
- "n_words": "vocab_size", # For backward compatibility
- }
-
- def __init__(
- self,
- pre_norm=False,
- layerdrop=0.0,
- vocab_size=30145,
- emb_dim=2048,
- n_layers=12,
- n_heads=16,
- dropout=0.1,
- attention_dropout=0.1,
- gelu_activation=True,
- sinusoidal_embeddings=False,
- causal=False,
- asm=False,
- n_langs=1,
- use_lang_emb=True,
- max_position_embeddings=512,
- embed_init_std=2048**-0.5,
- layer_norm_eps=1e-12,
- init_std=0.02,
- bos_index=0,
- eos_index=1,
- pad_index=2,
- unk_index=3,
- mask_index=5,
- is_encoder=True,
- summary_type="first",
- summary_use_proj=True,
- summary_activation=None,
- summary_proj_to_labels=True,
- summary_first_dropout=0.1,
- start_n_top=5,
- end_n_top=5,
- mask_token_id=0,
- lang_id=0,
- pad_token_id=2,
- bos_token_id=0,
- **kwargs,
- ):
- """Constructs FlaubertConfig."""
- self.pre_norm = pre_norm
- self.layerdrop = layerdrop
- self.vocab_size = vocab_size
- self.emb_dim = emb_dim
- self.n_layers = n_layers
- self.n_heads = n_heads
- self.dropout = dropout
- self.attention_dropout = attention_dropout
- self.gelu_activation = gelu_activation
- self.sinusoidal_embeddings = sinusoidal_embeddings
- self.causal = causal
- self.asm = asm
- self.n_langs = n_langs
- self.use_lang_emb = use_lang_emb
- self.layer_norm_eps = layer_norm_eps
- self.bos_index = bos_index
- self.eos_index = eos_index
- self.pad_index = pad_index
- self.unk_index = unk_index
- self.mask_index = mask_index
- self.is_encoder = is_encoder
- self.max_position_embeddings = max_position_embeddings
- self.embed_init_std = embed_init_std
- self.init_std = init_std
- self.summary_type = summary_type
- self.summary_use_proj = summary_use_proj
- self.summary_activation = summary_activation
- self.summary_proj_to_labels = summary_proj_to_labels
- self.summary_first_dropout = summary_first_dropout
- self.start_n_top = start_n_top
- self.end_n_top = end_n_top
- self.mask_token_id = mask_token_id
- self.lang_id = lang_id
-
- if "n_words" in kwargs:
- self.n_words = kwargs["n_words"]
-
- super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, **kwargs)
-
-
-class FlaubertOnnxConfig(OnnxConfig):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- if self.task == "multiple-choice":
- dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
- else:
- dynamic_axis = {0: "batch", 1: "sequence"}
- return OrderedDict(
- [
- ("input_ids", dynamic_axis),
- ("attention_mask", dynamic_axis),
- ]
- )
diff --git a/transformers/models/flaubert/modeling_flaubert.py b/transformers/models/flaubert/modeling_flaubert.py
deleted file mode 100644
index 49c2008cd10ac618e1fd66685e36a04fe07b5aaf..0000000000000000000000000000000000000000
--- a/transformers/models/flaubert/modeling_flaubert.py
+++ /dev/null
@@ -1,1302 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present CNRS, Facebook Inc. and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Flaubert model, based on XLM."""
-
-import itertools
-import math
-from dataclasses import dataclass
-from typing import Dict, Optional, Tuple, Union
-
-import numpy as np
-import torch
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import gelu
-from ...modeling_outputs import (
- BaseModelOutput,
- MaskedLMOutput,
- MultipleChoiceModelOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel, SequenceSummary, SQuADHead
-from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_flaubert import FlaubertConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "flaubert/flaubert_base_cased"
-_CONFIG_FOR_DOC = "FlaubertConfig"
-
-
-from ..deprecated._archive_maps import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.xlm.modeling_xlm.create_sinusoidal_embeddings
-def create_sinusoidal_embeddings(n_pos, dim, out):
- position_enc = np.array([[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)])
- out.requires_grad = False
- out[:, 0::2] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))
- out[:, 1::2] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))
- out.detach_()
-
-
-# Copied from transformers.models.xlm.modeling_xlm.get_masks
-def get_masks(slen, lengths, causal, padding_mask=None):
- """
- Generate hidden states mask, and optionally an attention mask.
- """
- alen = torch.arange(slen, dtype=torch.long, device=lengths.device)
- if padding_mask is not None:
- mask = padding_mask
- else:
- assert lengths.max().item() <= slen
- mask = alen < lengths[:, None]
-
- # attention mask is the same as mask, or triangular inferior attention (causal)
- bs = lengths.size(0)
- if causal:
- attn_mask = alen[None, None, :].repeat(bs, slen, 1) <= alen[None, :, None]
- else:
- attn_mask = mask
-
- # sanity check
- assert mask.size() == (bs, slen)
- assert causal is False or attn_mask.size() == (bs, slen, slen)
-
- return mask, attn_mask
-
-
-# Copied from transformers.models.xlm.modeling_xlm.MultiHeadAttention
-class MultiHeadAttention(nn.Module):
- NEW_ID = itertools.count()
-
- def __init__(self, n_heads, dim, config):
- super().__init__()
- self.layer_id = next(MultiHeadAttention.NEW_ID)
- self.dim = dim
- self.n_heads = n_heads
- self.dropout = config.attention_dropout
- assert self.dim % self.n_heads == 0
-
- self.q_lin = nn.Linear(dim, dim)
- self.k_lin = nn.Linear(dim, dim)
- self.v_lin = nn.Linear(dim, dim)
- self.out_lin = nn.Linear(dim, dim)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- attention_head_size = self.dim // self.n_heads
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(heads, self.n_heads, attention_head_size, self.pruned_heads)
- # Prune linear layers
- self.q_lin = prune_linear_layer(self.q_lin, index)
- self.k_lin = prune_linear_layer(self.k_lin, index)
- self.v_lin = prune_linear_layer(self.v_lin, index)
- self.out_lin = prune_linear_layer(self.out_lin, index, dim=1)
- # Update hyper params
- self.n_heads = self.n_heads - len(heads)
- self.dim = attention_head_size * self.n_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(self, input, mask, kv=None, cache=None, head_mask=None, output_attentions=False):
- """
- Self-attention (if kv is None) or attention over source sentence (provided by kv).
- """
- # Input is (bs, qlen, dim)
- # Mask is (bs, klen) (non-causal) or (bs, klen, klen)
- bs, qlen, dim = input.size()
- if kv is None:
- klen = qlen if cache is None else cache["slen"] + qlen
- else:
- klen = kv.size(1)
- # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
- n_heads = self.n_heads
- dim_per_head = self.dim // n_heads
- mask_reshape = (bs, 1, qlen, klen) if mask.dim() == 3 else (bs, 1, 1, klen)
-
- def shape(x):
- """projection"""
- return x.view(bs, -1, self.n_heads, dim_per_head).transpose(1, 2)
-
- def unshape(x):
- """compute context"""
- return x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head)
-
- q = shape(self.q_lin(input)) # (bs, n_heads, qlen, dim_per_head)
- if kv is None:
- k = shape(self.k_lin(input)) # (bs, n_heads, qlen, dim_per_head)
- v = shape(self.v_lin(input)) # (bs, n_heads, qlen, dim_per_head)
- elif cache is None or self.layer_id not in cache:
- k = v = kv
- k = shape(self.k_lin(k)) # (bs, n_heads, qlen, dim_per_head)
- v = shape(self.v_lin(v)) # (bs, n_heads, qlen, dim_per_head)
-
- if cache is not None:
- if self.layer_id in cache:
- if kv is None:
- k_, v_ = cache[self.layer_id]
- k = torch.cat([k_, k], dim=2) # (bs, n_heads, klen, dim_per_head)
- v = torch.cat([v_, v], dim=2) # (bs, n_heads, klen, dim_per_head)
- else:
- k, v = cache[self.layer_id]
- cache[self.layer_id] = (k, v)
-
- q = q / math.sqrt(dim_per_head) # (bs, n_heads, qlen, dim_per_head)
- scores = torch.matmul(q, k.transpose(2, 3)) # (bs, n_heads, qlen, klen)
- mask = (mask == 0).view(mask_reshape).expand_as(scores) # (bs, n_heads, qlen, klen)
- scores.masked_fill_(mask, torch.finfo(scores.dtype).min) # (bs, n_heads, qlen, klen)
-
- weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores) # (bs, n_heads, qlen, klen)
- weights = nn.functional.dropout(weights, p=self.dropout, training=self.training) # (bs, n_heads, qlen, klen)
-
- # Mask heads if we want to
- if head_mask is not None:
- weights = weights * head_mask
-
- context = torch.matmul(weights, v) # (bs, n_heads, qlen, dim_per_head)
- context = unshape(context) # (bs, qlen, dim)
-
- outputs = (self.out_lin(context),)
- if output_attentions:
- outputs = outputs + (weights,)
- return outputs
-
-
-# Copied from transformers.models.xlm.modeling_xlm.TransformerFFN
-class TransformerFFN(nn.Module):
- def __init__(self, in_dim, dim_hidden, out_dim, config):
- super().__init__()
- self.dropout = config.dropout
- self.lin1 = nn.Linear(in_dim, dim_hidden)
- self.lin2 = nn.Linear(dim_hidden, out_dim)
- self.act = gelu if config.gelu_activation else nn.functional.relu
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
-
- def forward(self, input):
- return apply_chunking_to_forward(self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, input)
-
- def ff_chunk(self, input):
- x = self.lin1(input)
- x = self.act(x)
- x = self.lin2(x)
- x = nn.functional.dropout(x, p=self.dropout, training=self.training)
- return x
-
-
-FLAUBERT_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`FlaubertConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FLAUBERT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Length of each sentence that can be used to avoid performing attention on padding token indices. You can
- also use `attention_mask` for the same result (see above), kept here for compatibility. Indices selected in
- `[0, ..., input_ids.size(-1)]`:
- cache (`Dict[str, torch.FloatTensor]`, *optional*):
- Dictionary strings to `torch.FloatTensor` that contains precomputed hidden-states (key and values in the
- attention blocks) as computed by the model (see `cache` output below). Can be used to speed up sequential
- decoding. The dictionary object will be modified in-place during the forward pass to add newly computed
- hidden-states.
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare Flaubert Model transformer outputting raw hidden-states without any specific head on top.",
- FLAUBERT_START_DOCSTRING,
-)
-# Copied from transformers.models.xlm.modeling_xlm.XLMPredLayer with XLM->Flaubert
-class FlaubertPredLayer(nn.Module):
- """
- Prediction layer (cross_entropy or adaptive_softmax).
- """
-
- def __init__(self, config):
- super().__init__()
- self.asm = config.asm
- self.n_words = config.n_words
- self.pad_index = config.pad_index
- dim = config.emb_dim
-
- if config.asm is False:
- self.proj = nn.Linear(dim, config.n_words, bias=True)
- else:
- self.proj = nn.AdaptiveLogSoftmaxWithLoss(
- in_features=dim,
- n_classes=config.n_words,
- cutoffs=config.asm_cutoffs,
- div_value=config.asm_div_value,
- head_bias=True, # default is False
- )
-
- def forward(self, x, y=None):
- """Compute the loss, and optionally the scores."""
- outputs = ()
- if self.asm is False:
- scores = self.proj(x)
- outputs = (scores,) + outputs
- if y is not None:
- loss = nn.functional.cross_entropy(scores.view(-1, self.n_words), y.view(-1), reduction="mean")
- outputs = (loss,) + outputs
- else:
- scores = self.proj.log_prob(x)
- outputs = (scores,) + outputs
- if y is not None:
- _, loss = self.proj(x, y)
- outputs = (loss,) + outputs
-
- return outputs
-
-
-# Copied from transformers.models.xlm.modeling_xlm.XLMPreTrainedModel with XLM->Flaubert
-class FlaubertPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = FlaubertConfig
- load_tf_weights = None
- base_model_prefix = "transformer"
-
- def __init__(self, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- @property
- def dummy_inputs(self):
- inputs_list = torch.tensor([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]])
- attns_list = torch.tensor([[1, 1, 0, 0, 1], [1, 1, 1, 0, 0], [1, 0, 0, 1, 1]])
- if self.config.use_lang_emb and self.config.n_langs > 1:
- langs_list = torch.tensor([[1, 1, 0, 0, 1], [1, 1, 1, 0, 0], [1, 0, 0, 1, 1]])
- else:
- langs_list = None
- return {"input_ids": inputs_list, "attention_mask": attns_list, "langs": langs_list}
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, nn.Embedding):
- if self.config is not None and self.config.embed_init_std is not None:
- nn.init.normal_(module.weight, mean=0, std=self.config.embed_init_std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- if isinstance(module, nn.Linear):
- if self.config is not None and self.config.init_std is not None:
- nn.init.normal_(module.weight, mean=0, std=self.config.init_std)
- if module.bias is not None:
- nn.init.constant_(module.bias, 0.0)
- if isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
- if isinstance(module, FlaubertModel) and self.config.sinusoidal_embeddings:
- create_sinusoidal_embeddings(
- self.config.max_position_embeddings, self.config.emb_dim, out=module.position_embeddings.weight
- )
-
-
-class FlaubertModel(FlaubertPreTrainedModel):
- def __init__(self, config): # , dico, is_encoder, with_output):
- super().__init__(config)
-
- # encoder / decoder, output layer
- self.is_encoder = config.is_encoder
- self.is_decoder = not config.is_encoder
- if self.is_decoder:
- raise NotImplementedError("Currently Flaubert can only be used as an encoder")
- # self.with_output = with_output
- self.causal = config.causal
-
- # dictionary / languages
- self.n_langs = config.n_langs
- self.use_lang_emb = config.use_lang_emb
- self.n_words = config.n_words
- self.eos_index = config.eos_index
- self.pad_index = config.pad_index
- # self.dico = dico
- # self.id2lang = config.id2lang
- # self.lang2id = config.lang2id
- # assert len(self.dico) == self.n_words
- # assert len(self.id2lang) == len(self.lang2id) == self.n_langs
-
- # model parameters
- self.dim = config.emb_dim # 512 by default
- self.hidden_dim = self.dim * 4 # 2048 by default
- self.n_heads = config.n_heads # 8 by default
- self.n_layers = config.n_layers
- self.dropout = config.dropout
- self.attention_dropout = config.attention_dropout
- assert self.dim % self.n_heads == 0, "transformer dim must be a multiple of n_heads"
-
- # embeddings
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.dim)
- if config.n_langs > 1 and config.use_lang_emb:
- self.lang_embeddings = nn.Embedding(self.n_langs, self.dim)
- self.embeddings = nn.Embedding(self.n_words, self.dim, padding_idx=self.pad_index)
- self.layer_norm_emb = nn.LayerNorm(self.dim, eps=config.layer_norm_eps)
-
- # transformer layers
- self.attentions = nn.ModuleList()
- self.layer_norm1 = nn.ModuleList()
- self.ffns = nn.ModuleList()
- self.layer_norm2 = nn.ModuleList()
- # if self.is_decoder:
- # self.layer_norm15 = nn.ModuleList()
- # self.encoder_attn = nn.ModuleList()
-
- for _ in range(self.n_layers):
- self.attentions.append(MultiHeadAttention(self.n_heads, self.dim, config=config))
- self.layer_norm1.append(nn.LayerNorm(self.dim, eps=config.layer_norm_eps))
- # if self.is_decoder:
- # self.layer_norm15.append(nn.LayerNorm(self.dim, eps=config.layer_norm_eps))
- # self.encoder_attn.append(MultiHeadAttention(self.n_heads, self.dim, dropout=self.attention_dropout))
- self.ffns.append(TransformerFFN(self.dim, self.hidden_dim, self.dim, config=config))
- self.layer_norm2.append(nn.LayerNorm(self.dim, eps=config.layer_norm_eps))
-
- if hasattr(config, "pruned_heads"):
- pruned_heads = config.pruned_heads.copy().items()
- config.pruned_heads = {}
- for layer, heads in pruned_heads:
- if self.attentions[int(layer)].n_heads == config.n_heads:
- self.prune_heads({int(layer): list(map(int, heads))})
-
- # Initialize weights and apply final processing
- self.post_init()
-
- self.layerdrop = getattr(config, "layerdrop", 0.0)
- self.pre_norm = getattr(config, "pre_norm", False)
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
-
- # Copied from transformers.models.xlm.modeling_xlm.XLMModel.get_input_embeddings
- def get_input_embeddings(self):
- return self.embeddings
-
- # Copied from transformers.models.xlm.modeling_xlm.XLMModel.set_input_embeddings
- def set_input_embeddings(self, new_embeddings):
- self.embeddings = new_embeddings
-
- # Copied from transformers.models.xlm.modeling_xlm.XLMModel._prune_heads
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.attentions[layer].prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- langs: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- lengths: Optional[torch.LongTensor] = None,
- cache: Optional[Dict[str, torch.FloatTensor]] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # removed: src_enc=None, src_len=None
- if input_ids is not None:
- bs, slen = input_ids.size()
- else:
- bs, slen = inputs_embeds.size()[:-1]
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if lengths is None:
- if input_ids is not None:
- lengths = (input_ids != self.pad_index).sum(dim=1).long()
- else:
- lengths = torch.tensor([slen] * bs, device=device)
- # mask = input_ids != self.pad_index
-
- # check inputs
- assert lengths.size(0) == bs
- assert lengths.max().item() <= slen
- # input_ids = input_ids.transpose(0, 1) # batch size as dimension 0
- # assert (src_enc is None) == (src_len is None)
- # if src_enc is not None:
- # assert self.is_decoder
- # assert src_enc.size(0) == bs
-
- # generate masks
- mask, attn_mask = get_masks(slen, lengths, self.causal, padding_mask=attention_mask)
- # if self.is_decoder and src_enc is not None:
- # src_mask = torch.arange(src_len.max(), dtype=torch.long, device=lengths.device) < src_len[:, None]
-
- # Setting the position-ids to the registered buffer in constructor, it helps
- # when tracing the model without passing position-ids, solves
- # isues similar to issue #5664
- if position_ids is None:
- if hasattr(self, "position_ids"):
- position_ids = self.position_ids[:, :slen]
- position_ids = position_ids.expand((bs, slen))
- else:
- position_ids = torch.arange(slen, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0).expand((bs, slen))
- else:
- assert position_ids.size() == (bs, slen) # (slen, bs)
- # position_ids = position_ids.transpose(0, 1)
-
- # langs
- if langs is not None:
- assert langs.size() == (bs, slen) # (slen, bs)
- # langs = langs.transpose(0, 1)
-
- # Prepare head mask if needed
- head_mask = self.get_head_mask(head_mask, self.config.n_layers)
-
- # do not recompute cached elements
- if cache is not None and input_ids is not None:
- _slen = slen - cache["slen"]
- input_ids = input_ids[:, -_slen:]
- position_ids = position_ids[:, -_slen:]
- if langs is not None:
- langs = langs[:, -_slen:]
- mask = mask[:, -_slen:]
- attn_mask = attn_mask[:, -_slen:]
-
- # embeddings
- if inputs_embeds is None:
- inputs_embeds = self.embeddings(input_ids)
-
- tensor = inputs_embeds + self.position_embeddings(position_ids).expand_as(inputs_embeds)
- if langs is not None and self.use_lang_emb and self.config.n_langs > 1:
- tensor = tensor + self.lang_embeddings(langs)
- if token_type_ids is not None:
- tensor = tensor + self.embeddings(token_type_ids)
- tensor = self.layer_norm_emb(tensor)
- tensor = nn.functional.dropout(tensor, p=self.dropout, training=self.training)
- tensor *= mask.unsqueeze(-1).to(tensor.dtype)
-
- # transformer layers
- hidden_states = () if output_hidden_states else None
- attentions = () if output_attentions else None
- for i in range(self.n_layers):
- # LayerDrop
- if self.training:
- dropout_probability = torch.rand([])
- if dropout_probability < self.layerdrop:
- continue
-
- if output_hidden_states:
- hidden_states = hidden_states + (tensor,)
-
- # self attention
- if not self.pre_norm:
- attn_outputs = self.attentions[i](
- tensor,
- attn_mask,
- cache=cache,
- head_mask=head_mask[i],
- output_attentions=output_attentions,
- )
- attn = attn_outputs[0]
- if output_attentions:
- attentions = attentions + (attn_outputs[1],)
- attn = nn.functional.dropout(attn, p=self.dropout, training=self.training)
- tensor = tensor + attn
- tensor = self.layer_norm1[i](tensor)
- else:
- tensor_normalized = self.layer_norm1[i](tensor)
- attn_outputs = self.attentions[i](tensor_normalized, attn_mask, cache=cache, head_mask=head_mask[i])
- attn = attn_outputs[0]
- if output_attentions:
- attentions = attentions + (attn_outputs[1],)
- attn = nn.functional.dropout(attn, p=self.dropout, training=self.training)
- tensor = tensor + attn
-
- # encoder attention (for decoder only)
- # if self.is_decoder and src_enc is not None:
- # attn = self.encoder_attn[i](tensor, src_mask, kv=src_enc, cache=cache)
- # attn = nn.functional.dropout(attn, p=self.dropout, training=self.training)
- # tensor = tensor + attn
- # tensor = self.layer_norm15[i](tensor)
-
- # FFN
- if not self.pre_norm:
- tensor = tensor + self.ffns[i](tensor)
- tensor = self.layer_norm2[i](tensor)
- else:
- tensor_normalized = self.layer_norm2[i](tensor)
- tensor = tensor + self.ffns[i](tensor_normalized)
-
- tensor *= mask.unsqueeze(-1).to(tensor.dtype)
-
- # Add last hidden state
- if output_hidden_states:
- hidden_states = hidden_states + (tensor,)
-
- # update cache length
- if cache is not None:
- cache["slen"] += tensor.size(1)
-
- # move back sequence length to dimension 0
- # tensor = tensor.transpose(0, 1)
-
- if not return_dict:
- return tuple(v for v in [tensor, hidden_states, attentions] if v is not None)
-
- return BaseModelOutput(last_hidden_state=tensor, hidden_states=hidden_states, attentions=attentions)
-
-
-@add_start_docstrings(
- """
- The Flaubert Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied transformers.models.xlm.modeling_xlm.XLMWithLMHeadModel with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class FlaubertWithLMHeadModel(FlaubertPreTrainedModel):
- _tied_weights_keys = ["pred_layer.proj.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.transformer = FlaubertModel(config)
- self.pred_layer = FlaubertPredLayer(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.pred_layer.proj
-
- def set_output_embeddings(self, new_embeddings):
- self.pred_layer.proj = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, **kwargs):
- mask_token_id = self.config.mask_token_id
- lang_id = self.config.lang_id
-
- effective_batch_size = input_ids.shape[0]
- mask_token = torch.full((effective_batch_size, 1), mask_token_id, dtype=torch.long, device=input_ids.device)
- input_ids = torch.cat([input_ids, mask_token], dim=1)
- if lang_id is not None:
- langs = torch.full_like(input_ids, lang_id)
- else:
- langs = None
- return {"input_ids": input_ids, "langs": langs}
-
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="",
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- langs: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- lengths: Optional[torch.Tensor] = None,
- cache: Optional[Dict[str, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- output = transformer_outputs[0]
- outputs = self.pred_layer(output, labels) # (loss, logits) or (logits,) depending on if labels are provided.
-
- if not return_dict:
- return outputs + transformer_outputs[1:]
-
- return MaskedLMOutput(
- loss=outputs[0] if labels is not None else None,
- logits=outputs[0] if labels is None else outputs[1],
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a sequence classification/regression head on top (a linear layer on top of the pooled output)
- e.g. for GLUE tasks.
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied transformers.models.xlm.modeling_xlm.XLMForSequenceClassification with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class FlaubertForSequenceClassification(FlaubertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
-
- self.transformer = FlaubertModel(config)
- self.sequence_summary = SequenceSummary(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- langs: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- lengths: Optional[torch.Tensor] = None,
- cache: Optional[Dict[str, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- output = transformer_outputs[0]
- logits = self.sequence_summary(output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied from transformers.models.xlm.modeling_xlm.XLMForTokenClassification with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class FlaubertForTokenClassification(FlaubertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.transformer = FlaubertModel(config)
- self.dropout = nn.Dropout(config.dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- langs: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- lengths: Optional[torch.Tensor] = None,
- cache: Optional[Dict[str, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied from transformers.models.xlm.modeling_xlm.XLMForQuestionAnsweringSimple with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class FlaubertForQuestionAnsweringSimple(FlaubertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.transformer = FlaubertModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- langs: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- lengths: Optional[torch.Tensor] = None,
- cache: Optional[Dict[str, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = transformer_outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + transformer_outputs[1:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a beam-search span classification head on top for extractive question-answering tasks like
- SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- FLAUBERT_START_DOCSTRING,
-)
-@dataclass
-# Copied from transformer.models.xlm.modeling_xlm.XLMForQuestionAnsweringOutput with XLM->Flaubert
-class FlaubertForQuestionAnsweringOutput(ModelOutput):
- """
- Base class for outputs of question answering models using a `SquadHead`.
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned if both `start_positions` and `end_positions` are provided):
- Classification loss as the sum of start token, end token (and is_impossible if provided) classification
- losses.
- start_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
- Log probabilities for the top config.start_n_top start token possibilities (beam-search).
- start_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
- Indices for the top config.start_n_top start token possibilities (beam-search).
- end_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
- Log probabilities for the top `config.start_n_top * config.end_n_top` end token possibilities
- (beam-search).
- end_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
- Indices for the top `config.start_n_top * config.end_n_top` end token possibilities (beam-search).
- cls_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
- Log probabilities for the `is_impossible` label of the answers.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- start_top_log_probs: Optional[torch.FloatTensor] = None
- start_top_index: Optional[torch.LongTensor] = None
- end_top_log_probs: Optional[torch.FloatTensor] = None
- end_top_index: Optional[torch.LongTensor] = None
- cls_logits: Optional[torch.FloatTensor] = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-# Copied from transformer.models.xlm.modeling_xlm.XLMForQuestionAnswering with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class FlaubertForQuestionAnswering(FlaubertPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.transformer = FlaubertModel(config)
- self.qa_outputs = SQuADHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=FlaubertForQuestionAnsweringOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- langs: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- lengths: Optional[torch.Tensor] = None,
- cache: Optional[Dict[str, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- is_impossible: Optional[torch.Tensor] = None,
- cls_index: Optional[torch.Tensor] = None,
- p_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, FlaubertForQuestionAnsweringOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels whether a question has an answer or no answer (SQuAD 2.0)
- cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the classification token to use as input for computing plausibility of the
- answer.
- p_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Optional mask of tokens which can't be in answers (e.g. [CLS], [PAD], ...). 1.0 means token should be
- masked. 0.0 mean token is not masked.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import XLMTokenizer, XLMForQuestionAnswering
- >>> import torch
-
- >>> tokenizer = XLMTokenizer.from_pretrained("FacebookAI/xlm-mlm-en-2048")
- >>> model = XLMForQuestionAnswering.from_pretrained("FacebookAI/xlm-mlm-en-2048")
-
- >>> input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(
- ... 0
- ... ) # Batch size 1
- >>> start_positions = torch.tensor([1])
- >>> end_positions = torch.tensor([3])
-
- >>> outputs = model(input_ids, start_positions=start_positions, end_positions=end_positions)
- >>> loss = outputs.loss
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- output = transformer_outputs[0]
-
- outputs = self.qa_outputs(
- output,
- start_positions=start_positions,
- end_positions=end_positions,
- cls_index=cls_index,
- is_impossible=is_impossible,
- p_mask=p_mask,
- return_dict=return_dict,
- )
-
- if not return_dict:
- return outputs + transformer_outputs[1:]
-
- return FlaubertForQuestionAnsweringOutput(
- loss=outputs.loss,
- start_top_log_probs=outputs.start_top_log_probs,
- start_top_index=outputs.start_top_index,
- end_top_log_probs=outputs.end_top_log_probs,
- end_top_index=outputs.end_top_index,
- cls_logits=outputs.cls_logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied from transformer.models.xlm.modeling_xlm.XLMForMultipleChoice with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class FlaubertForMultipleChoice(FlaubertPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.transformer = FlaubertModel(config)
- self.sequence_summary = SequenceSummary(config)
- self.logits_proj = nn.Linear(config.num_labels, 1)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(
- FLAUBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
- )
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- langs: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- lengths: Optional[torch.Tensor] = None,
- cache: Optional[Dict[str, torch.Tensor]] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
- position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
- langs = langs.view(-1, langs.size(-1)) if langs is not None else None
- inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- if lengths is not None:
- logger.warning(
- "The `lengths` parameter cannot be used with the Flaubert multiple choice models. Please use the "
- "attention mask instead."
- )
- lengths = None
-
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- output = transformer_outputs[0]
- logits = self.sequence_summary(output)
- logits = self.logits_proj(logits)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
diff --git a/transformers/models/flaubert/modeling_tf_flaubert.py b/transformers/models/flaubert/modeling_tf_flaubert.py
deleted file mode 100644
index 08e573daa99458280037cc8dd142f0540b295a91..0000000000000000000000000000000000000000
--- a/transformers/models/flaubert/modeling_tf_flaubert.py
+++ /dev/null
@@ -1,1337 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present, Facebook, Inc and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
- TF 2.0 Flaubert model.
-"""
-
-
-from __future__ import annotations
-
-import itertools
-import random
-import warnings
-from dataclasses import dataclass
-from typing import Dict, Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFMultipleChoiceModelOutput,
- TFQuestionAnsweringModelOutput,
- TFSequenceClassifierOutput,
- TFTokenClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFModelInputType,
- TFMultipleChoiceLoss,
- TFPreTrainedModel,
- TFQuestionAnsweringLoss,
- TFSequenceClassificationLoss,
- TFSequenceSummary,
- TFSharedEmbeddings,
- TFTokenClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import (
- MULTIPLE_CHOICE_DUMMY_INPUTS,
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
-)
-from .configuration_flaubert import FlaubertConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "flaubert/flaubert_base_cased"
-_CONFIG_FOR_DOC = "FlaubertConfig"
-
-
-from ..deprecated._archive_maps import TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-FLAUBERT_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`FlaubertConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FLAUBERT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`Numpy array` or `tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - `1` for tokens that are **not masked**,
- - `0` for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- langs (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- A parallel sequence of tokens to be used to indicate the language of each token in the input. Indices are
- languages ids which can be obtained from the language names by using two conversion mappings provided in
- the configuration of the model (only provided for multilingual models). More precisely, the *language name
- to language id* mapping is in `model.config.lang2id` (which is a dictionary string to int) and the
- *language id to language name* mapping is in `model.config.id2lang` (dictionary int to string).
-
- See usage examples detailed in the [multilingual documentation](../multilingual).
- token_type_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - `0` corresponds to a *sentence A* token,
- - `1` corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- lengths (`tf.Tensor` or `Numpy array` of shape `(batch_size,)`, *optional*):
- Length of each sentence that can be used to avoid performing attention on padding token indices. You can
- also use *attention_mask* for the same result (see above), kept here for compatibility Indices selected in
- `[0, ..., input_ids.size(-1)]`:
- cache (`Dict[str, tf.Tensor]`, *optional*):
- Dictionary string to `tf.FloatTensor` that contains precomputed hidden states (key and values in the
- attention blocks) as computed by the model (see `cache` output below). Can be used to speed up sequential
- decoding.
-
- The dictionary object will be modified in-place during the forward pass to add newly computed
- hidden-states.
- head_mask (`Numpy array` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - `1` indicates the head is **not masked**,
- - `0` indicates the head is **masked**.
-
- inputs_embeds (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-def get_masks(slen, lengths, causal, padding_mask=None):
- """
- Generate hidden states mask, and optionally an attention mask.
- """
- bs = shape_list(lengths)[0]
- if padding_mask is not None:
- mask = padding_mask
- else:
- # assert lengths.max().item() <= slen
- alen = tf.range(slen, dtype=lengths.dtype)
- mask = alen < tf.expand_dims(lengths, axis=1)
-
- # attention mask is the same as mask, or triangular inferior attention (causal)
- if causal:
- attn_mask = tf.less_equal(
- tf.tile(tf.reshape(alen, (1, 1, slen)), (bs, slen, 1)), tf.reshape(alen, (1, slen, 1))
- )
- else:
- attn_mask = mask
-
- # sanity check
- # assert shape_list(mask) == [bs, slen]
- tf.debugging.assert_equal(shape_list(mask), [bs, slen])
- if causal:
- tf.debugging.assert_equal(shape_list(attn_mask), [bs, slen, slen])
-
- return mask, attn_mask
-
-
-class TFFlaubertPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = FlaubertConfig
- base_model_prefix = "transformer"
-
- @property
- def dummy_inputs(self):
- # Sometimes Flaubert has language embeddings so don't forget to build them as well if needed
- inputs_list = tf.constant([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]], dtype=tf.int32)
- attns_list = tf.constant([[1, 1, 0, 0, 1], [1, 1, 1, 0, 0], [1, 0, 0, 1, 1]], dtype=tf.int32)
- if self.config.use_lang_emb and self.config.n_langs > 1:
- return {
- "input_ids": inputs_list,
- "attention_mask": attns_list,
- "langs": tf.constant([[1, 1, 0, 0, 1], [1, 1, 1, 0, 0], [1, 0, 0, 1, 1]], dtype=tf.int32),
- }
- else:
- return {"input_ids": inputs_list, "attention_mask": attns_list}
-
-
-@add_start_docstrings(
- "The bare Flaubert Model transformer outputting raw hidden-states without any specific head on top.",
- FLAUBERT_START_DOCSTRING,
-)
-class TFFlaubertModel(TFFlaubertPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFFlaubertMainLayer(config, name="transformer")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: np.ndarray | tf.Tensor | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- langs: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- lengths: np.ndarray | tf.Tensor | None = None,
- cache: Optional[Dict[str, tf.Tensor]] = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFBaseModelOutput]:
- outputs = self.transformer(
- input_ids=input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
-
-
-# Copied from transformers.models.xlm.modeling_tf_xlm.TFXLMMultiHeadAttention with XLM->Flaubert
-class TFFlaubertMultiHeadAttention(keras.layers.Layer):
- NEW_ID = itertools.count()
-
- def __init__(self, n_heads, dim, config, **kwargs):
- super().__init__(**kwargs)
- self.layer_id = next(TFFlaubertMultiHeadAttention.NEW_ID)
- self.dim = dim
- self.n_heads = n_heads
- self.output_attentions = config.output_attentions
- assert self.dim % self.n_heads == 0
-
- self.q_lin = keras.layers.Dense(dim, kernel_initializer=get_initializer(config.init_std), name="q_lin")
- self.k_lin = keras.layers.Dense(dim, kernel_initializer=get_initializer(config.init_std), name="k_lin")
- self.v_lin = keras.layers.Dense(dim, kernel_initializer=get_initializer(config.init_std), name="v_lin")
- self.out_lin = keras.layers.Dense(dim, kernel_initializer=get_initializer(config.init_std), name="out_lin")
- self.dropout = keras.layers.Dropout(config.attention_dropout)
- self.pruned_heads = set()
- self.dim = dim
-
- def prune_heads(self, heads):
- raise NotImplementedError
-
- def call(self, input, mask, kv, cache, head_mask, output_attentions, training=False):
- """
- Self-attention (if kv is None) or attention over source sentence (provided by kv).
- """
- # Input is (bs, qlen, dim)
- # Mask is (bs, klen) (non-causal) or (bs, klen, klen)
- bs, qlen, dim = shape_list(input)
-
- if kv is None:
- klen = qlen if cache is None else cache["slen"] + qlen
- else:
- klen = shape_list(kv)[1]
-
- # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
- dim_per_head = self.dim // self.n_heads
- mask_reshape = (bs, 1, qlen, klen) if len(shape_list(mask)) == 3 else (bs, 1, 1, klen)
-
- def shape(x):
- """projection"""
- return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))
-
- def unshape(x):
- """compute context"""
- return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))
-
- q = shape(self.q_lin(input)) # (bs, n_heads, qlen, dim_per_head)
-
- if kv is None:
- k = shape(self.k_lin(input)) # (bs, n_heads, qlen, dim_per_head)
- v = shape(self.v_lin(input)) # (bs, n_heads, qlen, dim_per_head)
- elif cache is None or self.layer_id not in cache:
- k = v = kv
- k = shape(self.k_lin(k)) # (bs, n_heads, qlen, dim_per_head)
- v = shape(self.v_lin(v)) # (bs, n_heads, qlen, dim_per_head)
-
- if cache is not None:
- if self.layer_id in cache:
- if kv is None:
- k_, v_ = cache[self.layer_id]
- k = tf.concat([k_, k], axis=2) # (bs, n_heads, klen, dim_per_head)
- v = tf.concat([v_, v], axis=2) # (bs, n_heads, klen, dim_per_head)
- else:
- k, v = cache[self.layer_id]
-
- cache[self.layer_id] = (k, v)
-
- f_dim_per_head = tf.cast(dim_per_head, dtype=q.dtype)
- q = tf.multiply(q, tf.math.rsqrt(f_dim_per_head)) # (bs, n_heads, qlen, dim_per_head)
- k = tf.cast(k, dtype=q.dtype)
- scores = tf.matmul(q, k, transpose_b=True) # (bs, n_heads, qlen, klen)
- mask = tf.reshape(mask, mask_reshape) # (bs, n_heads, qlen, klen)
- # scores.masked_fill_(mask, -float('inf')) # (bs, n_heads, qlen, klen)
- mask = tf.cast(mask, dtype=scores.dtype)
- scores = scores - 1e30 * (1.0 - mask)
- weights = stable_softmax(scores, axis=-1) # (bs, n_heads, qlen, klen)
- weights = self.dropout(weights, training=training) # (bs, n_heads, qlen, klen)
-
- # Mask heads if we want to
- if head_mask is not None:
- weights = weights * head_mask
-
- context = tf.matmul(weights, v) # (bs, n_heads, qlen, dim_per_head)
- context = unshape(context) # (bs, qlen, dim)
- outputs = (self.out_lin(context),)
-
- if output_attentions:
- outputs = outputs + (weights,)
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "q_lin", None) is not None:
- with tf.name_scope(self.q_lin.name):
- self.q_lin.build([None, None, self.dim])
- if getattr(self, "k_lin", None) is not None:
- with tf.name_scope(self.k_lin.name):
- self.k_lin.build([None, None, self.dim])
- if getattr(self, "v_lin", None) is not None:
- with tf.name_scope(self.v_lin.name):
- self.v_lin.build([None, None, self.dim])
- if getattr(self, "out_lin", None) is not None:
- with tf.name_scope(self.out_lin.name):
- self.out_lin.build([None, None, self.dim])
-
-
-# Copied from transformers.models.xlm.modeling_tf_xlm.TFXLMTransformerFFN
-class TFFlaubertTransformerFFN(keras.layers.Layer):
- def __init__(self, in_dim, dim_hidden, out_dim, config, **kwargs):
- super().__init__(**kwargs)
-
- self.lin1 = keras.layers.Dense(dim_hidden, kernel_initializer=get_initializer(config.init_std), name="lin1")
- self.lin2 = keras.layers.Dense(out_dim, kernel_initializer=get_initializer(config.init_std), name="lin2")
- self.act = get_tf_activation("gelu") if config.gelu_activation else get_tf_activation("relu")
- self.dropout = keras.layers.Dropout(config.dropout)
- self.in_dim = in_dim
- self.dim_hidden = dim_hidden
-
- def call(self, input, training=False):
- x = self.lin1(input)
- x = self.act(x)
- x = self.lin2(x)
- x = self.dropout(x, training=training)
-
- return x
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "lin1", None) is not None:
- with tf.name_scope(self.lin1.name):
- self.lin1.build([None, None, self.in_dim])
- if getattr(self, "lin2", None) is not None:
- with tf.name_scope(self.lin2.name):
- self.lin2.build([None, None, self.dim_hidden])
-
-
-@keras_serializable
-class TFFlaubertMainLayer(keras.layers.Layer):
- config_class = FlaubertConfig
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.n_heads = config.n_heads
- self.n_langs = config.n_langs
- self.dim = config.emb_dim
- self.hidden_dim = self.dim * 4
- self.n_words = config.n_words
- self.pad_index = config.pad_index
- self.causal = config.causal
- self.n_layers = config.n_layers
- self.use_lang_emb = config.use_lang_emb
- self.layerdrop = getattr(config, "layerdrop", 0.0)
- self.pre_norm = getattr(config, "pre_norm", False)
- self.output_attentions = config.output_attentions
- self.output_hidden_states = config.output_hidden_states
- self.return_dict = config.use_return_dict
- self.max_position_embeddings = config.max_position_embeddings
- self.embed_init_std = config.embed_init_std
- self.dropout = keras.layers.Dropout(config.dropout)
- self.embeddings = TFSharedEmbeddings(
- self.n_words, self.dim, initializer_range=config.embed_init_std, name="embeddings"
- )
- self.layer_norm_emb = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm_emb")
- self.attentions = []
- self.layer_norm1 = []
- self.ffns = []
- self.layer_norm2 = []
-
- for i in range(self.n_layers):
- self.attentions.append(
- TFFlaubertMultiHeadAttention(self.n_heads, self.dim, config=config, name=f"attentions_._{i}")
- )
- self.layer_norm1.append(
- keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name=f"layer_norm1_._{i}")
- )
- # if self.is_decoder:
- # self.layer_norm15.append(nn.LayerNorm(self.dim, eps=config.layer_norm_eps))
- # self.encoder_attn.append(MultiHeadAttention(self.n_heads, self.dim, dropout=self.attention_dropout))
- self.ffns.append(
- TFFlaubertTransformerFFN(self.dim, self.hidden_dim, self.dim, config=config, name=f"ffns_._{i}")
- )
- self.layer_norm2.append(
- keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name=f"layer_norm2_._{i}")
- )
-
- def build(self, input_shape=None):
- with tf.name_scope("position_embeddings"):
- self.position_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.max_position_embeddings, self.dim],
- initializer=get_initializer(self.embed_init_std),
- )
-
- if self.n_langs > 1 and self.use_lang_emb:
- with tf.name_scope("lang_embeddings"):
- self.lang_embeddings = self.add_weight(
- name="embeddings",
- shape=[self.n_langs, self.dim],
- initializer=get_initializer(self.embed_init_std),
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "layer_norm_emb", None) is not None:
- with tf.name_scope(self.layer_norm_emb.name):
- self.layer_norm_emb.build([None, None, self.dim])
- for layer in self.attentions:
- with tf.name_scope(layer.name):
- layer.build(None)
- for layer in self.layer_norm1:
- with tf.name_scope(layer.name):
- layer.build([None, None, self.dim])
- for layer in self.ffns:
- with tf.name_scope(layer.name):
- layer.build(None)
- for layer in self.layer_norm2:
- with tf.name_scope(layer.name):
- layer.build([None, None, self.dim])
-
- def get_input_embeddings(self):
- return self.embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.weight = value
- self.embeddings.vocab_size = shape_list(value)[0]
-
- @unpack_inputs
- def call(
- self,
- input_ids: np.ndarray | tf.Tensor | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- langs: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- lengths: np.ndarray | tf.Tensor | None = None,
- cache: Optional[Dict[str, tf.Tensor]] = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFBaseModelOutput]:
- # removed: src_enc=None, src_len=None
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- bs, slen = shape_list(input_ids)
- elif inputs_embeds is not None:
- bs, slen = shape_list(inputs_embeds)[:2]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if lengths is None:
- if input_ids is not None:
- lengths = tf.reduce_sum(
- tf.cast(tf.not_equal(input_ids, self.pad_index), dtype=input_ids.dtype), axis=1
- )
- else:
- lengths = tf.convert_to_tensor([slen] * bs)
- # mask = input_ids != self.pad_index
-
- # check inputs
- # assert shape_list(lengths)[0] == bs
- (
- tf.debugging.assert_equal(shape_list(lengths)[0], bs),
- f"Expected batch size {shape_list(lengths)[0]} and received batch size {bs} mismatched",
- )
- # assert lengths.max().item() <= slen
- # input_ids = input_ids.transpose(0, 1) # batch size as dimension 0
- # assert (src_enc is None) == (src_len is None)
- # if src_enc is not None:
- # assert self.is_decoder
- # assert src_enc.size(0) == bs
-
- # generate masks
- mask, attn_mask = get_masks(slen, lengths, self.causal, padding_mask=attention_mask)
- # if self.is_decoder and src_enc is not None:
- # src_mask = torch.arange(src_len.max(), dtype=torch.long, device=lengths.device) < src_len[:, None]
-
- # position_ids
- if position_ids is None:
- position_ids = tf.expand_dims(tf.range(slen), axis=0)
- position_ids = tf.tile(position_ids, (bs, 1))
-
- # assert shape_list(position_ids) == [bs, slen] # (slen, bs)
- (
- tf.debugging.assert_equal(shape_list(position_ids), [bs, slen]),
- f"Position id shape {shape_list(position_ids)} and input shape {[bs, slen]} mismatched",
- )
- # position_ids = position_ids.transpose(0, 1)
-
- # langs
- if langs is not None:
- # assert shape_list(langs) == [bs, slen] # (slen, bs)
- (
- tf.debugging.assert_equal(shape_list(langs), [bs, slen]),
- f"Lang shape {shape_list(langs)} and input shape {[bs, slen]} mismatched",
- )
- # langs = langs.transpose(0, 1)
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x qlen x klen]
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.n_layers
-
- # do not recompute cached elements
- if cache is not None and input_ids is not None:
- _slen = slen - cache["slen"]
- input_ids = input_ids[:, -_slen:]
- position_ids = position_ids[:, -_slen:]
- if langs is not None:
- langs = langs[:, -_slen:]
- mask = mask[:, -_slen:]
- attn_mask = attn_mask[:, -_slen:]
-
- # embeddings
- if inputs_embeds is None:
- check_embeddings_within_bounds(input_ids, self.embeddings.vocab_size)
- inputs_embeds = self.embeddings(input_ids)
-
- tensor = inputs_embeds + tf.gather(self.position_embeddings, position_ids)
-
- if langs is not None and self.use_lang_emb:
- tensor = tensor + tf.gather(self.lang_embeddings, langs)
- if token_type_ids is not None:
- tensor = tensor + self.embeddings(token_type_ids)
-
- tensor = self.layer_norm_emb(tensor)
- tensor = self.dropout(tensor, training=training)
- mask = tf.cast(mask, dtype=tensor.dtype)
- tensor = tensor * tf.expand_dims(mask, axis=-1)
-
- # hidden_states and attentions cannot be None in graph mode.
- hidden_states = () if output_hidden_states else None
- attentions = () if output_attentions else None
-
- # transformer layers
- for i in range(self.n_layers):
- # LayerDrop
- dropout_probability = random.uniform(0, 1)
-
- if training and (dropout_probability < self.layerdrop):
- continue
-
- if output_hidden_states:
- hidden_states = hidden_states + (tensor,)
-
- # self attention
- if not self.pre_norm:
- attn_outputs = self.attentions[i](
- tensor,
- attn_mask,
- None,
- cache,
- head_mask[i],
- output_attentions,
- training=training,
- )
- attn = attn_outputs[0]
-
- if output_attentions:
- attentions = attentions + (attn_outputs[1],)
-
- attn = self.dropout(attn, training=training)
- tensor = tensor + attn
- tensor = self.layer_norm1[i](tensor)
- else:
- tensor_normalized = self.layer_norm1[i](tensor)
- attn_outputs = self.attentions[i](
- tensor_normalized,
- attn_mask,
- None,
- cache,
- head_mask[i],
- output_attentions,
- training=training,
- )
- attn = attn_outputs[0]
-
- if output_attentions:
- attentions = attentions + (attn_outputs[1],)
-
- attn = self.dropout(attn, training=training)
- tensor = tensor + attn
-
- # encoder attention (for decoder only)
- # if self.is_decoder and src_enc is not None:
- # attn = self.encoder_attn[i](tensor, src_mask, kv=src_enc, cache=cache)
- # attn = nn.functional.dropout(attn, p=self.dropout, training=self.training)
- # tensor = tensor + attn
- # tensor = self.layer_norm15[i](tensor)
-
- # FFN
- if not self.pre_norm:
- tensor = tensor + self.ffns[i](tensor)
- tensor = self.layer_norm2[i](tensor)
- else:
- tensor_normalized = self.layer_norm2[i](tensor)
- tensor = tensor + self.ffns[i](tensor_normalized)
-
- tensor = tensor * tf.expand_dims(mask, axis=-1)
-
- # Add last hidden state
- if output_hidden_states:
- hidden_states = hidden_states + (tensor,)
-
- # update cache length
- if cache is not None:
- cache["slen"] += tensor.size(1)
-
- # move back sequence length to dimension 0
- # tensor = tensor.transpose(0, 1)
-
- if not return_dict:
- return tuple(v for v in [tensor, hidden_states, attentions] if v is not None)
-
- return TFBaseModelOutput(last_hidden_state=tensor, hidden_states=hidden_states, attentions=attentions)
-
-
-# Copied from transformers.models.xlm.modeling_tf_xlm.TFXLMPredLayer
-class TFFlaubertPredLayer(keras.layers.Layer):
- """
- Prediction layer (cross_entropy or adaptive_softmax).
- """
-
- def __init__(self, config, input_embeddings, **kwargs):
- super().__init__(**kwargs)
-
- self.asm = config.asm
- self.n_words = config.n_words
- self.pad_index = config.pad_index
-
- if config.asm is False:
- self.input_embeddings = input_embeddings
- else:
- raise NotImplementedError
- # self.proj = nn.AdaptiveLogSoftmaxWithLoss(
- # in_features=dim,
- # n_classes=config.n_words,
- # cutoffs=config.asm_cutoffs,
- # div_value=config.asm_div_value,
- # head_bias=True, # default is False
- # )
-
- def build(self, input_shape):
- # The output weights are the same as the input embeddings, but there is an output-only bias for each token.
- self.bias = self.add_weight(shape=(self.n_words,), initializer="zeros", trainable=True, name="bias")
-
- super().build(input_shape)
-
- def get_output_embeddings(self):
- return self.input_embeddings
-
- def set_output_embeddings(self, value):
- self.input_embeddings.weight = value
- self.input_embeddings.vocab_size = shape_list(value)[0]
-
- def get_bias(self):
- return {"bias": self.bias}
-
- def set_bias(self, value):
- self.bias = value["bias"]
- self.vocab_size = shape_list(value["bias"])[0]
-
- def call(self, hidden_states):
- hidden_states = self.input_embeddings(hidden_states, mode="linear")
- hidden_states = hidden_states + self.bias
-
- return hidden_states
-
-
-@dataclass
-class TFFlaubertWithLMHeadModelOutput(ModelOutput):
- """
- Base class for [`TFFlaubertWithLMHeadModel`] outputs.
-
- Args:
- logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- logits: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-@add_start_docstrings(
- """
- The Flaubert Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- FLAUBERT_START_DOCSTRING,
-)
-class TFFlaubertWithLMHeadModel(TFFlaubertPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFFlaubertMainLayer(config, name="transformer")
- self.pred_layer = TFFlaubertPredLayer(config, self.transformer.embeddings, name="pred_layer_._proj")
- # Flaubert does not have past caching features
- self.supports_xla_generation = False
-
- def get_lm_head(self):
- return self.pred_layer
-
- def get_prefix_bias_name(self):
- warnings.warn("The method get_prefix_bias_name is deprecated. Please use `get_bias` instead.", FutureWarning)
- return self.name + "/" + self.pred_layer.name
-
- def prepare_inputs_for_generation(self, inputs, **kwargs):
- mask_token_id = self.config.mask_token_id
- lang_id = self.config.lang_id
-
- effective_batch_size = inputs.shape[0]
- mask_token = tf.fill((effective_batch_size, 1), 1) * mask_token_id
- inputs = tf.concat([inputs, mask_token], axis=1)
-
- if lang_id is not None:
- langs = tf.ones_like(inputs) * lang_id
- else:
- langs = None
- return {"input_ids": inputs, "langs": langs}
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFFlaubertWithLMHeadModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: np.ndarray | tf.Tensor | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- langs: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- lengths: np.ndarray | tf.Tensor | None = None,
- cache: Optional[Dict[str, tf.Tensor]] = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[Tuple, TFFlaubertWithLMHeadModelOutput]:
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- output = transformer_outputs[0]
- outputs = self.pred_layer(output)
-
- if not return_dict:
- return (outputs,) + transformer_outputs[1:]
-
- return TFFlaubertWithLMHeadModelOutput(
- logits=outputs, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "pred_layer", None) is not None:
- with tf.name_scope(self.pred_layer.name):
- self.pred_layer.build(None)
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a sequence classification/regression head on top (a linear layer on top of the pooled output)
- e.g. for GLUE tasks.
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied from transformers.models.xlm.modeling_tf_xlm.TFXLMForSequenceClassification with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class TFFlaubertForSequenceClassification(TFFlaubertPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
-
- self.transformer = TFFlaubertMainLayer(config, name="transformer")
- self.sequence_summary = TFSequenceSummary(config, initializer_range=config.init_std, name="sequence_summary")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- langs: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- lengths: np.ndarray | tf.Tensor | None = None,
- cache: Optional[Dict[str, tf.Tensor]] = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- output = transformer_outputs[0]
-
- logits = self.sequence_summary(output)
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "sequence_summary", None) is not None:
- with tf.name_scope(self.sequence_summary.name):
- self.sequence_summary.build(None)
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied from transformers.models.xlm.modeling_tf_xlm.TFXLMForQuestionAnsweringSimple with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class TFFlaubertForQuestionAnsweringSimple(TFFlaubertPreTrainedModel, TFQuestionAnsweringLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFFlaubertMainLayer(config, name="transformer")
- self.qa_outputs = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.init_std), name="qa_outputs"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFQuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- langs: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- lengths: np.ndarray | tf.Tensor | None = None,
- cache: Optional[Dict[str, tf.Tensor]] = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- start_positions: np.ndarray | tf.Tensor | None = None,
- end_positions: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[TFQuestionAnsweringModelOutput, Tuple[tf.Tensor]]:
- r"""
- start_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = transformer_outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = tf.split(logits, 2, axis=-1)
- start_logits = tf.squeeze(start_logits, axis=-1)
- end_logits = tf.squeeze(end_logits, axis=-1)
-
- loss = None
- if start_positions is not None and end_positions is not None:
- labels = {"start_position": start_positions}
- labels["end_position"] = end_positions
- loss = self.hf_compute_loss(labels, (start_logits, end_logits))
-
- if not return_dict:
- output = (start_logits, end_logits) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFQuestionAnsweringModelOutput(
- loss=loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied from transformers.models.xlm.modeling_tf_xlm.TFXLMForTokenClassification with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class TFFlaubertForTokenClassification(TFFlaubertPreTrainedModel, TFTokenClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
-
- self.transformer = TFFlaubertMainLayer(config, name="transformer")
- self.dropout = keras.layers.Dropout(config.dropout)
- self.classifier = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.init_std), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FLAUBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFTokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- langs: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- lengths: np.ndarray | tf.Tensor | None = None,
- cache: Optional[Dict[str, tf.Tensor]] = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[TFTokenClassifierOutput, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- attention_mask=attention_mask,
- langs=langs,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- lengths=lengths,
- cache=cache,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = transformer_outputs[0]
-
- sequence_output = self.dropout(sequence_output, training=training)
- logits = self.classifier(sequence_output)
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFTokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- Flaubert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- FLAUBERT_START_DOCSTRING,
-)
-# Copied from transformers.models.xlm.modeling_tf_xlm.TFXLMForMultipleChoice with XLM_INPUTS->FLAUBERT_INPUTS,XLM->Flaubert
-class TFFlaubertForMultipleChoice(TFFlaubertPreTrainedModel, TFMultipleChoiceLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
-
- self.transformer = TFFlaubertMainLayer(config, name="transformer")
- self.sequence_summary = TFSequenceSummary(config, initializer_range=config.init_std, name="sequence_summary")
- self.logits_proj = keras.layers.Dense(
- 1, kernel_initializer=get_initializer(config.initializer_range), name="logits_proj"
- )
- self.config = config
-
- @property
- def dummy_inputs(self):
- """
- Dummy inputs to build the network.
-
- Returns:
- tf.Tensor with dummy inputs
- """
- # Sometimes Flaubert has language embeddings so don't forget to build them as well if needed
- if self.config.use_lang_emb and self.config.n_langs > 1:
- return {
- "input_ids": tf.constant(MULTIPLE_CHOICE_DUMMY_INPUTS, dtype=tf.int32),
- "langs": tf.constant(MULTIPLE_CHOICE_DUMMY_INPUTS, dtype=tf.int32),
- }
- else:
- return {
- "input_ids": tf.constant(MULTIPLE_CHOICE_DUMMY_INPUTS, dtype=tf.int32),
- }
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(
- FLAUBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
- )
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFMultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- langs: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- lengths: np.ndarray | tf.Tensor | None = None,
- cache: Optional[Dict[str, tf.Tensor]] = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[TFMultipleChoiceModelOutput, Tuple[tf.Tensor]]:
- if input_ids is not None:
- num_choices = shape_list(input_ids)[1]
- seq_length = shape_list(input_ids)[2]
- else:
- num_choices = shape_list(inputs_embeds)[1]
- seq_length = shape_list(inputs_embeds)[2]
-
- flat_input_ids = tf.reshape(input_ids, (-1, seq_length)) if input_ids is not None else None
- flat_attention_mask = tf.reshape(attention_mask, (-1, seq_length)) if attention_mask is not None else None
- flat_token_type_ids = tf.reshape(token_type_ids, (-1, seq_length)) if token_type_ids is not None else None
- flat_position_ids = tf.reshape(position_ids, (-1, seq_length)) if position_ids is not None else None
- flat_langs = tf.reshape(langs, (-1, seq_length)) if langs is not None else None
- flat_inputs_embeds = (
- tf.reshape(inputs_embeds, (-1, seq_length, shape_list(inputs_embeds)[3]))
- if inputs_embeds is not None
- else None
- )
-
- if lengths is not None:
- logger.warning(
- "The `lengths` parameter cannot be used with the Flaubert multiple choice models. Please use the "
- "attention mask instead.",
- )
- lengths = None
-
- transformer_outputs = self.transformer(
- flat_input_ids,
- flat_attention_mask,
- flat_langs,
- flat_token_type_ids,
- flat_position_ids,
- lengths,
- cache,
- head_mask,
- flat_inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- output = transformer_outputs[0]
- logits = self.sequence_summary(output)
- logits = self.logits_proj(logits)
- reshaped_logits = tf.reshape(logits, (-1, num_choices))
-
- loss = None if labels is None else self.hf_compute_loss(labels, reshaped_logits)
-
- if not return_dict:
- output = (reshaped_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFMultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "sequence_summary", None) is not None:
- with tf.name_scope(self.sequence_summary.name):
- self.sequence_summary.build(None)
- if getattr(self, "logits_proj", None) is not None:
- with tf.name_scope(self.logits_proj.name):
- self.logits_proj.build([None, None, self.config.num_labels])
diff --git a/transformers/models/flaubert/tokenization_flaubert.py b/transformers/models/flaubert/tokenization_flaubert.py
deleted file mode 100644
index 20f9926422064d61fa483275c45bf8f087e85cb0..0000000000000000000000000000000000000000
--- a/transformers/models/flaubert/tokenization_flaubert.py
+++ /dev/null
@@ -1,565 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present CNRS, Facebook Inc. and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for Flaubert."""
-
-
-import json
-import os
-import re
-import unicodedata
-from typing import List, Optional, Tuple
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {
- "vocab_file": "vocab.json",
- "merges_file": "merges.txt",
-}
-
-
-def convert_to_unicode(text):
- """
- Converts `text` to Unicode (if it's not already), assuming UTF-8 input.
- """
-
- def ensure_text(s, encoding="utf-8", errors="strict"):
- if isinstance(s, bytes):
- return s.decode(encoding, errors)
- elif isinstance(s, str):
- return s
- else:
- raise TypeError(f"not expecting type '{type(s)}'")
-
- return ensure_text(text, encoding="utf-8", errors="ignore")
-
-
-# Copied from transformers.models.xlm.tokenization_xlm.get_pairs
-def get_pairs(word):
- """
- Return set of symbol pairs in a word. word is represented as tuple of symbols (symbols being variable-length
- strings)
- """
- pairs = set()
- prev_char = word[0]
- for char in word[1:]:
- pairs.add((prev_char, char))
- prev_char = char
- return pairs
-
-
-# Copied from transformers.models.xlm.tokenization_xlm.replace_unicode_punct
-def replace_unicode_punct(text):
- """
- Port of https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/replace-unicode-punctuation.perl
- """
- text = text.replace("，", ",")
- text = re.sub(r"。\s*", ". ", text)
- text = text.replace("、", ",")
- text = text.replace("”", '"')
- text = text.replace("“", '"')
- text = text.replace("∶", ":")
- text = text.replace("：", ":")
- text = text.replace("？", "?")
- text = text.replace("《", '"')
- text = text.replace("》", '"')
- text = text.replace("）", ")")
- text = text.replace("！", "!")
- text = text.replace("（", "(")
- text = text.replace("；", ";")
- text = text.replace("１", "1")
- text = text.replace("」", '"')
- text = text.replace("「", '"')
- text = text.replace("０", "0")
- text = text.replace("３", "3")
- text = text.replace("２", "2")
- text = text.replace("５", "5")
- text = text.replace("６", "6")
- text = text.replace("９", "9")
- text = text.replace("７", "7")
- text = text.replace("８", "8")
- text = text.replace("４", "4")
- text = re.sub(r"．\s*", ". ", text)
- text = text.replace("～", "~")
- text = text.replace("’", "'")
- text = text.replace("…", "...")
- text = text.replace("━", "-")
- text = text.replace("〈", "<")
- text = text.replace("〉", ">")
- text = text.replace("【", "[")
- text = text.replace("】", "]")
- text = text.replace("％", "%")
- return text
-
-
-# Copied from transformers.models.xlm.tokenization_xlm.remove_non_printing_char
-def remove_non_printing_char(text):
- """
- Port of https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/remove-non-printing-char.perl
- """
- output = []
- for char in text:
- cat = unicodedata.category(char)
- if cat.startswith("C"):
- continue
- output.append(char)
- return "".join(output)
-
-
-class FlaubertTokenizer(PreTrainedTokenizer):
- """
- Construct a Flaubert tokenizer. Based on Byte-Pair Encoding. The tokenization process is the following:
-
- - Moses preprocessing and tokenization.
- - Normalizing all inputs text.
- - The arguments `special_tokens` and the function `set_special_tokens`, can be used to add additional symbols (like
- "__classify__") to a vocabulary.
- - The argument `do_lowercase` controls lower casing (automatically set for pretrained vocabularies).
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- Vocabulary file.
- merges_file (`str`):
- Merges file.
- do_lowercase (`bool`, *optional*, defaults to `False`):
- Controls lower casing.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the beginning of
- sequence. The token used is the `cls_token`.
-
-
-
- sep_token (`str`, *optional*, defaults to `""`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `""`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `""`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- additional_special_tokens (`List[str]`, *optional*, defaults to `['', '', '', '', '', '', '', '', '', '']`):
- List of additional special tokens.
- lang2id (`Dict[str, int]`, *optional*):
- Dictionary mapping languages string identifiers to their IDs.
- id2lang (`Dict[int, str]`, *optional*):
- Dictionary mapping language IDs to their string identifiers.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
-
- def __init__(
- self,
- vocab_file,
- merges_file,
- do_lowercase=False,
- unk_token="",
- bos_token="",
- sep_token="",
- pad_token="",
- cls_token="",
- mask_token="",
- additional_special_tokens=[
- "",
- "",
- "",
- "",
- "",
- "",
- "",
- "",
- "",
- "",
- ],
- lang2id=None,
- id2lang=None,
- **kwargs,
- ):
- do_lowercase_and_remove_accent = kwargs.pop("do_lowercase_and_remove_accent", None)
- if do_lowercase_and_remove_accent is not None:
- logger.warning(
- "`do_lowercase_and_remove_accent` is passed as a keyword argument, but this won't do anything."
- " `FlaubertTokenizer` will always set it to `False`."
- )
- # always `False`
- self.do_lowercase_and_remove_accent = False
-
- self.do_lowercase = do_lowercase
-
- try:
- import sacremoses
- except ImportError:
- raise ImportError(
- "You need to install sacremoses to use FlaubertTokenizer. "
- "See https://pypi.org/project/sacremoses/ for installation."
- )
-
- self.sm = sacremoses
-
- # cache of sm.MosesPunctNormalizer instance
- self.cache_moses_punct_normalizer = {}
- # cache of sm.MosesTokenizer instance
- self.cache_moses_tokenizer = {}
- self.lang_with_custom_tokenizer = {"zh", "th", "ja"}
- self.lang2id = lang2id
- self.id2lang = id2lang
- if lang2id is not None and id2lang is not None:
- assert len(lang2id) == len(id2lang)
-
- self.ja_word_tokenizer = None
- self.zh_word_tokenizer = None
-
- with open(vocab_file, encoding="utf-8") as vocab_handle:
- self.encoder = json.load(vocab_handle)
- self.decoder = {v: k for k, v in self.encoder.items()}
- with open(merges_file, encoding="utf-8") as merges_handle:
- merges = merges_handle.read().split("\n")[:-1]
- merges = [tuple(merge.split()[:2]) for merge in merges]
- self.bpe_ranks = dict(zip(merges, range(len(merges))))
- self.cache = {}
-
- super().__init__(
- unk_token=unk_token,
- bos_token=bos_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- additional_special_tokens=additional_special_tokens,
- lang2id=lang2id,
- id2lang=id2lang,
- **kwargs,
- )
-
- @property
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.do_lower_case
- def do_lower_case(self):
- return self.do_lowercase_and_remove_accent
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.moses_punct_norm
- def moses_punct_norm(self, text, lang):
- if lang not in self.cache_moses_punct_normalizer:
- punct_normalizer = self.sm.MosesPunctNormalizer(lang=lang)
- self.cache_moses_punct_normalizer[lang] = punct_normalizer
- else:
- punct_normalizer = self.cache_moses_punct_normalizer[lang]
- return punct_normalizer.normalize(text)
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.moses_tokenize
- def moses_tokenize(self, text, lang):
- if lang not in self.cache_moses_tokenizer:
- moses_tokenizer = self.sm.MosesTokenizer(lang=lang)
- self.cache_moses_tokenizer[lang] = moses_tokenizer
- else:
- moses_tokenizer = self.cache_moses_tokenizer[lang]
- return moses_tokenizer.tokenize(text, return_str=False, escape=False)
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.moses_pipeline
- def moses_pipeline(self, text, lang):
- text = replace_unicode_punct(text)
- text = self.moses_punct_norm(text, lang)
- text = remove_non_printing_char(text)
- return text
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.ja_tokenize
- def ja_tokenize(self, text):
- if self.ja_word_tokenizer is None:
- try:
- import Mykytea
-
- self.ja_word_tokenizer = Mykytea.Mykytea(
- f"-model {os.path.expanduser('~')}/local/share/kytea/model.bin"
- )
- except (AttributeError, ImportError):
- logger.error(
- "Make sure you install KyTea (https://github.com/neubig/kytea) and it's python wrapper"
- " (https://github.com/chezou/Mykytea-python) with the following steps"
- )
- logger.error("1. git clone git@github.com:neubig/kytea.git && cd kytea")
- logger.error("2. autoreconf -i")
- logger.error("3. ./configure --prefix=$HOME/local")
- logger.error("4. make && make install")
- logger.error("5. pip install kytea")
- raise
- return list(self.ja_word_tokenizer.getWS(text))
-
- @property
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.vocab_size
- def vocab_size(self):
- return len(self.encoder)
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.get_vocab
- def get_vocab(self):
- return dict(self.encoder, **self.added_tokens_encoder)
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.bpe
- def bpe(self, token):
- word = tuple(token[:-1]) + (token[-1] + "",)
- if token in self.cache:
- return self.cache[token]
- pairs = get_pairs(word)
-
- if not pairs:
- return token + ""
-
- while True:
- bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
- if bigram not in self.bpe_ranks:
- break
- first, second = bigram
- new_word = []
- i = 0
- while i < len(word):
- try:
- j = word.index(first, i)
- except ValueError:
- new_word.extend(word[i:])
- break
- else:
- new_word.extend(word[i:j])
- i = j
-
- if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
- new_word.append(first + second)
- i += 2
- else:
- new_word.append(word[i])
- i += 1
- new_word = tuple(new_word)
- word = new_word
- if len(word) == 1:
- break
- else:
- pairs = get_pairs(word)
- word = " ".join(word)
- if word == "\n ":
- word = "\n"
- self.cache[token] = word
- return word
-
- def preprocess_text(self, text):
- text = text.replace("``", '"').replace("''", '"')
- text = convert_to_unicode(text)
- text = unicodedata.normalize("NFC", text)
-
- if self.do_lowercase:
- text = text.lower()
-
- return text
-
- def _tokenize(self, text, bypass_tokenizer=False):
- """
- Tokenize a string given language code using Moses.
-
- Details of tokenization:
-
- - [sacremoses](https://github.com/alvations/sacremoses): port of Moses
- - Install with `pip install sacremoses`
-
- Args:
- - bypass_tokenizer: Allow users to preprocess and tokenize the sentences externally (default = False)
- (bool). If True, we only apply BPE.
-
- Returns:
- List of tokens.
- """
- lang = "fr"
- if lang and self.lang2id and lang not in self.lang2id:
- logger.error(
- "Supplied language code not found in lang2id mapping. Please check that your language is supported by"
- " the loaded pretrained model."
- )
-
- if bypass_tokenizer:
- text = text.split()
- else:
- text = self.preprocess_text(text)
- text = self.moses_pipeline(text, lang=lang)
- text = self.moses_tokenize(text, lang=lang)
-
- split_tokens = []
- for token in text:
- if token:
- split_tokens.extend(list(self.bpe(token).split(" ")))
-
- return split_tokens
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index, self.unk_token)
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.convert_tokens_to_string
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- out_string = "".join(tokens).replace("", " ").strip()
- return out_string
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. An XLM sequence has the following format:
-
- - single sequence: ` X `
- - pair of sequences: ` A B `
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
-
- """
- bos = [self.bos_token_id]
- sep = [self.sep_token_id]
-
- if token_ids_1 is None:
- return bos + token_ids_0 + sep
- return bos + token_ids_0 + sep + token_ids_1 + sep
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.get_special_tokens_mask
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. An XLM sequence
- pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- merge_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
- )
-
- with open(vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
- index = 0
- with open(merge_file, "w", encoding="utf-8") as writer:
- for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
- " Please check that the tokenizer is not corrupted!"
- )
- index = token_index
- writer.write(" ".join(bpe_tokens) + "\n")
- index += 1
-
- return vocab_file, merge_file
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.__getstate__
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sm"] = None
- return state
-
- # Copied from transformers.models.xlm.tokenization_xlm.XLMTokenizer.__setstate__
- def __setstate__(self, d):
- self.__dict__ = d
-
- try:
- import sacremoses
- except ImportError:
- raise ImportError(
- "You need to install sacremoses to use XLMTokenizer. "
- "See https://pypi.org/project/sacremoses/ for installation."
- )
-
- self.sm = sacremoses
diff --git a/transformers/models/flava/__init__.py b/transformers/models/flava/__init__.py
deleted file mode 100644
index 8d026a9443271c6f750bbe204abd777c1195ee07..0000000000000000000000000000000000000000
--- a/transformers/models/flava/__init__.py
+++ /dev/null
@@ -1,97 +0,0 @@
-# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {
- "configuration_flava": [
- "FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP",
- "FlavaConfig",
- "FlavaImageCodebookConfig",
- "FlavaImageConfig",
- "FlavaMultimodalConfig",
- "FlavaTextConfig",
- ],
-}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_flava"] = ["FlavaFeatureExtractor"]
- _import_structure["image_processing_flava"] = ["FlavaImageProcessor"]
- _import_structure["processing_flava"] = ["FlavaProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flava"] = [
- "FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST",
- "FlavaForPreTraining",
- "FlavaImageCodebook",
- "FlavaImageModel",
- "FlavaModel",
- "FlavaMultimodalModel",
- "FlavaPreTrainedModel",
- "FlavaTextModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_flava import (
- FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP,
- FlavaConfig,
- FlavaImageCodebookConfig,
- FlavaImageConfig,
- FlavaMultimodalConfig,
- FlavaTextConfig,
- )
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_flava import FlavaFeatureExtractor
- from .image_processing_flava import FlavaImageProcessor
- from .processing_flava import FlavaProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flava import (
- FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST,
- FlavaForPreTraining,
- FlavaImageCodebook,
- FlavaImageModel,
- FlavaModel,
- FlavaMultimodalModel,
- FlavaPreTrainedModel,
- FlavaTextModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/flava/__pycache__/__init__.cpython-310.pyc b/transformers/models/flava/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 28509517bb1358258b247411e0e4b01a92192eaa..0000000000000000000000000000000000000000
Binary files a/transformers/models/flava/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flava/__pycache__/configuration_flava.cpython-310.pyc b/transformers/models/flava/__pycache__/configuration_flava.cpython-310.pyc
deleted file mode 100644
index 012b449e5b3086af27a09b0ed8d18baf4fdb73d3..0000000000000000000000000000000000000000
Binary files a/transformers/models/flava/__pycache__/configuration_flava.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flava/__pycache__/convert_dalle_to_flava_codebook.cpython-310.pyc b/transformers/models/flava/__pycache__/convert_dalle_to_flava_codebook.cpython-310.pyc
deleted file mode 100644
index 9ecacce458f82f60579de2a3e9e47f6500d1e451..0000000000000000000000000000000000000000
Binary files a/transformers/models/flava/__pycache__/convert_dalle_to_flava_codebook.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flava/__pycache__/convert_flava_original_pytorch_to_hf.cpython-310.pyc b/transformers/models/flava/__pycache__/convert_flava_original_pytorch_to_hf.cpython-310.pyc
deleted file mode 100644
index 7d200c4027068cfbe0f0ff7a2dfc07c52cafae03..0000000000000000000000000000000000000000
Binary files a/transformers/models/flava/__pycache__/convert_flava_original_pytorch_to_hf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flava/__pycache__/feature_extraction_flava.cpython-310.pyc b/transformers/models/flava/__pycache__/feature_extraction_flava.cpython-310.pyc
deleted file mode 100644
index 39cd8f069f28c638daa1e343895574be001bb762..0000000000000000000000000000000000000000
Binary files a/transformers/models/flava/__pycache__/feature_extraction_flava.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flava/__pycache__/image_processing_flava.cpython-310.pyc b/transformers/models/flava/__pycache__/image_processing_flava.cpython-310.pyc
deleted file mode 100644
index 027a49056a7ec65c6545261f987d51858afffbee..0000000000000000000000000000000000000000
Binary files a/transformers/models/flava/__pycache__/image_processing_flava.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flava/__pycache__/modeling_flava.cpython-310.pyc b/transformers/models/flava/__pycache__/modeling_flava.cpython-310.pyc
deleted file mode 100644
index 18063abbfdfda20758fa23537c1acb96fa100bd8..0000000000000000000000000000000000000000
Binary files a/transformers/models/flava/__pycache__/modeling_flava.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flava/__pycache__/processing_flava.cpython-310.pyc b/transformers/models/flava/__pycache__/processing_flava.cpython-310.pyc
deleted file mode 100644
index aba44d9328483516d135130fbb0dc775eebdd30a..0000000000000000000000000000000000000000
Binary files a/transformers/models/flava/__pycache__/processing_flava.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/flava/configuration_flava.py b/transformers/models/flava/configuration_flava.py
deleted file mode 100644
index 2c8642bfd2759f0ce797611635b90fcbb160b12c..0000000000000000000000000000000000000000
--- a/transformers/models/flava/configuration_flava.py
+++ /dev/null
@@ -1,764 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" FLAVA model configurations"""
-
-import os
-from typing import Any, Dict, Union
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class FlavaImageConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FlavaImageModel`]. It is used to instantiate an
- FLAVA model according to the specified arguments, defining the model architecture.
-
- Instantiating a configuration with the defaults will yield a similar configuration to that of the FLAVA
- [facebook/flava-full](https://huggingface.co/facebook/flava-full) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 16):
- The size (resolution) of each patch.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- qkv_bias (`bool`, *optional*, defaults to `True`):
- Whether to add a bias to the queries, keys and values.
- mask_token (`bool`, *optional*, defaults to `True`):
- Whether to use a mask token or not. Used in MIM (Masked Image Modeling) loss for FLAVA.
- vocab_size (`int`, *optional*, defaults to 8192):
- Vocabulary size of the [`FlavaImageCodebook`] used in conjunction with [`FlavaImageModel`] for MIM (Masked
- Image Modeling) loss for FLAVA.
-
- Example:
-
- ```python
- >>> from transformers import FlavaImageConfig, FlavaImageModel
-
- >>> # Initializing a FlavaImageModel with style configuration
- >>> configuration = FlavaImageConfig()
-
- >>> # Initializing a FlavaImageModel model (with random weights) from the style configuration
- >>> model = FlavaImageModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "flava_image_model"
-
- def __init__(
- self,
- hidden_size: int = 768,
- num_hidden_layers: int = 12,
- num_attention_heads: int = 12,
- intermediate_size: int = 3072,
- hidden_act: int = "gelu",
- hidden_dropout_prob: float = 0.0,
- attention_probs_dropout_prob: float = 0.0,
- initializer_range: float = 0.02,
- layer_norm_eps: float = 1e-12,
- image_size: int = 224,
- patch_size: int = 16,
- num_channels: int = 3,
- qkv_bias: bool = True,
- mask_token: bool = True,
- vocab_size: int = 8192,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.qkv_bias = qkv_bias
- self.mask_token = mask_token
- self.vocab_size = vocab_size
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
- cls._set_token_in_kwargs(kwargs)
-
- config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
-
- # get the image config dict if we are loading from FlavaConfig
- if config_dict.get("model_type") == "flava":
- config_dict = config_dict["image_config"]
-
- if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
- logger.warning(
- f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
- f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
- )
-
- return cls.from_dict(config_dict, **kwargs)
-
-
-class FlavaTextConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FlavaTextModel`]. It is used to instantiate an
- FLAVA model according to the specified arguments, defining the model architecture.
-
- Instantiating a configuration with the defaults will yield a similar configuration to that of the FLAVA
- [facebook/flava-full](https://huggingface.co/facebook/flava-full) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`FlavaTextModel`].
- type_vocab_size (`int`, *optional*, defaults to 2):
- The vocabulary size of the `token_type_ids` passed when calling [`FlavaTextModel`]. Note that even though
- text encoder allows `token_type_ids`'s value as 2, for text-only pretraining and fine-tuning, only 1 is
- used similar to RoBERTa.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048). For VL, max_length passed to model is 77.
- position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
- Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
- positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
- [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
- For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
- with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 16):
- The size (resolution) of each patch.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- qkv_bias (`bool`, *optional*, defaults to `True`):
- Whether to add a bias to the queries, keys and values.
-
- Example:
-
- ```python
- >>> from transformers import FlavaTextConfig, FlavaTextModel
-
- >>> # Initializing a FlavaTextModel with style configuration
- >>> configuration = FlavaTextConfig()
-
- >>> # Initializing a FlavaTextModel model (with random weights) from the style configuration
- >>> model = FlavaTextModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "flava_text_model"
-
- def __init__(
- self,
- vocab_size: int = 30522,
- type_vocab_size: int = 2,
- max_position_embeddings: int = 512,
- position_embedding_type: str = "absolute",
- hidden_size: int = 768,
- num_hidden_layers: int = 12,
- num_attention_heads: int = 12,
- intermediate_size: int = 3072,
- hidden_act: str = "gelu",
- hidden_dropout_prob: float = 0.0,
- attention_probs_dropout_prob: float = 0.0,
- initializer_range: float = 0.02,
- layer_norm_eps: float = 1e-12,
- pad_token_id: int = 0,
- qkv_bias: bool = True,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.vocab_size = vocab_size
- self.type_vocab_size = type_vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.position_embedding_type = position_embedding_type
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.qkv_bias = qkv_bias
- self.pad_token_id = pad_token_id
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
- cls._set_token_in_kwargs(kwargs)
-
- config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
-
- # get the text config dict if we are loading from FlavaConfig
- if config_dict.get("model_type") == "flava":
- config_dict = config_dict["text_config"]
-
- if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
- logger.warning(
- f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
- f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
- )
-
- return cls.from_dict(config_dict, **kwargs)
-
-
-class FlavaMultimodalConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FlavaMultimodalModel`]. It is used to instantiate
- an FLAVA model according to the specified arguments, defining the model architecture.
-
- Instantiating a configuration with the defaults will yield a similar configuration to that of the FLAVA
- [facebook/flava-full](https://huggingface.co/facebook/flava-full) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 6):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- qkv_bias (`bool`, *optional*, defaults to `True`):
- Whether to add a bias to the queries, keys and values.
- use_cls_token (`bool`, *optional*, defaults to `True`):
- Whether to use an extra CLS token for multimodal settings. Usually needed by the FLAVA model.
-
-
- Example:
-
- ```python
- >>> from transformers import FlavaMultimodalConfig, FlavaMultimodalModel
-
- >>> # Initializing a FlavaMultimodalModel with style configuration
- >>> configuration = FlavaMultimodalConfig()
-
- >>> # Initializing a FlavaMultimodalModel model (with random weights) from the style configuration
- >>> model = FlavaMultimodalModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "flava_multimodal_model"
-
- def __init__(
- self,
- hidden_size: int = 768,
- num_hidden_layers: int = 6,
- num_attention_heads: int = 12,
- intermediate_size: int = 3072,
- hidden_act: int = "gelu",
- hidden_dropout_prob: int = 0.0,
- attention_probs_dropout_prob: int = 0.0,
- initializer_range: float = 0.02,
- layer_norm_eps: float = 1e-12,
- qkv_bias: bool = True,
- use_cls_token: bool = True,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.qkv_bias = qkv_bias
- self.use_cls_token = use_cls_token
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
- cls._set_token_in_kwargs(kwargs)
-
- config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
-
- # get the multimodal config dict if we are loading from FlavaConfig
- if config_dict.get("model_type") == "flava":
- config_dict = config_dict["multimodal_config"]
-
- if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
- logger.warning(
- f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
- f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
- )
-
- return cls.from_dict(config_dict, **kwargs)
-
-
-class FlavaImageCodebookConfig(PretrainedConfig):
- model_type = "flava_image_codebook"
-
- r"""
- [`FlavaImageCodebookConfig`] is the configuration class to store the configuration of a [`FlavaImageCodebook`]. It
- is used to instantiate an FLAVA model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the FLAVA
- [facebook/flava-image-codebook](https://huggingface.co/facebook/flava-image-codebook) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- num_groups (`int`, defaults to 4):
- Number of groups to be created. This parameter as of now doesn't affect the model and is used for some
- internal calculation and estimations.
- input_channels (`int`, defaults to 3):
- Number of channels in the image to be passed.
- num_blocks_per_group (`int`, defaults to 2):
- Number of conv-based blocks per group.
- hidden_size (`int`, defaults to 256):
- Size of hidden dim for the blocks.
- vocab_size (`int`, defaults to 8192):
- Size of the output vocabulary for the codebook.
- freeze (`bool`, defaults to `True`):
- Whether to freeze the weights of the model.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- kwargs (*optional*):
- Dictionary of keyword arguments.
-
- Example:
-
- ```python
- >>> from transformers import FlavaImageCodebookConfig, FlavaImageCodebook
-
- >>> # Initializing a FlavaImageCodebook with style configuration
- >>> configuration = FlavaImageCodebookConfig()
-
- >>> # Initializing a FlavaImageCodebook model (with random weights) from the style configuration
- >>> model = FlavaImageCodebook(configuration)
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```
- """
-
- def __init__(
- self,
- num_groups: int = 4,
- input_channels: int = 3,
- num_blocks_per_group: int = 2,
- hidden_size: int = 256,
- vocab_size: int = 8192,
- freeze: int = True,
- initializer_range: float = 0.02,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.num_groups = num_groups
- self.input_channels = input_channels
- self.num_blocks_per_group = num_blocks_per_group
- self.hidden_size = hidden_size
- self.vocab_size = vocab_size
- self.freeze = freeze
- self.initializer_range = initializer_range
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
- cls._set_token_in_kwargs(kwargs)
-
- config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
-
- # get the image codebook config dict if we are loading from FlavaConfig
- if config_dict.get("model_type") == "flava":
- config_dict = config_dict["image_codebook_config"]
-
- if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
- logger.warning(
- f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
- f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
- )
-
- return cls.from_dict(config_dict, **kwargs)
-
-
-class FlavaConfig(PretrainedConfig):
- r"""
- [`FlavaConfig`] is the configuration class to store the configuration of a [`FlavaModel`]. It is used to
- instantiate FLAVA model according to the specified arguments, defining the text model, image model, image codebook
- and multimodal model configs. Instantiating a configuration with the defaults will yield a similar configuration to
- that of the FLAVA [facebook/flava-full](https://huggingface.co/facebook/flava-full) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- text_config (`dict`, *optional*):
- Dictionary of configuration options used to initialize [`FlavaTextConfig`].
- image_config (`dict`, *optional*):
- Dictionary of configuration options used to initialize [`FlavaImageConfig`].
- multimodal_config (`dict`, *optional*):
- Dictionary of configuration options used to initialize [`FlavaMultimodalConfig`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- projection_dim (`int`, *optional*, defaults to 512):
- Dimentionality of text and image projection layers.
- logit_scale_init_value (`float`, *optional*, defaults to 2.6592):
- The inital value of the *logit_scale* paramter. Default is used as per the original FLAVA/CLIP
- implementation.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- ce_ignore_index (`int`, *optional*, defaults to -100):
- Cross entropy index to ignore.
- mim_weight (`float`, *optional*, defaults to 1.0):
- Weight to be assigned to MIM (Masked Image Modeling) unimodal loss
- mlm_weight (`float`, *optional*, defaults to 1.0):
- Weight to be assigned to MLM (Masked Language Modeling) unimodal loss
- global_contrastive_weight (`float`, *optional*, defaults to 1.0):
- Weight to be assigned to global contrastive cross-alignment loss.
- itm_weight (`float`, *optional*, defaults to 1.0):
- Weight to be assigned to image-text matching multimodal loss.
- mmm_image_weight (`float`, *optional*, defaults to 1.0):
- Weight to be assigned to MMM loss's image part.
- mmm_text_weight (`float`, *optional*, defaults to 1.0):
- Weight to be assigned to MMM loss's text part.
- global_backprop_contrastive (`bool`, *optional*, defaults to `True`):
- Whether to use global backpropgation through all workers in contrastive loss.
- skip_unmasked_multimodal_encoder (`bool`, *optional*, defaults to `True`):
- Whether to skip running unmasked multimodal encoder whose outputs are not used by FLAVA losses.
- return_loss (`bool`, *optional*, defaults to `True`):
- Whether to return loss or not
-
- kwargs (*optional*):
- Dictionary of keyword arguments.
-
- Example:
-
- ```python
- >>> from transformers import FlavaConfig, FlavaModel, FlavaForPreTraining
-
- >>> # Initializing a FlavaConfig with style configuration
- >>> configuration = FlavaConfig()
-
- >>> # Initializing a FlavaModel and FlavaForPreTraining model (with random weights) from the style configuration
- >>> model = FlavaModel(configuration)
- >>> model_pre = FlavaForPreTraining(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- >>> configuration_pre = model_pre.config
- ```
- """
-
- model_type = "flava"
-
- def __init__(
- self,
- image_config: Dict[str, Any] = None,
- text_config: Dict[str, Any] = None,
- multimodal_config: Dict[str, Any] = None,
- image_codebook_config: Dict[str, Any] = None,
- hidden_size: int = 768,
- layer_norm_eps: float = 1e-12,
- projection_dim: int = 768,
- init_codebook: bool = True,
- logit_scale_init_value: float = 2.6592,
- initializer_range: float = 0.02,
- ce_ignore_index: int = -100,
- mim_weight: float = 1.0,
- mlm_weight: float = 1.0,
- global_contrastive_weight: float = 1.0,
- itm_weight: float = 1.0,
- mmm_image_weight: float = 1.0,
- mmm_text_weight: float = 1.0,
- global_backprop_contrastive: bool = True,
- skip_unmasked_multimodal_encoder: bool = True,
- return_loss: bool = True,
- **kwargs,
- ):
- # If `_config_dict` exist, we use them for the backward compatibility.
- # We pop out these 2 attributes before calling `super().__init__` to avoid them being saved (which causes a lot
- # of confusion!).
- text_config_dict = kwargs.pop("text_config_dict", None)
- image_config_dict = kwargs.pop("image_config_dict", None)
- multimodal_config_dict = kwargs.pop("multimodal_config_dict", None)
- image_codebook_config_dict = kwargs.pop("image_codebook_config_dict", None)
-
- super().__init__(**kwargs)
-
- # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
- # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
- # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
- if text_config_dict is not None:
- if text_config is None:
- text_config = {}
-
- # This is the complete result when using `text_config_dict`.
- _text_config_dict = FlavaTextConfig(**text_config_dict).to_dict()
-
- # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
- for key, value in _text_config_dict.items():
- if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
- # If specified in `text_config_dict`
- if key in text_config_dict:
- message = (
- f"`{key}` is found in both `text_config_dict` and `text_config` but with different values. "
- f'The value `text_config_dict["{key}"]` will be used instead.'
- )
- # If inferred from default argument values (just to be super careful)
- else:
- message = (
- f"`text_config_dict` is provided which will be used to initialize `FlavaTextConfig`. The "
- f'value `text_config["{key}"]` will be overriden.'
- )
- logger.info(message)
-
- # Update all values in `text_config` with the ones in `_text_config_dict`.
- text_config.update(_text_config_dict)
-
- if image_config_dict is not None:
- if image_config is None:
- image_config = {}
-
- # This is the complete result when using `image_config_dict`.
- _image_config_dict = FlavaImageConfig(**image_config_dict).to_dict()
- # convert keys to string instead of integer
- if "id2label" in _image_config_dict:
- _image_config_dict["id2label"] = {
- str(key): value for key, value in _image_config_dict["id2label"].items()
- }
-
- # Give a warning if the values exist in both `_image_config_dict` and `image_config` but being different.
- for key, value in _image_config_dict.items():
- if key in image_config and value != image_config[key] and key not in ["transformers_version"]:
- # If specified in `image_config_dict`
- if key in image_config_dict:
- message = (
- f"`{key}` is found in both `image_config_dict` and `image_config` but with different "
- f'values. The value `image_config_dict["{key}"]` will be used instead.'
- )
- # If inferred from default argument values (just to be super careful)
- else:
- message = (
- f"`image_config_dict` is provided which will be used to initialize `FlavaImageConfig`. "
- f'The value `image_config["{key}"]` will be overriden.'
- )
- logger.info(message)
-
- # Update all values in `image_config` with the ones in `_image_config_dict`.
- image_config.update(_image_config_dict)
-
- if multimodal_config_dict is not None:
- if multimodal_config is None:
- multimodal_config = {}
-
- # This is the complete result when using `multimodal_config_dict`.
- _multimodal_config_dict = FlavaMultimodalConfig(**multimodal_config_dict).to_dict()
-
- # Give a warning if the values exist in both `_multimodal_config_dict` and `multimodal_config` but being
- # different.
- for key, value in _multimodal_config_dict.items():
- if (
- key in multimodal_config
- and value != multimodal_config[key]
- and key not in ["transformers_version"]
- ):
- # If specified in `multimodal_config_dict`
- if key in multimodal_config_dict:
- message = (
- f"`{key}` is found in both `multimodal_config_dict` and `multimodal_config` but with "
- f'different values. The value `multimodal_config_dict["{key}"]` will be used instead.'
- )
- # If inferred from default argument values (just to be super careful)
- else:
- message = (
- f"`multimodal_config_dict` is provided which will be used to initialize "
- f'`FlavaMultimodalConfig`. The value `multimodal_config["{key}"]` will be overriden.'
- )
- logger.info(message)
-
- # Update all values in `multimodal_config` with the ones in `_multimodal_config_dict`.
- multimodal_config.update(_multimodal_config_dict)
-
- if image_codebook_config_dict is not None:
- if image_codebook_config is None:
- image_codebook_config = {}
-
- # This is the complete result when using `image_codebook_config_dict`.
- _image_codebook_config_dict = FlavaImageCodebookConfig(**image_codebook_config_dict).to_dict()
-
- # Give a warning if the values exist in both `_image_codebook_config_dict` and `image_codebook_config` but
- # being different.
- for key, value in _image_codebook_config_dict.items():
- if (
- key in image_codebook_config
- and value != image_codebook_config[key]
- and key not in ["transformers_version"]
- ):
- # If specified in `image_codebook_config_dict`
- if key in image_codebook_config_dict:
- message = (
- f"`{key}` is found in both `image_codebook_config_dict` and `image_codebook_config` but "
- f'with different values. The value `image_codebook_config_dict["{key}"]` will be used '
- "instead."
- )
- # If inferred from default argument values (just to be super careful)
- else:
- message = (
- f"`image_codebook_config_dict` is provided which will be used to initialize "
- f'`FlavaImageCodebookConfig`. The value `image_codebook_config["{key}"]` will be overriden.'
- )
- logger.info(message)
-
- # Update all values in `image_codebook_config` with the ones in `_image_codebook_config_dict`.
- image_codebook_config.update(_image_codebook_config_dict)
-
- if image_config is None:
- image_config = {}
- logger.info("`image_config` is `None`. initializing the `FlavaImageConfig` with default values.")
-
- if text_config is None:
- text_config = {}
- logger.info("`text_config` is `None`. Initializing the `FlavaTextConfig` with default values.")
-
- if multimodal_config is None:
- multimodal_config = {}
- logger.info("`multimodal_config` is `None`. initializing the `FlavaMultimodalConfig` with default values.")
-
- if image_codebook_config is None:
- image_codebook_config = {}
- logger.info(
- "`image_codebook_config` is `None`. initializing the `FlavaImageCodebookConfig` with default values."
- )
-
- self.image_config = FlavaImageConfig(**image_config)
- self.text_config = FlavaTextConfig(**text_config)
- self.multimodal_config = FlavaMultimodalConfig(**multimodal_config)
- self.image_codebook_config = FlavaImageCodebookConfig(**image_codebook_config)
- self.projection_dim = projection_dim
- self.init_codebook = init_codebook
-
- self.hidden_size = hidden_size
- self.layer_norm_eps = layer_norm_eps
- self.initializer_range = initializer_range
- self.logit_scale_init_value = logit_scale_init_value
- self.initializer_factor = 1.0
- self.ce_ignore_index = ce_ignore_index
- self.mim_weight = mim_weight
- self.mlm_weight = mlm_weight
- self.global_contrastive_weight = global_contrastive_weight
- self.itm_weight = itm_weight
- self.mmm_image_weight = mmm_image_weight
- self.mmm_text_weight = mmm_text_weight
- self.global_backprop_contrastive = global_backprop_contrastive
- self.skip_unmasked_multimodal_encoder = skip_unmasked_multimodal_encoder
- self.return_loss = return_loss
-
- @classmethod
- def from_configs(
- cls,
- image_config: FlavaImageConfig,
- text_config: FlavaTextConfig,
- multimodal_config: FlavaMultimodalConfig,
- image_codebook_config: FlavaImageCodebookConfig,
- **kwargs,
- ):
- r"""
- Instantiate a [`FlavaConfig`] (or a derived class) from flava text model configuration, flava image model
- configuration, flava multimodal model and flava codebook model configuration.
-
- Returns:
- [`FlavaConfig`]: An instance of a configuration object
- """
-
- return cls(
- image_config=image_config.to_dict(),
- text_config=text_config.to_dict(),
- multimodal_config=multimodal_config.to_dict(),
- image_codebook_config=image_codebook_config.to_dict(),
- **kwargs,
- )
diff --git a/transformers/models/flava/convert_dalle_to_flava_codebook.py b/transformers/models/flava/convert_dalle_to_flava_codebook.py
deleted file mode 100644
index 7b544125114c85fcf01a881f460ae70472148c85..0000000000000000000000000000000000000000
--- a/transformers/models/flava/convert_dalle_to_flava_codebook.py
+++ /dev/null
@@ -1,102 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import os
-
-import torch
-
-from transformers import FlavaImageCodebook, FlavaImageCodebookConfig
-
-
-def rreplace(s, old, new, occurrence):
- li = s.rsplit(old, occurrence)
- return new.join(li)
-
-
-def count_parameters(state_dict):
- # encoder.embeddings are double copied in original FLAVA
- return sum(param.float().sum() if "encoder.embeddings" not in key else 0 for key, param in state_dict.items())
-
-
-def upgrade_state_dict(state_dict):
- upgrade = {}
-
- group_keys = ["group_1", "group_2", "group_3", "group_4"]
- for key, value in state_dict.items():
- for group_key in group_keys:
- if group_key in key:
- key = key.replace(f"{group_key}.", f"{group_key}.group.")
-
- if "res_path" in key:
- key = key.replace("res_path.", "res_path.path.")
-
- if key.endswith(".w"):
- key = rreplace(key, ".w", ".weight", 1)
- if key.endswith(".b"):
- key = rreplace(key, ".b", ".bias", 1)
-
- upgrade[key] = value.float()
-
- return upgrade
-
-
-@torch.no_grad()
-def convert_dalle_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_path=None, save_checkpoint=True):
- """
- Copy/paste/tweak model's weights to transformers design.
- """
- from dall_e import Encoder
-
- encoder = Encoder()
- if os.path.exists(checkpoint_path):
- ckpt = torch.load(checkpoint_path)
- else:
- ckpt = torch.hub.load_state_dict_from_url(checkpoint_path)
-
- if isinstance(ckpt, Encoder):
- ckpt = ckpt.state_dict()
- encoder.load_state_dict(ckpt)
-
- if config_path is not None:
- config = FlavaImageCodebookConfig.from_pretrained(config_path)
- else:
- config = FlavaImageCodebookConfig()
-
- hf_model = FlavaImageCodebook(config).eval()
- state_dict = encoder.state_dict()
-
- hf_state_dict = upgrade_state_dict(state_dict)
- hf_model.load_state_dict(hf_state_dict)
- hf_state_dict = hf_model.state_dict()
- hf_count = count_parameters(hf_state_dict)
- state_dict_count = count_parameters(state_dict)
-
- assert torch.allclose(hf_count, state_dict_count, atol=1e-3)
-
- if save_checkpoint:
- hf_model.save_pretrained(pytorch_dump_folder_path)
- else:
- return hf_state_dict
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
- parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to flava checkpoint")
- parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
- args = parser.parse_args()
-
- convert_dalle_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)
diff --git a/transformers/models/flava/convert_flava_original_pytorch_to_hf.py b/transformers/models/flava/convert_flava_original_pytorch_to_hf.py
deleted file mode 100644
index 95ebb2bfdb236060037fc91c355dc4f7fe2f62d7..0000000000000000000000000000000000000000
--- a/transformers/models/flava/convert_flava_original_pytorch_to_hf.py
+++ /dev/null
@@ -1,99 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import os
-
-import torch
-
-from transformers import FlavaConfig, FlavaForPreTraining
-from transformers.models.flava.convert_dalle_to_flava_codebook import convert_dalle_checkpoint
-
-
-def count_parameters(state_dict):
- # encoder.embeddings are double copied in original FLAVA
- return sum(param.float().sum() if "encoder.embeddings" not in key else 0 for key, param in state_dict.items())
-
-
-def upgrade_state_dict(state_dict, codebook_state_dict):
- upgrade = {}
-
- for key, value in state_dict.items():
- if "text_encoder.embeddings" in key or "image_encoder.embeddings" in key:
- continue
-
- key = key.replace("heads.cmd.mim_head.cls.predictions", "mmm_image_head")
- key = key.replace("heads.cmd.mlm_head.cls.predictions", "mmm_text_head")
- key = key.replace("heads.cmd.itm_head.cls", "itm_head")
- key = key.replace("heads.cmd.itm_head.pooler", "itm_head.pooler")
- key = key.replace("heads.cmd.clip_head.logit_scale", "flava.logit_scale")
- key = key.replace("heads.fairseq_mlm.cls.predictions", "mlm_head")
- key = key.replace("heads.imagenet.mim_head.cls.predictions", "mim_head")
- key = key.replace("mm_text_projection", "flava.text_to_mm_projection")
- key = key.replace("mm_image_projection", "flava.image_to_mm_projection")
- key = key.replace("image_encoder.module", "flava.image_model")
- key = key.replace("text_encoder.module", "flava.text_model")
- key = key.replace("mm_encoder.module.encoder.cls_token", "flava.multimodal_model.cls_token")
- key = key.replace("mm_encoder.module", "flava.multimodal_model")
- key = key.replace("text_projection", "flava.text_projection")
- key = key.replace("image_projection", "flava.image_projection")
-
- upgrade[key] = value.float()
-
- for key, value in codebook_state_dict.items():
- upgrade[f"image_codebook.{key}"] = value
-
- return upgrade
-
-
-@torch.no_grad()
-def convert_flava_checkpoint(checkpoint_path, codebook_path, pytorch_dump_folder_path, config_path=None):
- """
- Copy/paste/tweak model's weights to transformers design.
- """
- if config_path is not None:
- config = FlavaConfig.from_pretrained(config_path)
- else:
- config = FlavaConfig()
-
- hf_model = FlavaForPreTraining(config).eval()
-
- codebook_state_dict = convert_dalle_checkpoint(codebook_path, None, save_checkpoint=False)
-
- if os.path.exists(checkpoint_path):
- state_dict = torch.load(checkpoint_path, map_location="cpu")
- else:
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_path, map_location="cpu")
-
- hf_state_dict = upgrade_state_dict(state_dict, codebook_state_dict)
- hf_model.load_state_dict(hf_state_dict)
- hf_state_dict = hf_model.state_dict()
- hf_count = count_parameters(hf_state_dict)
- state_dict_count = count_parameters(state_dict) + count_parameters(codebook_state_dict)
-
- assert torch.allclose(hf_count, state_dict_count, atol=1e-3)
-
- hf_model.save_pretrained(pytorch_dump_folder_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
- parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to flava checkpoint")
- parser.add_argument("--codebook_path", default=None, type=str, help="Path to flava codebook checkpoint")
- parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
- args = parser.parse_args()
-
- convert_flava_checkpoint(args.checkpoint_path, args.codebook_path, args.pytorch_dump_folder_path, args.config_path)
diff --git a/transformers/models/flava/feature_extraction_flava.py b/transformers/models/flava/feature_extraction_flava.py
deleted file mode 100644
index c707b575cef2eff9d3dff7e122cc6a875f3e3931..0000000000000000000000000000000000000000
--- a/transformers/models/flava/feature_extraction_flava.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for FLAVA."""
-
-import warnings
-
-from ...utils import logging
-from .image_processing_flava import FlavaImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-class FlavaFeatureExtractor(FlavaImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class FlavaFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
- " use FlavaImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/flava/image_processing_flava.py b/transformers/models/flava/image_processing_flava.py
deleted file mode 100644
index d6a7c8080bb6b4aa9e89f693dd96d3483b6e0e44..0000000000000000000000000000000000000000
--- a/transformers/models/flava/image_processing_flava.py
+++ /dev/null
@@ -1,738 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for Flava."""
-
-import math
-import random
-from functools import lru_cache
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
-from ...image_transforms import resize, to_channel_dimension_format
-from ...image_utils import (
- OPENAI_CLIP_MEAN,
- OPENAI_CLIP_STD,
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import TensorType, is_vision_available, logging
-
-
-if is_vision_available():
- import PIL
-
-
-logger = logging.get_logger(__name__)
-
-
-# These values are taken from CLIP
-FLAVA_IMAGE_MEAN = OPENAI_CLIP_MEAN
-FLAVA_IMAGE_STD = OPENAI_CLIP_STD
-FLAVA_CODEBOOK_MEAN = [0.0, 0.0, 0.0]
-FLAVA_CODEBOOK_STD = [1.0, 1.0, 1.0]
-LOGIT_LAPLACE_EPS: float = 0.1
-
-
-# Inspired from https://github.com/microsoft/unilm/blob/master/beit/masking_generator.py
-class FlavaMaskingGenerator:
- def __init__(
- self,
- input_size: Union[int, Tuple[int, int]] = 14,
- total_mask_patches: int = 75,
- mask_group_max_patches: Optional[int] = None,
- mask_group_min_patches: int = 16,
- mask_group_min_aspect_ratio: Optional[float] = 0.3,
- mask_group_max_aspect_ratio: float = None,
- ):
- if not isinstance(input_size, tuple):
- input_size = (input_size,) * 2
- self.height, self.width = input_size
-
- self.num_patches = self.height * self.width
- self.total_mask_patches = total_mask_patches
-
- self.mask_group_min_patches = mask_group_min_patches
- self.mask_group_max_patches = total_mask_patches if mask_group_max_patches is None else mask_group_max_patches
-
- mask_group_max_aspect_ratio = mask_group_max_aspect_ratio or 1 / mask_group_min_aspect_ratio
- self.log_aspect_ratio = (math.log(mask_group_min_aspect_ratio), math.log(mask_group_max_aspect_ratio))
-
- def __repr__(self):
- repr_str = "MaskingGenerator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
- self.height,
- self.width,
- self.mask_group_min_patches,
- self.mask_group_max_patches,
- self.total_mask_patches,
- self.log_aspect_ratio[0],
- self.log_aspect_ratio[1],
- )
- return repr_str
-
- def get_shape(self):
- return self.height, self.width
-
- def _mask(self, mask, max_mask_patches):
- delta = 0
- for _attempt in range(10):
- target_area = random.uniform(self.mask_group_min_patches, max_mask_patches)
- aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
- height = int(round(math.sqrt(target_area * aspect_ratio)))
- width = int(round(math.sqrt(target_area / aspect_ratio)))
- if width < self.width and height < self.height:
- top = random.randint(0, self.height - height)
- left = random.randint(0, self.width - width)
-
- num_masked = mask[top : top + height, left : left + width].sum()
- # Overlap
- if 0 < height * width - num_masked <= max_mask_patches:
- for i in range(top, top + height):
- for j in range(left, left + width):
- if mask[i, j] == 0:
- mask[i, j] = 1
- delta += 1
-
- if delta > 0:
- break
- return delta
-
- def __call__(self):
- mask = np.zeros(shape=self.get_shape(), dtype=int)
- mask_count = 0
- while mask_count < self.total_mask_patches:
- max_mask_patches = self.total_mask_patches - mask_count
- max_mask_patches = min(max_mask_patches, self.mask_group_max_patches)
-
- delta = self._mask(mask, max_mask_patches)
- if delta == 0:
- break
- else:
- mask_count += delta
-
- return mask
-
-
-class FlavaImageProcessor(BaseImageProcessor):
- r"""
- Constructs a Flava image processor.
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by the
- `do_resize` parameter in `preprocess`.
- size (`Dict[str, int]` *optional*, defaults to `{"height": 224, "width": 224}`):
- Size of the image after resizing. Can be overridden by the `size` parameter in `preprocess`.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
- Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in
- `preprocess`.
- do_center_crop (`bool`, *optional*, defaults to `True`):
- Whether to center crop the images. Can be overridden by the `do_center_crop` parameter in `preprocess`.
- crop_size (`Dict[str, int]` *optional*, defaults to `{"height": 224, "width": 224}`):
- Size of image after the center crop `(crop_size["height"], crop_size["width"])`. Can be overridden by the
- `crop_size` parameter in `preprocess`.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`
- parameter in `preprocess`.
- rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in
- `preprocess`.
- do_normalize (`bool`, *optional*, defaults to `True`):
- Whether to normalize the image. Can be overridden by the `do_normalize` parameter in `preprocess`.
- image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
- Mean to use if normalizing the image. This is a float or list of floats the length of the number of
- channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
- Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
- number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
- return_image_mask (`bool`, *optional*, defaults to `False`):
- Whether to return the image mask. Can be overridden by the `return_image_mask` parameter in `preprocess`.
- input_size_patches (`int`, *optional*, defaults to 14):
- Number of patches in the image in height and width direction. 14x14 = 196 total patches. Can be overridden
- by the `input_size_patches` parameter in `preprocess`.
- total_mask_patches (`int`, *optional*, defaults to 75):
- Total number of patches that should be masked. Can be overridden by the `total_mask_patches` parameter in
- `preprocess`.
- mask_group_min_patches (`int`, *optional*, defaults to 16):
- Minimum number of patches that should be masked. Can be overridden by the `mask_group_min_patches`
- parameter in `preprocess`.
- mask_group_max_patches (`int`, *optional*):
- Maximum number of patches that should be masked. Can be overridden by the `mask_group_max_patches`
- parameter in `preprocess`.
- mask_group_min_aspect_ratio (`float`, *optional*, defaults to 0.3):
- Minimum aspect ratio of the mask window. Can be overridden by the `mask_group_min_aspect_ratio` parameter
- in `preprocess`.
- mask_group_max_aspect_ratio (`float`, *optional*):
- Maximum aspect ratio of the mask window. Can be overridden by the `mask_group_max_aspect_ratio` parameter
- in `preprocess`.
- codebook_do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the input for codebook to a certain. Can be overridden by the `codebook_do_resize`
- parameter in `preprocess`. `codebook_size`.
- codebook_size (`Dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
- Resize the input for codebook to the given size. Can be overridden by the `codebook_size` parameter in
- `preprocess`.
- codebook_resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.LANCZOS`):
- Resampling filter to use if resizing the codebook image. Can be overridden by the `codebook_resample`
- parameter in `preprocess`.
- codebook_do_center_crop (`bool`, *optional*, defaults to `True`):
- Whether to crop the input for codebook at the center. If the input size is smaller than
- `codebook_crop_size` along any edge, the image is padded with 0's and then center cropped. Can be
- overridden by the `codebook_do_center_crop` parameter in `preprocess`.
- codebook_crop_size (`Dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
- Desired output size for codebook input when applying center-cropping. Can be overridden by the
- `codebook_crop_size` parameter in `preprocess`.
- codebook_do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the input for codebook by the specified scale `codebook_rescale_factor`. Can be
- overridden by the `codebook_do_rescale` parameter in `preprocess`.
- codebook_rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
- Defines the scale factor to use if rescaling the codebook image. Can be overridden by the
- `codebook_rescale_factor` parameter in `preprocess`.
- codebook_do_map_pixels (`bool`, *optional*, defaults to `True`):
- Whether to map the pixel values of the codebook input to (1 - 2e)x + e. Can be overridden by the
- `codebook_do_map_pixels` parameter in `preprocess`.
- codebook_do_normalize (`bool`, *optional*, defaults to `True`):
- Whether or not to normalize the input for codebook with `codebook_image_mean` and `codebook_image_std`. Can
- be overridden by the `codebook_do_normalize` parameter in `preprocess`.
- codebook_image_mean (`Optional[Union[float, Iterable[float]]]`, *optional*, defaults to `[0, 0, 0]`):
- The sequence of means for each channel, to be used when normalizing images for codebook. Can be overridden
- by the `codebook_image_mean` parameter in `preprocess`.
- codebook_image_std (`Optional[Union[float, Iterable[float]]]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
- The sequence of standard deviations for each channel, to be used when normalizing images for codebook. Can
- be overridden by the `codebook_image_std` parameter in `preprocess`.
- """
-
- model_input_names = ["pixel_values"]
-
- def __init__(
- self,
- do_resize: bool = True,
- size: Dict[str, int] = None,
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- do_center_crop: bool = True,
- crop_size: Dict[str, int] = None,
- do_rescale: bool = True,
- rescale_factor: Union[int, float] = 1 / 255,
- do_normalize: bool = True,
- image_mean: Optional[Union[float, Iterable[float]]] = None,
- image_std: Optional[Union[float, Iterable[float]]] = None,
- # Mask related params
- return_image_mask: bool = False,
- input_size_patches: int = 14,
- total_mask_patches: int = 75,
- mask_group_min_patches: int = 16,
- mask_group_max_patches: Optional[int] = None,
- mask_group_min_aspect_ratio: float = 0.3,
- mask_group_max_aspect_ratio: Optional[float] = None,
- # Codebook related params
- return_codebook_pixels: bool = False,
- codebook_do_resize: bool = True,
- codebook_size: bool = None,
- codebook_resample: int = PILImageResampling.LANCZOS,
- codebook_do_center_crop: bool = True,
- codebook_crop_size: int = None,
- codebook_do_rescale: bool = True,
- codebook_rescale_factor: Union[int, float] = 1 / 255,
- codebook_do_map_pixels: bool = True,
- codebook_do_normalize: bool = True,
- codebook_image_mean: Optional[Union[float, Iterable[float]]] = None,
- codebook_image_std: Optional[Union[float, Iterable[float]]] = None,
- **kwargs,
- ) -> None:
- super().__init__(**kwargs)
- size = size if size is not None else {"height": 224, "width": 224}
- size = get_size_dict(size)
- crop_size = crop_size if crop_size is not None else {"height": 224, "width": 224}
- crop_size = get_size_dict(crop_size, param_name="crop_size")
-
- codebook_size = codebook_size if codebook_size is not None else {"height": 112, "width": 112}
- codebook_size = get_size_dict(codebook_size, param_name="codebook_size")
- codebook_crop_size = codebook_crop_size if codebook_crop_size is not None else {"height": 112, "width": 112}
- codebook_crop_size = get_size_dict(codebook_crop_size, param_name="codebook_crop_size")
-
- self.do_resize = do_resize
- self.size = size
- self.resample = resample
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.do_center_crop = do_center_crop
- self.crop_size = crop_size
- self.do_normalize = do_normalize
- self.image_mean = image_mean if image_mean is not None else FLAVA_IMAGE_MEAN
- self.image_std = image_std if image_std is not None else FLAVA_IMAGE_STD
-
- self.return_image_mask = return_image_mask
- self.input_size_patches = input_size_patches
- self.total_mask_patches = total_mask_patches
- self.mask_group_min_patches = mask_group_min_patches
- self.mask_group_max_patches = mask_group_max_patches
- self.mask_group_min_aspect_ratio = mask_group_min_aspect_ratio
- self.mask_group_max_aspect_ratio = mask_group_max_aspect_ratio
-
- self.return_codebook_pixels = return_codebook_pixels
- self.codebook_do_resize = codebook_do_resize
- self.codebook_size = codebook_size
- self.codebook_resample = codebook_resample
- self.codebook_do_center_crop = codebook_do_center_crop
- self.codebook_crop_size = codebook_crop_size
- self.codebook_do_rescale = codebook_do_rescale
- self.codebook_rescale_factor = codebook_rescale_factor
- self.codebook_do_map_pixels = codebook_do_map_pixels
- self.codebook_do_normalize = codebook_do_normalize
- self.codebook_image_mean = codebook_image_mean
- self.codebook_image_mean = codebook_image_mean if codebook_image_mean is not None else FLAVA_CODEBOOK_MEAN
- self.codebook_image_std = codebook_image_std if codebook_image_std is not None else FLAVA_CODEBOOK_STD
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size",
- "resample",
- "do_center_crop",
- "crop_size",
- "do_rescale",
- "rescale_factor",
- "do_normalize",
- "image_mean",
- "image_std",
- "return_image_mask",
- "input_size_patches",
- "total_mask_patches",
- "mask_group_min_patches",
- "mask_group_max_patches",
- "mask_group_min_aspect_ratio",
- "mask_group_max_aspect_ratio",
- "return_codebook_pixels",
- "codebook_do_resize",
- "codebook_size",
- "codebook_resample",
- "codebook_do_center_crop",
- "codebook_crop_size",
- "codebook_do_rescale",
- "codebook_rescale_factor",
- "codebook_do_map_pixels",
- "codebook_do_normalize",
- "codebook_image_mean",
- "codebook_image_std",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- @classmethod
- def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
- """
- Overrides the `from_dict` method from the base class to make sure parameters are updated if image processor is
- created using from_dict and kwargs e.g. `FlavaImageProcessor.from_pretrained(checkpoint, codebook_size=600)`
- """
- image_processor_dict = image_processor_dict.copy()
- if "codebook_size" in kwargs:
- image_processor_dict["codebook_size"] = kwargs.pop("codebook_size")
- if "codebook_crop_size" in kwargs:
- image_processor_dict["codebook_crop_size"] = kwargs.pop("codebook_crop_size")
- return super().from_dict(image_processor_dict, **kwargs)
-
- @lru_cache()
- def masking_generator(
- self,
- input_size_patches,
- total_mask_patches,
- mask_group_min_patches,
- mask_group_max_patches,
- mask_group_min_aspect_ratio,
- mask_group_max_aspect_ratio,
- ) -> FlavaMaskingGenerator:
- return FlavaMaskingGenerator(
- input_size=input_size_patches,
- total_mask_patches=total_mask_patches,
- mask_group_min_patches=mask_group_min_patches,
- mask_group_max_patches=mask_group_max_patches,
- mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
- mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
- )
-
- # Copied from transformers.models.vit.image_processing_vit.ViTImageProcessor.resize with PILImageResampling.BILINEAR->PILImageResampling.BICUBIC
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BICUBIC,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize an image to `(size["height"], size["width"])`.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
- `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BICUBIC`.
- data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
-
- Returns:
- `np.ndarray`: The resized image.
- """
- size = get_size_dict(size)
- if "height" not in size or "width" not in size:
- raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}")
- output_size = (size["height"], size["width"])
- return resize(
- image,
- size=output_size,
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
-
- def map_pixels(self, image: np.ndarray) -> np.ndarray:
- return (1 - 2 * LOGIT_LAPLACE_EPS) * image + LOGIT_LAPLACE_EPS
-
- def _preprocess_image(
- self,
- image: ImageInput,
- do_resize: bool = None,
- size: Dict[str, int] = None,
- resample: PILImageResampling = None,
- do_center_crop: bool = None,
- crop_size: Dict[str, int] = None,
- do_rescale: bool = None,
- rescale_factor: float = None,
- do_normalize: bool = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- do_map_pixels: bool = None,
- data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
- input_data_format: Optional[ChannelDimension] = None,
- ) -> np.ndarray:
- """Preprocesses a single image."""
-
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_center_crop=do_center_crop,
- crop_size=crop_size,
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
-
- # All transformations expect numpy arrays.
- image = to_numpy_array(image)
-
- if is_scaled_image(image) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(image)
-
- if do_resize:
- image = self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)
-
- if do_center_crop:
- image = self.center_crop(image=image, size=crop_size, input_data_format=input_data_format)
-
- if do_rescale:
- image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
-
- if do_normalize:
- image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
-
- if do_map_pixels:
- image = self.map_pixels(image)
-
- if data_format is not None:
- image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
- return image
-
- def preprocess(
- self,
- images: ImageInput,
- do_resize: Optional[bool] = None,
- size: Dict[str, int] = None,
- resample: PILImageResampling = None,
- do_center_crop: Optional[bool] = None,
- crop_size: Optional[Dict[str, int]] = None,
- do_rescale: Optional[bool] = None,
- rescale_factor: Optional[float] = None,
- do_normalize: Optional[bool] = None,
- image_mean: Optional[Union[float, List[float]]] = None,
- image_std: Optional[Union[float, List[float]]] = None,
- # Mask related params
- return_image_mask: Optional[bool] = None,
- input_size_patches: Optional[int] = None,
- total_mask_patches: Optional[int] = None,
- mask_group_min_patches: Optional[int] = None,
- mask_group_max_patches: Optional[int] = None,
- mask_group_min_aspect_ratio: Optional[float] = None,
- mask_group_max_aspect_ratio: Optional[float] = None,
- # Codebook related params
- return_codebook_pixels: Optional[bool] = None,
- codebook_do_resize: Optional[bool] = None,
- codebook_size: Optional[Dict[str, int]] = None,
- codebook_resample: Optional[int] = None,
- codebook_do_center_crop: Optional[bool] = None,
- codebook_crop_size: Optional[Dict[str, int]] = None,
- codebook_do_rescale: Optional[bool] = None,
- codebook_rescale_factor: Optional[float] = None,
- codebook_do_map_pixels: Optional[bool] = None,
- codebook_do_normalize: Optional[bool] = None,
- codebook_image_mean: Optional[Iterable[float]] = None,
- codebook_image_std: Optional[Iterable[float]] = None,
- return_tensors: Optional[Union[str, TensorType]] = None,
- data_format: ChannelDimension = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> PIL.Image.Image:
- """
- Preprocess an image or batch of images.
-
- Args:
- images (`ImageInput`):
- Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
- passing in images with pixel values between 0 and 1, set `do_rescale=False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the image.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
- Size of the image.
- resample (`int`, *optional*, defaults to `self.resample`):
- Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`, Only
- has an effect if `do_resize` is set to `True`.
- do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`):
- Whether to center crop the image.
- crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
- Size of the center crop. Only has an effect if `do_center_crop` is set to `True`.
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether to rescale the image values between [0 - 1].
- rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
- Rescale factor to rescale the image by if `do_rescale` is set to `True`.
- do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
- Whether to normalize the image.
- image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
- Image mean.
- image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
- Image standard deviation.
- return_image_mask (`bool`, *optional*, defaults to `self.return_image_mask`):
- Whether to return the image mask.
- input_size_patches (`int`, *optional*, defaults to `self.input_size_patches`):
- Size of the patches to extract from the image.
- total_mask_patches (`int`, *optional*, defaults to `self.total_mask_patches`):
- Total number of patches to extract from the image.
- mask_group_min_patches (`int`, *optional*, defaults to `self.mask_group_min_patches`):
- Minimum number of patches to extract from the image.
- mask_group_max_patches (`int`, *optional*, defaults to `self.mask_group_max_patches`):
- Maximum number of patches to extract from the image.
- mask_group_min_aspect_ratio (`float`, *optional*, defaults to `self.mask_group_min_aspect_ratio`):
- Minimum aspect ratio of the patches to extract from the image.
- mask_group_max_aspect_ratio (`float`, *optional*, defaults to `self.mask_group_max_aspect_ratio`):
- Maximum aspect ratio of the patches to extract from the image.
- return_codebook_pixels (`bool`, *optional*, defaults to `self.return_codebook_pixels`):
- Whether to return the codebook pixels.
- codebook_do_resize (`bool`, *optional*, defaults to `self.codebook_do_resize`):
- Whether to resize the codebook pixels.
- codebook_size (`Dict[str, int]`, *optional*, defaults to `self.codebook_size`):
- Size of the codebook pixels.
- codebook_resample (`int`, *optional*, defaults to `self.codebook_resample`):
- Resampling filter to use if resizing the codebook pixels. This can be one of the enum
- `PILImageResampling`, Only has an effect if `codebook_do_resize` is set to `True`.
- codebook_do_center_crop (`bool`, *optional*, defaults to `self.codebook_do_center_crop`):
- Whether to center crop the codebook pixels.
- codebook_crop_size (`Dict[str, int]`, *optional*, defaults to `self.codebook_crop_size`):
- Size of the center crop of the codebook pixels. Only has an effect if `codebook_do_center_crop` is set
- to `True`.
- codebook_do_rescale (`bool`, *optional*, defaults to `self.codebook_do_rescale`):
- Whether to rescale the codebook pixels values between [0 - 1].
- codebook_rescale_factor (`float`, *optional*, defaults to `self.codebook_rescale_factor`):
- Rescale factor to rescale the codebook pixels by if `codebook_do_rescale` is set to `True`.
- codebook_do_map_pixels (`bool`, *optional*, defaults to `self.codebook_do_map_pixels`):
- Whether to map the codebook pixels values.
- codebook_do_normalize (`bool`, *optional*, defaults to `self.codebook_do_normalize`):
- Whether to normalize the codebook pixels.
- codebook_image_mean (`float` or `List[float]`, *optional*, defaults to `self.codebook_image_mean`):
- Codebook pixels mean to normalize the codebook pixels by if `codebook_do_normalize` is set to `True`.
- codebook_image_std (`float` or `List[float]`, *optional*, defaults to `self.codebook_image_std`):
- Codebook pixels standard deviation to normalize the codebook pixels by if `codebook_do_normalize` is
- set to `True`.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- do_resize = do_resize if do_resize is not None else self.do_resize
- size = size if size is not None else self.size
- size = get_size_dict(size)
- resample = resample if resample is not None else self.resample
- do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
- crop_size = crop_size if crop_size is not None else self.crop_size
- crop_size = get_size_dict(crop_size, param_name="crop_size")
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- do_normalize = do_normalize if do_normalize is not None else self.do_normalize
- image_mean = image_mean if image_mean is not None else self.image_mean
- image_std = image_std if image_std is not None else self.image_std
-
- return_image_mask = return_image_mask if return_image_mask is not None else self.return_image_mask
- input_size_patches = input_size_patches if input_size_patches is not None else self.input_size_patches
- total_mask_patches = total_mask_patches if total_mask_patches is not None else self.total_mask_patches
- mask_group_min_patches = (
- mask_group_min_patches if mask_group_min_patches is not None else self.mask_group_min_patches
- )
- mask_group_max_patches = (
- mask_group_max_patches if mask_group_max_patches is not None else self.mask_group_max_patches
- )
- mask_group_min_aspect_ratio = (
- mask_group_min_aspect_ratio
- if mask_group_min_aspect_ratio is not None
- else self.mask_group_min_aspect_ratio
- )
- mask_group_max_aspect_ratio = (
- mask_group_max_aspect_ratio
- if mask_group_max_aspect_ratio is not None
- else self.mask_group_max_aspect_ratio
- )
-
- return_codebook_pixels = (
- return_codebook_pixels if return_codebook_pixels is not None else self.return_codebook_pixels
- )
- codebook_do_resize = codebook_do_resize if codebook_do_resize is not None else self.codebook_do_resize
- codebook_size = codebook_size if codebook_size is not None else self.codebook_size
- codebook_size = get_size_dict(codebook_size, param_name="codebook_size")
- codebook_resample = codebook_resample if codebook_resample is not None else self.codebook_resample
- codebook_do_rescale = codebook_do_rescale if codebook_do_rescale is not None else self.codebook_do_rescale
- codebook_rescale_factor = (
- codebook_rescale_factor if codebook_rescale_factor is not None else self.codebook_rescale_factor
- )
- codebook_do_center_crop = (
- codebook_do_center_crop if codebook_do_center_crop is not None else self.codebook_do_center_crop
- )
- codebook_crop_size = codebook_crop_size if codebook_crop_size is not None else self.codebook_crop_size
- codebook_crop_size = get_size_dict(codebook_crop_size, param_name="codebook_crop_size")
- codebook_do_map_pixels = (
- codebook_do_map_pixels if codebook_do_map_pixels is not None else self.codebook_do_map_pixels
- )
- codebook_do_normalize = (
- codebook_do_normalize if codebook_do_normalize is not None else self.codebook_do_normalize
- )
- codebook_image_mean = codebook_image_mean if codebook_image_mean is not None else self.codebook_image_mean
- codebook_image_std = codebook_image_std if codebook_image_std is not None else self.codebook_image_std
-
- images = make_list_of_images(images)
-
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
-
- processed_images = [
- self._preprocess_image(
- image=img,
- do_resize=do_resize,
- size=size,
- resample=resample,
- do_center_crop=do_center_crop,
- crop_size=crop_size,
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_map_pixels=False,
- data_format=data_format,
- input_data_format=input_data_format,
- )
- for img in images
- ]
- data = {"pixel_values": processed_images}
-
- if return_codebook_pixels:
- codebook_images = [
- self._preprocess_image(
- image=img,
- do_resize=codebook_do_resize,
- size=codebook_size,
- resample=codebook_resample,
- do_center_crop=codebook_do_center_crop,
- crop_size=codebook_crop_size,
- do_rescale=codebook_do_rescale,
- rescale_factor=codebook_rescale_factor,
- do_normalize=codebook_do_normalize,
- image_mean=codebook_image_mean,
- image_std=codebook_image_std,
- do_map_pixels=codebook_do_map_pixels,
- data_format=data_format,
- input_data_format=input_data_format,
- )
- for img in images
- ]
- data["codebook_pixel_values"] = codebook_images
-
- if return_image_mask:
- mask_generator = self.masking_generator(
- input_size_patches=input_size_patches,
- total_mask_patches=total_mask_patches,
- mask_group_min_patches=mask_group_min_patches,
- mask_group_max_patches=mask_group_max_patches,
- mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
- mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
- )
- masks = [mask_generator() for _ in images]
- data["bool_masked_pos"] = masks
-
- return BatchFeature(data=data, tensor_type=return_tensors)
diff --git a/transformers/models/flava/modeling_flava.py b/transformers/models/flava/modeling_flava.py
deleted file mode 100644
index 19f19d4c9d5666b31a888c3c5cdaf742f530801f..0000000000000000000000000000000000000000
--- a/transformers/models/flava/modeling_flava.py
+++ /dev/null
@@ -1,2098 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch FLAVA model."""
-
-import collections
-import math
-from collections import OrderedDict
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from ...activations import ACT2FN
-from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
-from ...modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_flava import (
- FlavaConfig,
- FlavaImageCodebookConfig,
- FlavaImageConfig,
- FlavaMultimodalConfig,
- FlavaTextConfig,
-)
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "facebook/flava-full"
-
-# Codebook docstring
-_CHECKPOINT_FOR_CODEBOOK_DOC = "facebook/flava-image-codebook"
-_CONFIG_CLASS_FOR_IMAGE_MODEL_DOC = "FlavaImageConfig"
-_CONFIG_CLASS_FOR_TEXT_MODEL_DOC = "FlavaTextConfig"
-_CONFIG_CLASS_FOR_MULTIMODAL_MODEL_DOC = "FlavaMultimodalConfig"
-_EXPECTED_IMAGE_OUTPUT_SHAPE = [1, 197, 768]
-
-from ..deprecated._archive_maps import FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-FLAVA_CODEBOOK_PRETRAINED_MODEL_ARCHIVE_LIST = ["facebook/flava-image-codebook"]
-LOGIT_SCALE_CLAMP_MIN = 0
-LOGIT_SCALE_CLAMP_MAX = 4.6052
-
-FlavaPossibleConfigs = Union[FlavaTextConfig, FlavaImageConfig, FlavaMultimodalConfig]
-
-
-@dataclass
-class FlavaModelOutput(ModelOutput):
- """
- Output from FlavaModel containing embeddings and outputs from individual encoders.
-
- Note that `image_embeddings` and `text_embeddigns` returned are similar to pooled output returned from a
- transformer. If you want embeddings for contrastive loss or retrieval use a FLAVA model's `image_projection` and
- `text_projection` layers on `image_embeddings` and `text_embeddings` respectively.
-
- Args:
- image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present):
- The image embeddings which are basically the pooled output of [`FlavaImageModel`].
- image_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present):
- The output of the [`FlavaImageModel`].
- text_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` are present):
- The text embeddings which are basically the pooled output of [`FlavaTextModel`].
- text_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids` are present):
- The output of the [`FlavaTextModel`].
- multimodal_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present and `skip_multimodal_encoder` is `None` or `False`):
- The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`].
- multimodal_output (`BaseModelOutputWithPooling`, returned when `input_ids` and `pixel_values` are present and `skip_multimodal_encoder` is `None` or `False`):
- The output of the [`FlavaMultimodalModel`].
- """
-
- image_embeddings: Optional[torch.FloatTensor] = None
- image_output: Optional[BaseModelOutputWithPooling] = None
- text_embeddings: Optional[torch.FloatTensor] = None
- text_output: Optional[BaseModelOutputWithPooling] = None
- multimodal_embeddings: Optional[torch.FloatTensor] = None
- multimodal_output: Optional[BaseModelOutputWithPooling] = None
-
- def to_tuple(self) -> Tuple[Any]:
- return tuple(
- self[k] if k not in ["text_output", "image_output", "multimodal_output"] else getattr(self, k).to_tuple()
- for k in self.keys()
- )
-
-
-@dataclass
-class FlavaLosses(ModelOutput):
- """Class representing pretraining losses from FLAVA model
-
- Args:
- mim (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mim_labels` and `pixel_values` are present, `input_ids_masked` is absent and `mim_weight` > 0.:
- Masked Image Modeling loss as used in BeIT calculated only for unimodal image data.
- mlm (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mlm_labels` and `input_ids_masked` are present, `pixel_values` is absent and `mlm_weight` > 0.:
- Masked Language Modeling loss as used in BERT calculated only for unimodal text data.
- itm (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `itm_labels`, `input_ids_masked`, `pixel_values` are present and `itm_weight` > 0.:
- Image Text Matching (ITM) loss calculated for paired image-text data. Note that ITM loss is calculated on
- masked pairs in FLAVA.
- global_contrastive (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `input_ids` and `pixel_values` are present and `global_contrastive_weight` > 0.:
- Contrastive loss for image-text similarity similar to CLIP but calculated globally for paired image-text
- data. This is calculated on unmasked images and texts.
- mmm_image (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mim_labels`, `pixel_values` and `input_ids_masked` are present and `mmm_image_weight` > 0.:
- Masked Multimodal Modeling loss's image component calculated on paired image-text data.
- mmm_text (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mlm_labels`, `pixel_values` and `input_ids_masked` are present and `mmm_text_weight` > 0.:
- Masked Multimodal Modeling loss's text component calculated on paired image-text data.
- """
-
- mim: Optional[torch.FloatTensor] = None
- mlm: Optional[torch.FloatTensor] = None
- itm: Optional[torch.FloatTensor] = None
- global_contrastive: Optional[torch.FloatTensor] = None
- mmm_image: Optional[torch.FloatTensor] = None
- mmm_text: Optional[torch.FloatTensor] = None
-
- def all_none(self) -> bool:
- all_none = True
- for v in self.values():
- if v is not None:
- all_none = False
- break
- return all_none
-
-
-@dataclass
-class FlavaForPreTrainingOutput(ModelOutput):
- """
- Output from FlavaForPreTraining containing embeddings, and outputs from individual encoders.
-
- Note that `image_embeddings` and `text_embeddings` returned are similar to pooled output returned from a
- transformer. If you want embeddings for contrastive loss or retrieval use a FLAVA model's `image_projection` and
- `text_projection` layers on `image_embeddings` and `text_embeddings` respectively.
-
- Args:
- loss (`torch.FloatTensor`, *optional*, returned when `return_loss` is True):
- Total loss calculated for this model.
- loss_info (`FlavaLosses`):
- Detailed info for FLAVA Pretraining losses. Check `FlavaLosses` class description for the information on
- the keys.
- image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present):
- The image embeddings which are basically the pooled output of [`FlavaImageModel`].
- image_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present):
- The output of the [`FlavaImageModel`].
- text_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` are present):
- The text embeddings which are basically the pooled output of [`FlavaTextModel`].
- text_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids` are present):
- The output of the [`FlavaTextModel`].
- multimodal_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present and `skip_unmasked_multimodal_encoder` is `None` or `False`):
- The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`].
- multimodal_output (`BaseModelOutputWithPooling`, returned when `input_ids` and `pixel_values` are present and `skip_unmasked_multimodal_encoder` is `None` or `False`):
- The output of the [`FlavaMultimodalModel`].
-
- image_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present):
- The image embeddings which are basically the pooled output of [`FlavaImageModel`]. Uses `bool_masked_pos`
- to create masked images.
- image_masked_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present):
- The output of the [`FlavaImageModel`]. Uses `bool_masked_pos` to create masked images.
- text_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids_masked` are present):
- The text embeddings which are basically the pooled output of [`FlavaTextModel`].
- text_masked_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids_masked` are present):
- The output of the [`FlavaTextModel`].
- multimodal_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present):
- The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`].
- multimodal_masked_output (`BaseModelOutputWithPooling`, returned when `input_ids_masked` and `pixel_values` are present):
- The output of the [`FlavaMultimodalModel`].
-
- mim_logits (`torch.FloatTensor` of shape `(batch_size, num_image_patches, image_vocab_size)` or of shape `(total_masked_patches, image_vocab_size)` , *optional*, returned when `pixel_values` are present and `input_ids_masked` are not):
- The logits for MIM unimodal loss. Uses `book_masked_pos` to get masked patches. The flattened output is
- returned when `bool_masked_pos` has some of the patches masked.
- mlm_logits (`torch.FloatTensor` of shape `(batch_size, text_seq_length, text_vocab_size)` or of shape `(total_masked_seq_length, text_vocab_size)`, *optional*, returned when `input_ids_masked` are present and `pixel_values` are not):
- The logits for MLM unimodal loss. The flattened output is returned when `input_ids_masked` has some of
- the tokens masked.
- itm_logits (`torch.FloatTensor` of shape `(batch_size, 2)`, *optional*, returned when `input_ids_masked` and `pixel_values` are present):
- The logits for ITM loss. Note that ITM loss is calculated on masked pairs in FLAVA.
- mmm_image_logits (`torch.FloatTensor` of shape `(batch_size, num_image_patches, image_vocab_size)` or of shape`(total_masked_patches, image_vocab_size)`, *optional*, returned when `pixel_values` and `input_ids_masked` are present):
- The logits for MMM image multimodal loss. Uses `book_masked_pos` to get masked patches. The flattened
- output is returned when `bool_masked_pos` has some of the patches masked.
- mmm_text_logits (`torch.FloatTensor` of shape `(batch_size, text_seq_length, text_vocab_size)` or of shape `(`(total_masked_seq_length, text_vocab_size)`), *optional*, returned when `pixel_values` and `input_ids_masked` are present):
- The logits for MMM text multimodal loss. The flattened output is returned when `input_ids_masked` has
- some of the tokens masked.
- contrastive_logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):
- The scaled dot product scores between `image_embeddings` and `text_embeddings` but passed through FLAVA's
- `image_projection` and `text_projection` layers respectively. This represents the image-text similarity
- scores. This is calculated on unmasked images and texts.
- contrastive_logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):
- The scaled dot product scores between `text_embeddings` and `image_embeddings` but passed through FLAVA's
- `text_projection` and `image_projection` layers respectively. This is calculated on unmasked images and
- texts.
- """
-
- loss: Optional[torch.FloatTensor] = None
- loss_info: FlavaLosses = None
- image_embeddings: Optional[torch.FloatTensor] = None
- image_output: Optional[BaseModelOutputWithPooling] = None
- text_embeddings: Optional[torch.FloatTensor] = None
- text_output: Optional[BaseModelOutputWithPooling] = None
- multimodal_embeddings: Optional[torch.FloatTensor] = None
- multimodal_output: Optional[BaseModelOutputWithPooling] = None
- image_masked_embeddings: Optional[torch.FloatTensor] = None
- image_masked_output: Optional[BaseModelOutputWithPooling] = None
- text_masked_embeddings: Optional[torch.FloatTensor] = None
- text_masked_output: Optional[BaseModelOutputWithPooling] = None
- multimodal_masked_embeddings: Optional[torch.FloatTensor] = None
- multimodal_masked_output: Optional[BaseModelOutputWithPooling] = None
- mim_logits: Optional[torch.FloatTensor] = None
- mlm_logits: Optional[torch.FloatTensor] = None
- itm_logits: Optional[torch.FloatTensor] = None
- contrastive_logits_per_image: Optional[torch.FloatTensor] = None
- contrastive_logits_per_text: Optional[torch.FloatTensor] = None
- mmm_image_logits: Optional[torch.FloatTensor] = None
- mmm_text_logits: Optional[torch.FloatTensor] = None
-
- def to_tuple(self) -> Tuple[Any]:
- transformer_outputs = [
- "text_output",
- "image_output",
- "multimodal_output",
- "text_masked_output",
- "image_masked_output",
- "multimodal_masked_output",
- ]
- return tuple(self[k] if k not in transformer_outputs else getattr(self, k).to_tuple() for k in self.keys())
-
-
-# Based on timm implementation, which can be found here:
-# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/image_transformer.py
-class FlavaImageEmbeddings(nn.Module):
- """
- Construct the CLS token, position and patch embeddings. Optionally, also the mask token.
- """
-
- def __init__(self, config: FlavaImageConfig, use_mask_token: bool = False) -> None:
- super().__init__()
-
- use_mask_token = use_mask_token or config.mask_token
- self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
- self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) if use_mask_token else None
- self.patch_embeddings = PatchEmbeddings(
- image_size=config.image_size,
- patch_size=config.patch_size,
- num_channels=config.num_channels,
- embed_dim=config.hidden_size,
- )
- num_patches = self.patch_embeddings.num_patches
- self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.config = config
-
- def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
- """
- This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
- resolution images.
-
- Source:
- https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/image_transformer.py#L174
- """
-
- npatch = embeddings.shape[1] - 1
- num_pos = self.position_embeddings.shape[1] - 1
- if npatch == num_pos and height == width:
- return self.position_embeddings
- class_pos_embed = self.position_embeddings[:, 0]
- patch_pos_embed = self.position_embeddings[:, 1:]
- dim = embeddings.shape[-1]
- num_h_patches = height // self.config.patch_size
- num_w_patches = width // self.config.patch_size
- # we add a small number to avoid floating point error in the interpolation
- # see discussion at https://github.com/facebookresearch/dino/issues/8
- num_h_patches, num_w_patches = num_h_patches + 0.1, num_w_patches + 0.1
- patch_pos_embed = nn.functional.interpolate(
- patch_pos_embed.reshape(1, int(math.sqrt(num_pos)), int(math.sqrt(num_pos)), dim).permute(0, 3, 1, 2),
- scale_factor=(num_h_patches / math.sqrt(num_pos), num_w_patches / math.sqrt(num_pos)),
- mode="bicubic",
- align_corners=False,
- )
- if int(num_h_patches) != patch_pos_embed.shape[-2] or int(num_w_patches) != patch_pos_embed.shape[-1]:
- raise ValueError(
- f"Number of patches for images ({int(num_h_patches), int(num_w_patches)}) don't match the "
- f"shape of position embedding ({patch_pos_embed.shape[-2], patch_pos_embed.shape[-1]})"
- )
- patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
- return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
-
- def forward(
- self,
- pixel_values: torch.Tensor,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- interpolate_pos_encoding: bool = False,
- ) -> torch.Tensor:
- batch_size, num_channels, height, width = pixel_values.shape
- embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
-
- batch_size, seq_len, _ = embeddings.size()
- if bool_masked_pos is not None:
- mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
- # B X H X W = B X HW
- if bool_masked_pos.dim() == 3:
- bool_masked_pos = bool_masked_pos.view(bool_masked_pos.size(0), -1)
- # replace the masked visual tokens by mask_tokens
- mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
- embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
-
- # add the [CLS] token to the embedded patch tokens
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- embeddings = torch.cat((cls_tokens, embeddings), dim=1)
-
- # add positional encoding to each token
- if interpolate_pos_encoding:
- embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
- else:
- embeddings = embeddings + self.position_embeddings
-
- embeddings = self.dropout(embeddings)
-
- return embeddings
-
-
-# Based on timm implementation, which can be found here:
-# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/image_transformer.py
-class PatchEmbeddings(nn.Module):
- """
- Image to Patch Embedding.
- """
-
- def __init__(
- self,
- image_size: int = 224,
- patch_size: Union[int, Tuple[int, int]] = 16,
- num_channels: int = 3,
- embed_dim: int = 768,
- ):
- super().__init__()
- if not isinstance(image_size, collections.abc.Iterable):
- image_size = (image_size, image_size)
- if not isinstance(patch_size, collections.abc.Iterable):
- patch_size = (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_patches = num_patches
-
- self.projection = nn.Conv2d(num_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
-
- def forward(self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
- batch_size, num_channels, height, width = pixel_values.shape
- if not interpolate_pos_encoding:
- if height != self.image_size[0] or width != self.image_size[1]:
- raise ValueError(
- f"Input image size ({height}*{width}) doesn't match model"
- f" ({self.image_size[0]}*{self.image_size[1]})."
- )
- x = self.projection(pixel_values).flatten(2).transpose(1, 2)
- return x
-
-
-class FlavaTextEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
-
- # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
- # any TensorFlow checkpoint file
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
- self.register_buffer(
- "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
- )
-
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- ):
- input_shape = input_ids.size()
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, :seq_length]
-
- # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
- # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
- # issue #5664
- if token_type_ids is None:
- if hasattr(self, "token_type_ids"):
- buffered_token_type_ids = self.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- inputs_embeds = self.word_embeddings(input_ids)
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
- embeddings = inputs_embeds + token_type_embeddings
- if self.position_embedding_type == "absolute":
- position_embeddings = self.position_embeddings(position_ids)
- embeddings += position_embeddings
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-class FlavaSelfAttention(nn.Module):
- def __init__(self, config: FlavaPossibleConfigs) -> None:
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
- f"heads {config.num_attention_heads}."
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
- self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
- self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(*new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- mixed_query_layer = self.query(hidden_states)
-
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(*new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-class FlavaSelfOutput(nn.Module):
- """
- The residual connection is defined in FlavaLayer (same as ViTLayer) instead of here (as is the case with other
- models), due to the layernorm applied before each block.
- """
-
- def __init__(self, config: FlavaPossibleConfigs) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- return hidden_states
-
-
-class FlavaAttention(nn.Module):
- def __init__(self, config: FlavaPossibleConfigs) -> None:
- super().__init__()
- self.attention = FlavaSelfAttention(config)
- self.output = FlavaSelfOutput(config)
- self.pruned_heads = set()
-
- def prune_heads(self, heads: Set[int]) -> None:
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.attention.query = prune_linear_layer(self.attention.query, index)
- self.attention.key = prune_linear_layer(self.attention.key, index)
- self.attention.value = prune_linear_layer(self.attention.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
- self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- self_outputs = self.attention(
- hidden_states, attention_mask=attention_mask, head_mask=head_mask, output_attentions=output_attentions
- )
-
- attention_output = self.output(self_outputs[0], hidden_states)
-
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-class FlavaIntermediate(nn.Module):
- def __init__(self, config: FlavaPossibleConfigs) -> None:
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- # Copied from transformers.models.vit.modeling_vit.ViTIntermediate.forward
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
-
- return hidden_states
-
-
-class FlavaOutput(nn.Module):
- def __init__(self, config: FlavaPossibleConfigs) -> None:
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- # Copied from transformers.models.vit.modeling_vit.ViTOutput.forward
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
-
- hidden_states = hidden_states + input_tensor
-
- return hidden_states
-
-
-class FlavaLayer(nn.Module):
- """This corresponds to the Block class in the timm implementation."""
-
- def __init__(self, config: FlavaPossibleConfigs) -> None:
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = FlavaAttention(config)
- self.intermediate = FlavaIntermediate(config)
- self.output = FlavaOutput(config)
-
- # TODO: Check fp32 layer norm possiblity
- self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
- self_attention_outputs = self.attention(
- self.layernorm_before(hidden_states), # in ViT, layernorm is applied before self-attention
- attention_mask=attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- )
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- # first residual connection
- hidden_states = attention_output + hidden_states
-
- # in ViT, layernorm is also applied after self-attention
- layer_output = self.layernorm_after(hidden_states)
- layer_output = self.intermediate(layer_output)
-
- # second residual connection is done here
- layer_output = self.output(layer_output, hidden_states)
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-class FlavaEncoder(nn.Module):
- def __init__(self, config: FlavaConfig) -> None:
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([FlavaLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[tuple, BaseModelOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- attention_mask,
- layer_head_mask,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(hidden_states, attention_mask, layer_head_mask, output_attentions)
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_self_attentions
- )
-
-
-class FlavaPooler(nn.Module):
- def __init__(self, config: FlavaPossibleConfigs):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = nn.Tanh()
-
- def forward(self, hidden_states: torch.Tensor):
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(first_token_tensor)
- pooled_output = self.activation(pooled_output)
- return pooled_output
-
-
-FLAVA_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
- as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`{config}`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FLAVA_INPUTS_DOCSTRING_COMMON = r"""
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
- [What are attention masks?](../glossary#attention-mask)
-
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
-
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-FLAVA_IMAGE_INPUTS_DOCSTRING_BASE = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`FlavaImageProcessor.__call__`] for details.
-
- bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, image_num_patches)`):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
-
- interpolate_pos_encoding (`bool`, *optional*):
- Whether to interpolate the pre-trained position encodings.
-"""
-
-FLAVA_IMAGE_INPUTS_DOCSTRING = FLAVA_IMAGE_INPUTS_DOCSTRING_BASE + FLAVA_INPUTS_DOCSTRING_COMMON
-
-FLAVA_TEXT_INPUTS_DOCSTRING_BASE = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See
- [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input
- IDs?](../glossary#input-ids)
-
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
- [What are token type IDs?](../glossary#token-type-ids)
-"""
-
-FLAVA_TEXT_INPUTS_DOCSTRING = FLAVA_TEXT_INPUTS_DOCSTRING_BASE + FLAVA_INPUTS_DOCSTRING_COMMON
-
-FLAVA_MULTIMODAL_INPUTS_DOCSTRING = (
- r"""
- Args:
- hidden_states (`torch.FloatTensor` of shape `(batch_size, image_num_patches + text_seq_len, hidden_size)`):
- The concatenated hidden states of unimodal encoders.
-"""
- + FLAVA_INPUTS_DOCSTRING_COMMON
-)
-
-FLAVA_MODEL_INPUTS_DOCSTRING_BASE = r"""
- Args:
- skip_multimodal_encoder (*bool*, *optional*):
- Skip any calculations for multimodal encoder. Useful if multimodal encoding is not going to be used.
-"""
-
-FLAVA_MODEL_INPUTS_DOCSTRING = (
- FLAVA_IMAGE_INPUTS_DOCSTRING_BASE
- + FLAVA_TEXT_INPUTS_DOCSTRING_BASE
- + FLAVA_INPUTS_DOCSTRING_COMMON
- + FLAVA_MODEL_INPUTS_DOCSTRING_BASE
-)
-
-
-FLAVA_PRETRAINING_INPUTS_DOCSTRING = (
- r"""
- Args:
- input_ids_masked (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary. These ones are the masked version of the original task
- to be used with MLM. Indices can be obtained using [`AutoTokenizer`] along with
- [`DataCollatorForMaskedLanguageModeling`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids)
-
-"""
- + FLAVA_TEXT_INPUTS_DOCSTRING_BASE
- + FLAVA_IMAGE_INPUTS_DOCSTRING_BASE
- + r"""
- image_attention_mask (`torch.FloatTensor` of shape `({1})`, *optional*):
- Mask to avoid performing attention on padding token indices specifically for images. Mask values selected
- in `[0, 1]`:
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
- [What are attention masks?](../glossary#attention-mask)
-
- skip_unmasked_multimodal_encoder (*bool*, *optional*):
- Skip any calculations for multimodal encoder for unmasked inputs. FLAVA pretraining doesn't need unmasked
- multimodal embeddings or outputs as of now.
-
- mlm_labels (`torch.LongTensor` of shape `(batch_size, text_seq_len)`, *optional*):
- Labels for computing the left-to-right language and multimodal masked modeling loss (next word prediction).
- Indices should be in `[-100, 0, ..., text_config.vocab_size - 1]` (see `input_ids` docstring). Tokens with
- indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0,
- ..., text_config.vocab_size - 1]`.
-
- mim_labels (`torch.LongTensor` of shape `(batch_size, image_num_patches)`, *optional*):
- Labels for computing the image and multimodal masked modeling loss. Indices should be in `[-100, 0, ...,
- image_config.vocab_size - 1]`. Tokens with indices set to `-100` are ignored (masked), the loss is only
- computed for the tokens with labels in `[0, ..., image_config.vocab_size - 1]`. If not passed, they are
- generated automatically using the image codebook assigned to the model. By default, it uses
- [`FlavaImageCodebook`]. See [`FlavaImageCodebook`] to understand how to generate mim_labels.
-
- itm_labels (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*):
- Labels for computing the image-text matching loss. 0 means the pairs don't match and 1 means they match.
- The pairs with 0 will be skipped for calculation of MMM and global contrastive losses as well.
-
- return_loss (`bool`, *optional*, default to None):
- Whether to return calculated loss or not.
-"""
- + FLAVA_INPUTS_DOCSTRING_COMMON
-)
-
-FLAVA_PRETRAINING_START_DOCSTRING_EXTRA = r"""
- Parameters:
- image_codebook ([`nn.Module`]): If passed, the image codebook will be set to this. Otherwise. it will
- be initialized using the image_codebook_config defined in the config first as the first parameter.
-"""
-
-
-class FlavaPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = FlavaConfig
- base_model_prefix = "flava"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-@add_start_docstrings(
- "The bare FLAVA Image Model transformer outputting raw hidden-states without any specific head on top.",
- FLAVA_START_DOCSTRING.format(config="FlavaImageConfig"),
-)
-class FlavaImageModel(FlavaPreTrainedModel):
- config_class = FlavaImageConfig
- # This override allows us to load FlavaImageModel from FlavaModel/FlavaForPreTraining checkpoints.
- base_model_prefix = "flava.image_model"
- main_input_name = "pixel_values"
-
- def __init__(self, config: FlavaImageConfig, add_pooling_layer: bool = True):
- super().__init__(config)
-
- self.config = config
-
- self.embeddings = FlavaImageEmbeddings(config)
- self.encoder = FlavaEncoder(config)
-
- self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.pooler = FlavaPooler(config) if add_pooling_layer else None
-
- self.post_init()
-
- def get_input_embeddings(self) -> nn.Module:
- return self.embeddings.patch_embeddings
-
- def set_input_embeddings(self, value: nn.Module):
- self.embeddings.patch_embeddings = value
-
- def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(FLAVA_IMAGE_INPUTS_DOCSTRING.format("batch_size, image_num_patches"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPooling,
- config_class=_CONFIG_CLASS_FOR_IMAGE_MODEL_DOC,
- modality="vision",
- expected_output=_EXPECTED_IMAGE_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- interpolate_pos_encoding: Optional[bool] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, BaseModelOutputWithPooling]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- embedding_output = self.embeddings(
- pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding
- )
-
- encoder_outputs = self.encoder(
- embedding_output,
- attention_mask=attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- return (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- "The bare FLAVA Text Model transformer outputting raw hidden-states without any specific head on top.",
- FLAVA_START_DOCSTRING.format(config="FlavaTextConfig"),
-)
-class FlavaTextModel(FlavaPreTrainedModel):
- config_class = FlavaTextConfig
- # This override allows us to load FlavaTextModel from FlavaModel/FlavaForPreTraining checkpoints.
- base_model_prefix = "flava.text_model"
-
- def __init__(self, config: FlavaTextConfig, add_pooling_layer: bool = True):
- super().__init__(config)
- self.config = config
-
- self.embeddings = FlavaTextEmbeddings(config)
- self.encoder = FlavaEncoder(config)
-
- self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.pooler = FlavaPooler(config) if add_pooling_layer else None
-
- self.post_init()
-
- def get_input_embeddings(self) -> PatchEmbeddings:
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value: nn.Module):
- self.embeddings.word_embeddings = value
-
- def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(FLAVA_TEXT_INPUTS_DOCSTRING.format("batch_size, text_seq_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPooling,
- config_class=_CONFIG_CLASS_FOR_TEXT_MODEL_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, BaseModelOutputWithPooling]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is None:
- raise ValueError("You have to specify input_ids")
-
- input_shape = input_ids.size()
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=input_ids.device)
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
- extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
- attention_mask, input_shape, input_ids.device
- )
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- )
-
- encoder_outputs = self.encoder(
- embedding_output,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- return (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- "The bare FLAVA Multimodal Model transformer outputting raw hidden-states without any specific head on top.",
- FLAVA_START_DOCSTRING.format(config="FlavaMultimodalConfig"),
-)
-class FlavaMultimodalModel(FlavaPreTrainedModel):
- config_class = FlavaMultimodalConfig
- # This override allows us to load FlavaMultimodalModel from FlavaModel/FlavaForPreTraining checkpoints.
- base_model_prefix = "flava.multimodal_model"
- main_input_name = "hidden_states"
-
- def __init__(self, config: FlavaMultimodalConfig, add_pooling_layer=True):
- super().__init__(config)
- self.config = config
- self.use_cls_token = self.config.use_cls_token
- if self.use_cls_token:
- self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
-
- self.encoder = FlavaEncoder(config)
-
- self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.pooler = FlavaPooler(config) if add_pooling_layer else None
-
- self.post_init()
-
- def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(
- FLAVA_MULTIMODAL_INPUTS_DOCSTRING.format("batch_size, image_num_patches + text_seq_len")
- )
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPooling,
- config_class=_CONFIG_CLASS_FOR_MULTIMODAL_MODEL_DOC,
- )
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, BaseModelOutputWithPooling]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- batch_size, seq_length, _ = hidden_states.size()
-
- if self.use_cls_token:
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- hidden_states = torch.cat((cls_tokens, hidden_states), dim=1)
- seq_length += 1
-
- if attention_mask is None:
- attention_mask = torch.ones((batch_size, seq_length), device=hidden_states.device)
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
- extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
- attention_mask, (batch_size, seq_length), hidden_states.device
- )
-
- encoder_outputs = self.encoder(
- hidden_states,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
- pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- return (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- "The bare FLAVA Model transformer outputting raw hidden-states without any specific head on top.",
- FLAVA_START_DOCSTRING.format(config="FlavaConfig"),
-)
-class FlavaModel(FlavaPreTrainedModel):
- config_class = FlavaConfig
-
- def __init__(self, config: FlavaConfig):
- super().__init__(config)
-
- if not isinstance(config.text_config, FlavaTextConfig):
- raise ValueError(
- "config.text_config is expected to be of type FlavaTextConfig but is of type"
- f" {type(config.text_config)}."
- )
-
- if not isinstance(config.image_config, FlavaImageConfig):
- raise ValueError(
- "config.image_config is expected to be of type FlavaImageConfig but is of type"
- f" {type(config.image_config)}."
- )
-
- if not isinstance(config.multimodal_config, FlavaMultimodalConfig):
- raise ValueError(
- "config.multimodal_config is expected to be of type FlavaMultimodalConfig but "
- + f"is of type {type(config.multimodal_config)}."
- )
-
- text_config = config.text_config
- image_config = config.image_config
- multimodal_config = config.multimodal_config
-
- self.projection_dim = config.projection_dim
- self.text_hidden_size = text_config.hidden_size
- self.image_hidden_size = image_config.hidden_size
- self.mm_hidden_size = multimodal_config.hidden_size
-
- self.text_model = FlavaTextModel(text_config)
- self.image_model = FlavaImageModel(image_config)
- self.multimodal_model = FlavaMultimodalModel(multimodal_config)
-
- self.image_projection = nn.Linear(self.image_hidden_size, self.projection_dim)
- self.text_projection = nn.Linear(self.text_hidden_size, self.projection_dim)
- self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
-
- self.image_to_mm_projection = nn.Linear(self.image_hidden_size, self.mm_hidden_size)
- self.text_to_mm_projection = nn.Linear(self.text_hidden_size, self.mm_hidden_size)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FLAVA_TEXT_INPUTS_DOCSTRING.format("batch_size, text_seq_length"))
- def get_text_features(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> torch.FloatTensor:
- r"""
- Returns:
- text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
- applying the projection layer to the pooled output of [`FlavaTextModel`].
-
- Examples:
-
- ```python
- >>> from transformers import AutoProcessor, FlavaModel
-
- >>> model = FlavaModel.from_pretrained("{0}")
- >>> processor = AutoProcessor.from_pretrained("{0}")
-
- >>> inputs = processor(
- ... text=["a photo of a cat", "a photo of a dog"], max_length=77, padding="max_length", return_tensors="pt"
- ... )
- >>> text_features = model.get_text_features(**inputs)
- ```""".format(_CHECKPOINT_FOR_DOC)
- text_outputs = self.text_model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = text_outputs[0] # last_hidden_state
- text_features = self.text_projection(pooled_output)
-
- return text_features
-
- @add_start_docstrings_to_model_forward(FLAVA_IMAGE_INPUTS_DOCSTRING.format("batch_size, image_num_patches"))
- def get_image_features(
- self,
- pixel_values: Optional[torch.Tensor] = None,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- interpolate_pos_encoding: Optional[bool] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> torch.FloatTensor:
- r"""
- Returns:
- image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by
- applying the projection layer to the pooled output of [`FlavaImageModel`].
-
- Examples:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoProcessor, FlavaModel
-
- >>> model = FlavaModel.from_pretrained("{0}")
- >>> processor = AutoProcessor.from_pretrained("{0}")
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> inputs = processor(images=image, return_tensors="pt")
-
- >>> image_features = model.get_image_features(**inputs)
- ```""".format(_CHECKPOINT_FOR_DOC)
- image_outputs = self.image_model(
- pixel_values=pixel_values,
- bool_masked_pos=bool_masked_pos,
- attention_mask=attention_mask,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- interpolate_pos_encoding=interpolate_pos_encoding,
- return_dict=return_dict,
- )
-
- pooled_output = image_outputs[0] # last_hidden_state
- image_features = self.image_projection(pooled_output)
-
- return image_features
-
- @add_start_docstrings_to_model_forward(
- FLAVA_MODEL_INPUTS_DOCSTRING.format("batch_size, image_num_patches + text_seq_len")
- )
- @replace_return_docstrings(output_type=FlavaModelOutput, config_class=FlavaConfig)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- pixel_values: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- bool_masked_pos: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- image_attention_mask: Optional[torch.Tensor] = None,
- skip_multimodal_encoder: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: bool = True,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, FlavaOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoProcessor, FlavaModel
-
- >>> model = FlavaModel.from_pretrained("facebook/flava-full")
- >>> processor = AutoProcessor.from_pretrained("facebook/flava-full")
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> inputs = processor(text=["a photo of a cat"], images=image, return_tensors="pt", padding=True)
-
- >>> outputs = model(**inputs)
-
- >>> image_embeddings = outputs.image_embeddings
- >>> text_embeddings = outputs.text_embeddings
- >>> multimodal_embeddings = outputs.multimodal_embeddings
-
- >>> outputs.image_embeddings.shape
- torch.Size([1, 197, 768])
-
- >>> text_embeddings.shape
- torch.Size([1, 7, 768])
-
- >>> multimodal_embeddings.shape
- torch.Size([1, 205, 768])
- ```
- """
-
- return_dict = return_dict if return_dict is not None else self.config.return_dict
- if not output_hidden_states:
- raise ValueError("FLAVA model requires hidden states to work. Please set `output_hidden_states=True`")
- image_embeddings = None
- image_states = None
- image_mm_projection = None
- image_output = None
- if pixel_values is not None:
- image_output = self.image_model(
- pixel_values=pixel_values,
- bool_masked_pos=bool_masked_pos,
- attention_mask=image_attention_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- image_embeddings, image_states = image_output[0], image_output[2]
- # Note that these states don't use final layernorm in the transformer model
- image_mm_projection = self.image_to_mm_projection(image_states[-1])
-
- text_embeddings = None
- text_states = None
- text_mm_projection = None
- text_output = None
- if input_ids is not None:
- text_output = self.text_model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- token_type_ids=token_type_ids,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- text_embeddings, text_states = text_output[0], text_output[2]
- # Note that these states don't use final layernorm in the transformer model
- text_mm_projection = self.text_to_mm_projection(text_states[-1])
-
- multimodal_embeddings = None
- multimodal_output = None
- if image_mm_projection is not None and text_mm_projection is not None and not skip_multimodal_encoder:
- if attention_mask is not None:
- batch_size, seq_len, _ = image_mm_projection.shape
- if self.multimodal_model.use_cls_token:
- seq_len += 1
- attention_mask_image = torch.ones(batch_size, seq_len, device=image_mm_projection.device)
- attention_multimodal = torch.cat([attention_mask_image, attention_mask], dim=1)
- else:
- attention_multimodal = None
- multimodal_input = torch.cat([image_mm_projection, text_mm_projection], dim=1)
- multimodal_output = self.multimodal_model(
- multimodal_input, attention_mask=attention_multimodal, return_dict=return_dict
- )
- multimodal_embeddings = multimodal_output[0]
-
- if not return_dict:
- return (
- image_embeddings,
- image_output,
- text_embeddings,
- text_output,
- multimodal_embeddings,
- multimodal_output,
- )
-
- return FlavaModelOutput(
- image_embeddings=image_embeddings,
- image_output=image_output,
- text_embeddings=text_embeddings,
- text_output=text_output,
- multimodal_embeddings=multimodal_embeddings,
- multimodal_output=multimodal_output,
- )
-
-
-class FlavaImageCodebookResPath(nn.Module):
- def __init__(self, in_size: int, out_size: int, **kwargs):
- super().__init__()
- hid_size = out_size // 4
-
- path = OrderedDict()
- path["relu_1"] = nn.ReLU()
- path["conv_1"] = nn.Conv2d(in_size, hid_size, kernel_size=3, padding=1)
- path["relu_2"] = nn.ReLU()
- path["conv_2"] = nn.Conv2d(hid_size, hid_size, kernel_size=3, padding=1)
- path["relu_3"] = nn.ReLU()
- path["conv_3"] = nn.Conv2d(hid_size, hid_size, kernel_size=3, padding=1)
- path["relu_4"] = nn.ReLU()
- path["conv_4"] = nn.Conv2d(hid_size, out_size, kernel_size=1, padding=0)
-
- self.path = nn.Sequential(path)
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- return self.path(x)
-
-
-class FlavaImageCodebookBlock(nn.Module):
- def __init__(self, in_size: int, out_size: int, num_layers: int, **kwargs):
- super().__init__()
-
- self.post_gain = 1 / (num_layers**2)
-
- if in_size != out_size:
- self.id_path = nn.Conv2d(in_size, out_size, kernel_size=1, padding=0)
- else:
- self.id_path = nn.Identity()
-
- self.res_path = FlavaImageCodebookResPath(in_size, out_size)
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- return self.id_path(x) + self.post_gain * self.res_path(x)
-
-
-class FlavaImageCodebookLayerGroup(nn.Module):
- def __init__(self, num_blocks: int, num_layers: int, in_size: int, out_size: int, use_pool: bool = True):
- super().__init__()
- blocks = OrderedDict()
- for i in range(num_blocks):
- if i == 0:
- blocks[f"block_{i+1}"] = FlavaImageCodebookBlock(in_size, out_size, num_layers)
- else:
- blocks[f"block_{i+1}"] = FlavaImageCodebookBlock(out_size, out_size, num_layers)
-
- if use_pool:
- blocks["pool"] = nn.MaxPool2d(kernel_size=2)
-
- self.group = nn.Sequential(blocks)
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- return self.group(x)
-
-
-# Inspired by DALLE Encoder in https://github.com/openai/DALL-E/blob/5be4b236bc3ade6943662354117a0e83752cc322/dall_e/encoder.py#L42
-@add_start_docstrings(
- """
- The FLAVA's image codebook model inspired from DALL-E's original encoder. Outputs raw hidden states and can be used
- to generate image tokens for an image based on DALL-E's vocab. Used to generate labels for MIM. Use
- `get_codebook_indices` to get image tokens for an image.
- """,
- FLAVA_START_DOCSTRING.format(config="FlavaImageCodebookConfig"),
-)
-class FlavaImageCodebook(FlavaPreTrainedModel):
- base_model_prefix = ""
- config_class = FlavaImageCodebookConfig
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = False
-
- def __init__(
- self,
- config: FlavaImageCodebookConfig,
- **kwargs: Any,
- ):
- super().__init__(config)
-
- self.config = config
- self.num_groups = config.num_groups
- self.input_channels = config.input_channels
- self.num_blocks_per_group = config.num_blocks_per_group
- self.hidden_size = config.hidden_size
- self.vocab_size = config.vocab_size
-
- num_layers = self.num_groups * self.num_blocks_per_group
-
- output_blocks = OrderedDict()
- output_blocks["relu"] = nn.ReLU()
- output_blocks["conv"] = nn.Conv2d(8 * self.hidden_size, self.vocab_size, kernel_size=1, padding=0)
-
- blocks = OrderedDict()
- blocks["input"] = nn.Conv2d(self.input_channels, 1 * self.hidden_size, kernel_size=7, padding=3)
- blocks["group_1"] = FlavaImageCodebookLayerGroup(
- self.num_blocks_per_group, num_layers, 1 * self.hidden_size, 1 * self.hidden_size
- )
- blocks["group_2"] = FlavaImageCodebookLayerGroup(
- self.num_blocks_per_group, num_layers, 1 * self.hidden_size, 2 * self.hidden_size
- )
- blocks["group_3"] = FlavaImageCodebookLayerGroup(
- self.num_blocks_per_group, num_layers, 2 * self.hidden_size, 4 * self.hidden_size
- )
- blocks["group_4"] = FlavaImageCodebookLayerGroup(
- self.num_blocks_per_group, num_layers, 4 * self.hidden_size, 8 * self.hidden_size, use_pool=False
- )
- blocks["output"] = nn.Sequential(output_blocks)
-
- self.blocks = nn.Sequential(blocks)
-
- self.post_init()
-
- if self.config.freeze:
- for param in self.parameters():
- param.requires_grad = False
-
- def get_codebook_indices(self, pixel_values: torch.Tensor) -> torch.Tensor:
- """
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Codebook pixel values can be obtained using [`AutoImageProcessor`] by passing
- `return_codebook_pixels=True`. See [`FlavaImageProcessor.__call__`] for details.
-
- Examples:
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoImageProcessor, FlavaImageCodebook
-
- >>> model = FlavaImageCodebook.from_pretrained("{0}")
- >>> image_processor = AutoImageProcessor.from_pretrained("{0}")
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> inputs = image_processor([image], return_codebook_pixels=True, return_tensors="pt")
- >>> inputs = dict(pixel_values=inputs.codebook_pixel_values)
-
- >>> outputs = model.get_codebook_indices(**inputs)
- ```
- """.format(_CHECKPOINT_FOR_CODEBOOK_DOC)
- z_logits = self.blocks(pixel_values)
- return torch.argmax(z_logits, axis=1)
-
- def get_codebook_probs(self, pixel_values: torch.Tensor) -> torch.Tensor:
- z_logits = self.blocks(pixel_values)
- return nn.Softmax(dim=1)(z_logits)
-
- def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
- """
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Codebook pixel values can be obtained using [`AutoImageProcessor`] by passing
- `return_codebook_pixels=True`. See [`FlavaImageProcessor.__call__`] for details.
-
- Examples:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoImageProcessor, FlavaImageCodebook
-
- >>> model = FlavaImageCodebook.from_pretrained("{0}")
- >>> image_processor = AutoImageProcessor.from_pretrained("{0}")
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> inputs = image_processor([image], return_codebook_pixels=True, return_tensors="pt")
- >>> inputs = dict(pixel_values=inputs.codebook_pixel_values)
-
- >>> outputs = model(**inputs)
- >>> print(outputs.shape)
- (1, 196)
- ```
- """.format(_CHECKPOINT_FOR_CODEBOOK_DOC)
- if len(pixel_values.shape) != 4:
- raise ValueError(f"input shape {pixel_values.shape} is not 4d")
- if pixel_values.shape[1] != self.input_channels:
- raise ValueError(f"input has {pixel_values.shape[1]} channels but model built for {self.input_channels}")
- return self.blocks(pixel_values)
-
-
-class FlavaPredictionHeadTransform(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = ACT2FN[config.hidden_act]
- else:
- self.transform_act_fn = config.hidden_act
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_states):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
- return hidden_states
-
-
-class FlavaMaskedPredictionHead(nn.Module):
- def __init__(self, config, weight=None):
- super().__init__()
- self.config = config
- self.transform = FlavaPredictionHeadTransform(config)
- self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
- self.bias = nn.Parameter(torch.zeros(config.vocab_size))
- if weight is not None:
- self.decoder.weight = weight
-
- # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
- self.decoder.bias = self.bias
-
- def forward(self, x):
- x = self.transform(x)
- x = self.decoder(x)
- return x
-
-
-class FlavaITMHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.pooler = FlavaPooler(config)
- self.seq_relationship = nn.Linear(config.hidden_size, 2)
-
- def forward(self, x):
- x = self.pooler(x)
- x = self.seq_relationship(x)
- return x
-
-
-class FlavaGlobalContrastiveHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.global_backprop_contrastive = config.global_backprop_contrastive
-
- def forward(self, image_embeddings, text_embeddings, logit_scale):
- temperature = torch.exp(logit_scale)
- if not torch.distributed.is_available() or not torch.distributed.is_initialized():
- labels = torch.arange(image_embeddings.size(0), device=image_embeddings.device)
- image_embeddings_all = [image_embeddings]
- text_embeddings_all = [text_embeddings]
- else:
- local_batch_size = image_embeddings.size(0)
- world_size = torch.distributed.get_world_size()
-
- if self.global_backprop_contrastive:
- # `torch.distributed.nn.functional.all_gather` does backprop on all active workers
- # whereas `torch.distributed.all_gather` does only backpropagates on the current worker.
- image_embeddings_all = torch.distributed.nn.functional.all_gather(image_embeddings)
- text_embeddings_all = torch.distributed.nn.functional.all_gather(text_embeddings)
- else:
- image_embeddings_all = [torch.zeros_like(text_embeddings) for _ in range(world_size)]
- text_embeddings_all = [torch.zeros_like(image_embeddings) for _ in range(world_size)]
- torch.distributed.all_gather(image_embeddings_all, image_embeddings)
- torch.distributed.all_gather(text_embeddings_all, text_embeddings)
-
- labels = local_batch_size * torch.distributed.get_rank() + torch.arange(
- local_batch_size, device=image_embeddings.device
- )
-
- image_embeddings_all = torch.cat(image_embeddings_all)
- text_embeddings_all = torch.cat(text_embeddings_all)
-
- logits_per_image = torch.matmul(image_embeddings, text_embeddings_all.transpose(0, 1)) * temperature
- logits_per_text = torch.matmul(text_embeddings, image_embeddings_all.transpose(0, 1)) * temperature
-
- return logits_per_image, logits_per_text, labels
-
-
-@add_start_docstrings(
- """
- The FLAVA model for pretraining which outputs losses, embeddings, logits and transformer outputs.
- """,
- FLAVA_START_DOCSTRING.format(config="FlavaConfig") + FLAVA_PRETRAINING_START_DOCSTRING_EXTRA,
-)
-class FlavaForPreTraining(FlavaPreTrainedModel):
- # Those are linked to xxx.bias
- _tied_weights_keys = [
- "mmm_text_head.decoder.bias",
- "mmm_image_head.decoder.bias",
- "mlm_head.decoder.bias",
- "mim_head.decoder.bias",
- ]
-
- def __init__(self, config: FlavaConfig, image_codebook: Optional[nn.Module] = None):
- super().__init__(config)
- self.flava = FlavaModel(config)
-
- self.image_codebook = image_codebook
- if self.image_codebook is None and config.init_codebook:
- self.image_codebook = FlavaImageCodebook(config.image_codebook_config)
-
- # Levarage text and image encoder configs to create the masked
- # head since it has the right vocab
- self.mim_head = FlavaMaskedPredictionHead(config.image_config)
- self.mlm_head = FlavaMaskedPredictionHead(config.text_config)
- self.itm_head = FlavaITMHead(config)
- self.mmm_image_head = FlavaMaskedPredictionHead(config.image_config)
- self.mmm_text_head = FlavaMaskedPredictionHead(config.text_config)
- self.global_contrastive_head = FlavaGlobalContrastiveHead(config)
-
- self.image_vocab_size = config.image_config.vocab_size
- self.text_vocab_size = config.text_config.vocab_size
- self.mlm_weight = config.mlm_weight
- self.mim_weight = config.mim_weight
- self.global_contrastive_weight = config.global_contrastive_weight
- self.ce_ignore_index = config.ce_ignore_index
- self.itm_weight = config.itm_weight
- self.mmm_image_weight = config.mmm_image_weight
- self.mmm_text_weight = config.mmm_text_weight
- self.skip_unmasked_multimodal_encoder = config.skip_unmasked_multimodal_encoder
-
- self.post_init()
-
- def _resize_to_2d(self, x: torch.Tensor):
- if x.dim() > 2:
- x = x.view(x.size(0), -1)
- return x
-
- @add_start_docstrings_to_model_forward(
- FLAVA_PRETRAINING_INPUTS_DOCSTRING.format("batch_size, text_seq_len", "batch_size, image_num_patches")
- )
- @replace_return_docstrings(output_type=FlavaForPreTrainingOutput, config_class=FlavaConfig)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- input_ids_masked: Optional[torch.LongTensor] = None,
- pixel_values: Optional[torch.FloatTensor] = None,
- codebook_pixel_values: Optional[torch.FloatTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- bool_masked_pos: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- image_attention_mask: Optional[torch.Tensor] = None,
- skip_unmasked_multimodal_encoder: bool = None,
- mlm_labels: Optional[torch.Tensor] = None,
- mim_labels: Optional[torch.Tensor] = None,
- itm_labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: bool = True,
- return_dict: Optional[bool] = None,
- return_loss: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], FlavaForPreTrainingOutput]:
- """
- Examples:
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import FlavaForPreTraining, AutoProcessor
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> model = FlavaForPreTraining.from_pretrained("facebook/flava-full")
- >>> processor = AutoProcessor.from_pretrained("facebook/flava-full")
-
- >>> text = ["a photo of a cat"]
-
- >>> inputs = processor(
- ... images=[image],
- ... text=text,
- ... return_masks=True,
- ... return_codebook_pixels=True,
- ... padding=True,
- ... max_length=77,
- ... return_tensors="pt",
- ... )
-
-
- >>> output = model(**inputs)
- ```
-
- Return:
-
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- return_loss = return_loss if return_loss is not None else self.config.return_loss
-
- skip_unmasked_multimodal_encoder = (
- skip_unmasked_multimodal_encoder
- if skip_unmasked_multimodal_encoder is not None
- else self.skip_unmasked_multimodal_encoder
- )
-
- if input_ids_masked is None and input_ids is not None:
- logger.warning(
- "`input_ids_masked` isn't passed which means MLM loss won't be calculated correctlySetting it to"
- " `input_ids` so that model can work. Please pass it if this is unintentional. This is usually OKAY if"
- " you are doing inference on unmasked text..."
- )
- input_ids_masked = input_ids
-
- flava_output = self.flava(
- input_ids=input_ids,
- pixel_values=pixel_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- image_attention_mask=image_attention_mask,
- # Don't need unmasked multimodal embedding for anything so skip it
- # NOTE: ITM uses masked version
- skip_multimodal_encoder=skip_unmasked_multimodal_encoder,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- # Pass true to have deterministic outputs
- return_dict=True,
- )
-
- flava_masked_output = self.flava(
- input_ids=input_ids_masked,
- pixel_values=pixel_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- image_attention_mask=image_attention_mask,
- bool_masked_pos=bool_masked_pos,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=True,
- )
-
- pos_mask = None
-
- image_embeddings = flava_output.image_embeddings
- text_embeddings = flava_output.text_embeddings
- image_masked_embeddings = flava_masked_output.image_embeddings
- text_masked_embeddings = flava_masked_output.text_embeddings
- multimodal_masked_embeddings = flava_masked_output.multimodal_embeddings
-
- total_loss = mim_loss = mlm_loss = mmm_text_loss = mmm_image_loss = gc_loss = itm_loss = None
- mim_logits = mlm_logits = mmm_text_logits = mmm_image_logits = None
- itm_logits = logits_per_image = logits_per_text = None
-
- # Calculate mim_labels if necessary from the image_codebook
- if image_masked_embeddings is not None or multimodal_masked_embeddings is not None:
- if mim_labels is None and return_loss:
- if self.image_codebook is None:
- raise RuntimeError(
- "`return_loss` is set to True but the image codebook is not initialized and no `mim_labels` "
- " have been passed. Reinstantiate the model with `init_codebook` set to True or "
- "pass in your custom `mim_labels`"
- )
- if codebook_pixel_values is None:
- raise ValueError(
- "`codebook_pixel_value` are required to generate `mim_labels` if loss is expected. "
- "Call `AutoProcessor` with `return_codebook_pixels` set to True"
- )
- mim_labels = self.image_codebook.get_codebook_indices(codebook_pixel_values)
- # Unimodal MIM Loss
- # If multimodal embeddings are present, we will calculate MMM loss
- if self.mim_weight > 0 and image_masked_embeddings is not None and multimodal_masked_embeddings is None:
- sequence_for_image = image_masked_embeddings
-
- if mim_labels is not None:
- mim_labels = self._resize_to_2d(mim_labels)
- bool_masked_pos = self._resize_to_2d(bool_masked_pos)
- mim_labels[bool_masked_pos.ne(True)] = self.ce_ignore_index
-
- sequence_for_image = sequence_for_image[:, -mim_labels.size(1) :, :]
- masked_tokens = mim_labels.ne(self.ce_ignore_index)
- mim_labels_filtered = mim_labels[masked_tokens]
- sequence_for_image = sequence_for_image[masked_tokens, :]
- mim_logits = self.mim_head(sequence_for_image)
- if return_loss:
- mim_loss = nn.functional.cross_entropy(
- mim_logits.view(-1, self.image_vocab_size), mim_labels_filtered.view(-1)
- )
- mim_loss *= self.mim_weight
- else:
- mim_logits = self.mim_head(sequence_for_image)
-
- # Unimodal MLM Loss
- if self.mlm_weight > 0 and text_masked_embeddings is not None and multimodal_masked_embeddings is None:
- sequence_for_text = text_masked_embeddings
- if mlm_labels is not None:
- mlm_labels = self._resize_to_2d(mlm_labels)
- sequence_for_text = sequence_for_text[:, -mlm_labels.size(1) :, :]
- masked_tokens = mlm_labels.ne(self.ce_ignore_index)
- mlm_labels_filtered = mlm_labels[masked_tokens]
- sequence_for_text = sequence_for_text[masked_tokens, :]
- mlm_logits = self.mlm_head(sequence_for_text)
- if return_loss:
- mlm_loss = nn.functional.cross_entropy(
- mlm_logits.view(-1, self.text_vocab_size), mlm_labels_filtered.view(-1)
- )
- mlm_loss *= self.mlm_weight
- else:
- mlm_logits = self.mlm_head(sequence_for_text)
-
- # ITM Loss
- if self.itm_weight > 0 and multimodal_masked_embeddings is not None:
- itm_logits = self.itm_head(multimodal_masked_embeddings)
-
- if itm_labels is not None:
- pos_pairs = itm_labels.ne(0)
- pos_mask = torch.where(pos_pairs.any(), pos_pairs, pos_pairs.new([True]))
- if return_loss:
- itm_loss = nn.functional.cross_entropy(itm_logits, itm_labels)
- itm_loss *= self.itm_weight
-
- if multimodal_masked_embeddings is not None:
- multimodal_masked_embeddings = multimodal_masked_embeddings[pos_mask]
-
- if mlm_labels is not None:
- mlm_labels = mlm_labels[pos_mask]
-
- if mim_labels is not None:
- mim_labels = mim_labels[pos_mask]
- bool_masked_pos = bool_masked_pos[pos_mask]
-
- # MMM Image Loss
- if multimodal_masked_embeddings is not None and self.mmm_image_weight > 0:
- sequence_for_image = multimodal_masked_embeddings
- end_index = image_masked_embeddings.size(1) - 1
- sequence_for_image = sequence_for_image[:, 2 : 2 + end_index, :]
-
- if mim_labels is not None:
- mim_labels = self._resize_to_2d(mim_labels)
- bool_masked_pos = self._resize_to_2d(bool_masked_pos)
- mim_labels[bool_masked_pos.ne(True)] = self.ce_ignore_index
-
- masked_tokens = mim_labels.ne(self.ce_ignore_index)
- mim_labels_filtered = mim_labels[masked_tokens]
- sequence_for_image = sequence_for_image[masked_tokens, :]
- mmm_image_logits = self.mmm_image_head(sequence_for_image)
- if return_loss:
- mmm_image_loss = nn.functional.cross_entropy(
- mmm_image_logits.view(-1, self.image_vocab_size), mim_labels_filtered.view(-1)
- )
- mmm_image_loss *= self.mmm_image_weight
- else:
- mmm_image_logits = self.mmm_image_head(sequence_for_image)
-
- # MMM Text Loss
- if multimodal_masked_embeddings is not None and self.mmm_text_weight > 0:
- sequence_for_text = multimodal_masked_embeddings
- sequence_for_text = sequence_for_text[:, -text_masked_embeddings.size(1) :, :]
-
- if mlm_labels is not None:
- mlm_labels = self._resize_to_2d(mlm_labels)
- masked_tokens = mlm_labels.ne(self.ce_ignore_index)
- mlm_labels_filtered = mlm_labels[masked_tokens]
- sequence_for_text = sequence_for_text[masked_tokens, :]
- mmm_text_logits = self.mmm_text_head(sequence_for_text)
- if return_loss:
- mmm_text_loss = nn.functional.cross_entropy(
- mmm_text_logits.view(-1, self.text_vocab_size), mlm_labels_filtered.view(-1)
- )
- mmm_text_loss *= self.mmm_text_weight
- else:
- mmm_text_logits = self.mmm_text_head(sequence_for_text)
-
- # Global Contrastive Loss
- if image_embeddings is not None and text_embeddings is not None and self.global_contrastive_weight > 0:
- text_embedding = self.flava.text_projection(text_embeddings[:, 0, :])
- text_embedding = nn.functional.normalize(text_embedding, dim=-1)
-
- image_embedding = self.flava.image_projection(image_embeddings[:, 0, :])
- image_embedding = nn.functional.normalize(image_embedding, dim=-1)
-
- self.flava.logit_scale.data.clamp_(LOGIT_SCALE_CLAMP_MIN, LOGIT_SCALE_CLAMP_MAX)
-
- logits_per_image, logits_per_text, gc_labels = self.global_contrastive_head(
- image_embedding, text_embedding, self.flava.logit_scale
- )
-
- # Apply ITM negative mask if any
- if pos_mask is not None:
- logits_per_image = logits_per_image[pos_mask]
- logits_per_text = logits_per_text[pos_mask]
- gc_labels = gc_labels[pos_mask]
-
- if return_loss:
- gc_loss_image = nn.functional.cross_entropy(logits_per_image, gc_labels)
- gc_loss_text = nn.functional.cross_entropy(logits_per_text, gc_labels)
- gc_loss = (gc_loss_image + gc_loss_text) / 2
- gc_loss *= self.global_contrastive_weight
-
- flava_losses = FlavaLosses(
- mim=mim_loss,
- mlm=mlm_loss,
- itm=itm_loss,
- global_contrastive=gc_loss,
- mmm_image=mmm_image_loss,
- mmm_text=mmm_text_loss,
- )
-
- if return_loss and not flava_losses.all_none():
- total_loss = sum(loss if loss is not None else 0 for loss in flava_losses.values())
-
- if not return_dict:
- output = (
- image_embeddings,
- flava_output.image_output.to_tuple() if flava_output.image_output is not None else None,
- text_embeddings,
- flava_output.text_output.to_tuple() if flava_output.text_output is not None else None,
- flava_output.multimodal_embeddings,
- flava_output.multimodal_output.to_tuple() if flava_output.multimodal_output is not None else None,
- image_masked_embeddings,
- flava_masked_output.image_output.to_tuple() if flava_masked_output.image_output is not None else None,
- text_masked_embeddings,
- flava_masked_output.text_output.to_tuple() if flava_masked_output.text_output is not None else None,
- multimodal_masked_embeddings,
- flava_masked_output.multimodal_output.to_tuple()
- if flava_masked_output.multimodal_output is not None
- else None,
- mim_logits,
- mlm_logits,
- itm_logits,
- logits_per_image,
- logits_per_image,
- mmm_image_logits,
- mmm_text_logits,
- )
- if return_loss and not flava_losses.all_none():
- output = (
- total_loss,
- flava_losses,
- ) + output
-
- # Filter None as transformer by default won't handle it
- return tuple(x for x in output if x is None)
-
- return FlavaForPreTrainingOutput(
- loss=total_loss,
- loss_info=flava_losses,
- image_embeddings=image_embeddings,
- image_output=flava_output.image_output,
- text_embeddings=text_embeddings,
- text_output=flava_output.text_output,
- multimodal_embeddings=flava_output.multimodal_embeddings,
- multimodal_output=flava_output.multimodal_output,
- image_masked_embeddings=image_masked_embeddings,
- image_masked_output=flava_masked_output.image_output,
- text_masked_embeddings=text_masked_embeddings,
- text_masked_output=flava_masked_output.text_output,
- multimodal_masked_embeddings=multimodal_masked_embeddings,
- multimodal_masked_output=flava_masked_output.multimodal_output,
- mim_logits=mim_logits,
- mlm_logits=mlm_logits,
- itm_logits=itm_logits,
- contrastive_logits_per_image=logits_per_image,
- contrastive_logits_per_text=logits_per_text,
- mmm_image_logits=mmm_image_logits,
- mmm_text_logits=mmm_text_logits,
- )
diff --git a/transformers/models/flava/processing_flava.py b/transformers/models/flava/processing_flava.py
deleted file mode 100644
index 7f439b040a8fd04e898075875cc96c7d26440959..0000000000000000000000000000000000000000
--- a/transformers/models/flava/processing_flava.py
+++ /dev/null
@@ -1,165 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Image/Text processor class for FLAVA
-"""
-
-import warnings
-from typing import List, Optional, Union
-
-from ...image_utils import ImageInput
-from ...processing_utils import ProcessorMixin
-from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
-from ...utils import TensorType
-
-
-class FlavaProcessor(ProcessorMixin):
- r"""
- Constructs a FLAVA processor which wraps a FLAVA image processor and a FLAVA tokenizer into a single processor.
-
- [`FlavaProcessor`] offers all the functionalities of [`FlavaImageProcessor`] and [`BertTokenizerFast`]. See the
- [`~FlavaProcessor.__call__`] and [`~FlavaProcessor.decode`] for more information.
-
- Args:
- image_processor ([`FlavaImageProcessor`], *optional*): The image processor is a required input.
- tokenizer ([`BertTokenizerFast`], *optional*): The tokenizer is a required input.
- """
-
- attributes = ["image_processor", "tokenizer"]
- image_processor_class = "FlavaImageProcessor"
- tokenizer_class = ("BertTokenizer", "BertTokenizerFast")
-
- def __init__(self, image_processor=None, tokenizer=None, **kwargs):
- feature_extractor = None
- if "feature_extractor" in kwargs:
- warnings.warn(
- "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
- " instead.",
- FutureWarning,
- )
- feature_extractor = kwargs.pop("feature_extractor")
-
- image_processor = image_processor if image_processor is not None else feature_extractor
- if image_processor is None:
- raise ValueError("You need to specify an `image_processor`.")
- if tokenizer is None:
- raise ValueError("You need to specify a `tokenizer`.")
-
- super().__init__(image_processor, tokenizer)
- self.current_processor = self.image_processor
-
- def __call__(
- self,
- images: Optional[ImageInput] = None,
- text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None,
- add_special_tokens: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str, TruncationStrategy] = False,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_image_mask: Optional[bool] = None,
- return_codebook_pixels: Optional[bool] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- return_tensors: Optional[Union[str, TensorType]] = None,
- **kwargs,
- ):
- """
- This method uses [`FlavaImageProcessor.__call__`] method to prepare image(s) for the model, and
- [`BertTokenizerFast.__call__`] to prepare text for the model.
-
- Please refer to the docstring of the above two methods for more information.
- """
-
- if text is None and images is None:
- raise ValueError("You have to specify either text or images. Both cannot be none.")
-
- if text is not None:
- encoding = self.tokenizer(
- text=text,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=pad_to_multiple_of,
- return_token_type_ids=return_token_type_ids,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_length=return_length,
- verbose=verbose,
- return_tensors=return_tensors,
- **kwargs,
- )
- if images is not None:
- image_features = self.image_processor(
- images,
- return_image_mask=return_image_mask,
- return_codebook_pixels=return_codebook_pixels,
- return_tensors=return_tensors,
- **kwargs,
- )
-
- if text is not None and images is not None:
- encoding.update(image_features)
- return encoding
- elif text is not None:
- return encoding
- else:
- return BatchEncoding(data=dict(**image_features), tensor_type=return_tensors)
-
- def batch_decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to BertTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
- refer to the docstring of this method for more information.
- """
- return self.tokenizer.batch_decode(*args, **kwargs)
-
- def decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to BertTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
- the docstring of this method for more information.
- """
- return self.tokenizer.decode(*args, **kwargs)
-
- @property
- def model_input_names(self):
- tokenizer_input_names = self.tokenizer.model_input_names
- image_processor_input_names = self.image_processor.model_input_names
- return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
-
- @property
- def feature_extractor_class(self):
- warnings.warn(
- "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.",
- FutureWarning,
- )
- return self.image_processor_class
-
- @property
- def feature_extractor(self):
- warnings.warn(
- "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.",
- FutureWarning,
- )
- return self.image_processor
diff --git a/transformers/models/fnet/__init__.py b/transformers/models/fnet/__init__.py
deleted file mode 100644
index 485160d1ccaa69b035e20c5710e9e5b319423816..0000000000000000000000000000000000000000
--- a/transformers/models/fnet/__init__.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# Copyright 2021 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_sentencepiece_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {"configuration_fnet": ["FNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "FNetConfig"]}
-
-try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_fnet"] = ["FNetTokenizer"]
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_fnet_fast"] = ["FNetTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_fnet"] = [
- "FNET_PRETRAINED_MODEL_ARCHIVE_LIST",
- "FNetForMaskedLM",
- "FNetForMultipleChoice",
- "FNetForNextSentencePrediction",
- "FNetForPreTraining",
- "FNetForQuestionAnswering",
- "FNetForSequenceClassification",
- "FNetForTokenClassification",
- "FNetLayer",
- "FNetModel",
- "FNetPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_fnet import FNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FNetConfig
-
- try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_fnet import FNetTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_fnet_fast import FNetTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_fnet import (
- FNET_PRETRAINED_MODEL_ARCHIVE_LIST,
- FNetForMaskedLM,
- FNetForMultipleChoice,
- FNetForNextSentencePrediction,
- FNetForPreTraining,
- FNetForQuestionAnswering,
- FNetForSequenceClassification,
- FNetForTokenClassification,
- FNetLayer,
- FNetModel,
- FNetPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/fnet/__pycache__/__init__.cpython-310.pyc b/transformers/models/fnet/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index dee2ead7ed5667fa6f3caca71431f8a17938d241..0000000000000000000000000000000000000000
Binary files a/transformers/models/fnet/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fnet/__pycache__/configuration_fnet.cpython-310.pyc b/transformers/models/fnet/__pycache__/configuration_fnet.cpython-310.pyc
deleted file mode 100644
index 81310373636e91c0b20f3ddc7d0f731baa7453f1..0000000000000000000000000000000000000000
Binary files a/transformers/models/fnet/__pycache__/configuration_fnet.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fnet/__pycache__/convert_fnet_original_flax_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/fnet/__pycache__/convert_fnet_original_flax_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 77b273b2b026811d93be05a7a952fdc9a75a72ee..0000000000000000000000000000000000000000
Binary files a/transformers/models/fnet/__pycache__/convert_fnet_original_flax_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fnet/__pycache__/modeling_fnet.cpython-310.pyc b/transformers/models/fnet/__pycache__/modeling_fnet.cpython-310.pyc
deleted file mode 100644
index a89007989c22ff69894a22687ac00aff8ea1b580..0000000000000000000000000000000000000000
Binary files a/transformers/models/fnet/__pycache__/modeling_fnet.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fnet/__pycache__/tokenization_fnet.cpython-310.pyc b/transformers/models/fnet/__pycache__/tokenization_fnet.cpython-310.pyc
deleted file mode 100644
index 7245703973a86a721f248eda6fd6e6c65fb1199c..0000000000000000000000000000000000000000
Binary files a/transformers/models/fnet/__pycache__/tokenization_fnet.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fnet/__pycache__/tokenization_fnet_fast.cpython-310.pyc b/transformers/models/fnet/__pycache__/tokenization_fnet_fast.cpython-310.pyc
deleted file mode 100644
index 0727f3a4a83611b1584bd7417a88b896172a0e9b..0000000000000000000000000000000000000000
Binary files a/transformers/models/fnet/__pycache__/tokenization_fnet_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fnet/configuration_fnet.py b/transformers/models/fnet/configuration_fnet.py
deleted file mode 100644
index 4678cae92e2a29a40317e61d3bd9371d941f7d30..0000000000000000000000000000000000000000
--- a/transformers/models/fnet/configuration_fnet.py
+++ /dev/null
@@ -1,119 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Google AI and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" FNet model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FNET_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class FNetConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FNetModel`]. It is used to instantiate an FNet
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the FNet
- [google/fnet-base](https://huggingface.co/google/fnet-base) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 32000):
- Vocabulary size of the FNet model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`FNetModel`] or [`TFFNetModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimension of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu_new"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- max_position_embeddings (`int`, *optional*, defaults to 512):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- type_vocab_size (`int`, *optional*, defaults to 4):
- The vocabulary size of the `token_type_ids` passed when calling [`FNetModel`] or [`TFFNetModel`].
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- use_tpu_fourier_optimizations (`bool`, *optional*, defaults to `False`):
- Determines whether to use TPU optimized FFTs. If `True`, the model will favor axis-wise FFTs transforms.
- Set to `False` for GPU/CPU hardware, in which case n-dimensional FFTs are used.
- tpu_short_seq_length (`int`, *optional*, defaults to 512):
- The sequence length that is expected by the model when using TPUs. This will be used to initialize the DFT
- matrix only when *use_tpu_fourier_optimizations* is set to `True` and the input sequence is shorter than or
- equal to 4096 tokens.
-
- Example:
-
- ```python
- >>> from transformers import FNetConfig, FNetModel
-
- >>> # Initializing a FNet fnet-base style configuration
- >>> configuration = FNetConfig()
-
- >>> # Initializing a model (with random weights) from the fnet-base style configuration
- >>> model = FNetModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "fnet"
-
- def __init__(
- self,
- vocab_size=32000,
- hidden_size=768,
- num_hidden_layers=12,
- intermediate_size=3072,
- hidden_act="gelu_new",
- hidden_dropout_prob=0.1,
- max_position_embeddings=512,
- type_vocab_size=4,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- use_tpu_fourier_optimizations=False,
- tpu_short_seq_length=512,
- pad_token_id=3,
- bos_token_id=1,
- eos_token_id=2,
- **kwargs,
- ):
- super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
-
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.initializer_range = initializer_range
- self.type_vocab_size = type_vocab_size
- self.layer_norm_eps = layer_norm_eps
- self.use_tpu_fourier_optimizations = use_tpu_fourier_optimizations
- self.tpu_short_seq_length = tpu_short_seq_length
diff --git a/transformers/models/fnet/convert_fnet_original_flax_checkpoint_to_pytorch.py b/transformers/models/fnet/convert_fnet_original_flax_checkpoint_to_pytorch.py
deleted file mode 100644
index f77a44874ae42919ccbdb32d35e8272074d80acc..0000000000000000000000000000000000000000
--- a/transformers/models/fnet/convert_fnet_original_flax_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert FNet checkpoint."""
-
-
-import argparse
-
-import torch
-from flax.training.checkpoints import restore_checkpoint
-
-from transformers import FNetConfig, FNetForPreTraining
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-
-
-def convert_flax_checkpoint_to_pytorch(flax_checkpoint_path, fnet_config_file, save_path):
- # Initialise PyTorch model
- config = FNetConfig.from_json_file(fnet_config_file)
- print(f"Building PyTorch model from configuration: {config}")
- fnet_pretraining_model = FNetForPreTraining(config)
-
- checkpoint_dict = restore_checkpoint(flax_checkpoint_path, None)
- pretrained_model_params = checkpoint_dict["target"]
-
- # Embeddings
- # Position IDs
- state_dict = fnet_pretraining_model.state_dict()
-
- position_ids = state_dict["fnet.embeddings.position_ids"]
- new_state_dict = {"fnet.embeddings.position_ids": position_ids}
- # Embedding Layers
- new_state_dict["fnet.embeddings.word_embeddings.weight"] = torch.tensor(
- pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]
- )
- new_state_dict["fnet.embeddings.position_embeddings.weight"] = torch.tensor(
- pretrained_model_params["encoder"]["embedder"]["position"]["embedding"][0]
- )
- new_state_dict["fnet.embeddings.token_type_embeddings.weight"] = torch.tensor(
- pretrained_model_params["encoder"]["embedder"]["type"]["embedding"]
- )
- new_state_dict["fnet.embeddings.projection.weight"] = torch.tensor(
- pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["kernel"]
- ).T
- new_state_dict["fnet.embeddings.projection.bias"] = torch.tensor(
- pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["bias"]
- )
- new_state_dict["fnet.embeddings.LayerNorm.weight"] = torch.tensor(
- pretrained_model_params["encoder"]["embedder"]["layer_norm"]["scale"]
- )
- new_state_dict["fnet.embeddings.LayerNorm.bias"] = torch.tensor(
- pretrained_model_params["encoder"]["embedder"]["layer_norm"]["bias"]
- )
-
- # Encoder Layers
- for layer in range(config.num_hidden_layers):
- new_state_dict[f"fnet.encoder.layer.{layer}.fourier.output.LayerNorm.weight"] = torch.tensor(
- pretrained_model_params["encoder"][f"encoder_{layer}"]["mixing_layer_norm"]["scale"]
- )
- new_state_dict[f"fnet.encoder.layer.{layer}.fourier.output.LayerNorm.bias"] = torch.tensor(
- pretrained_model_params["encoder"][f"encoder_{layer}"]["mixing_layer_norm"]["bias"]
- )
-
- new_state_dict[f"fnet.encoder.layer.{layer}.intermediate.dense.weight"] = torch.tensor(
- pretrained_model_params["encoder"][f"feed_forward_{layer}"]["intermediate"]["kernel"]
- ).T
- new_state_dict[f"fnet.encoder.layer.{layer}.intermediate.dense.bias"] = torch.tensor(
- pretrained_model_params["encoder"][f"feed_forward_{layer}"]["intermediate"]["bias"]
- )
-
- new_state_dict[f"fnet.encoder.layer.{layer}.output.dense.weight"] = torch.tensor(
- pretrained_model_params["encoder"][f"feed_forward_{layer}"]["output"]["kernel"]
- ).T
- new_state_dict[f"fnet.encoder.layer.{layer}.output.dense.bias"] = torch.tensor(
- pretrained_model_params["encoder"][f"feed_forward_{layer}"]["output"]["bias"]
- )
-
- new_state_dict[f"fnet.encoder.layer.{layer}.output.LayerNorm.weight"] = torch.tensor(
- pretrained_model_params["encoder"][f"encoder_{layer}"]["output_layer_norm"]["scale"]
- )
- new_state_dict[f"fnet.encoder.layer.{layer}.output.LayerNorm.bias"] = torch.tensor(
- pretrained_model_params["encoder"][f"encoder_{layer}"]["output_layer_norm"]["bias"]
- )
-
- # Pooler Layers
- new_state_dict["fnet.pooler.dense.weight"] = torch.tensor(pretrained_model_params["encoder"]["pooler"]["kernel"]).T
- new_state_dict["fnet.pooler.dense.bias"] = torch.tensor(pretrained_model_params["encoder"]["pooler"]["bias"])
-
- # Masked LM Layers
- new_state_dict["cls.predictions.transform.dense.weight"] = torch.tensor(
- pretrained_model_params["predictions_dense"]["kernel"]
- ).T
- new_state_dict["cls.predictions.transform.dense.bias"] = torch.tensor(
- pretrained_model_params["predictions_dense"]["bias"]
- )
- new_state_dict["cls.predictions.transform.LayerNorm.weight"] = torch.tensor(
- pretrained_model_params["predictions_layer_norm"]["scale"]
- )
- new_state_dict["cls.predictions.transform.LayerNorm.bias"] = torch.tensor(
- pretrained_model_params["predictions_layer_norm"]["bias"]
- )
- new_state_dict["cls.predictions.decoder.weight"] = torch.tensor(
- pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]
- )
- new_state_dict["cls.predictions.decoder.bias"] = torch.tensor(
- pretrained_model_params["predictions_output"]["output_bias"]
- )
- new_state_dict["cls.predictions.bias"] = torch.tensor(pretrained_model_params["predictions_output"]["output_bias"])
-
- # Seq Relationship Layers
- new_state_dict["cls.seq_relationship.weight"] = torch.tensor(
- pretrained_model_params["classification"]["output_kernel"]
- )
- new_state_dict["cls.seq_relationship.bias"] = torch.tensor(
- pretrained_model_params["classification"]["output_bias"]
- )
-
- # Load State Dict
- fnet_pretraining_model.load_state_dict(new_state_dict)
-
- # Save PreTrained
- print(f"Saving pretrained model to {save_path}")
- fnet_pretraining_model.save_pretrained(save_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--flax_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
- )
- parser.add_argument(
- "--fnet_config_file",
- default=None,
- type=str,
- required=True,
- help=(
- "The config json file corresponding to the pre-trained FNet model. \n"
- "This specifies the model architecture."
- ),
- )
- parser.add_argument("--save_path", default=None, type=str, required=True, help="Path to the output model.")
- args = parser.parse_args()
- convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)
diff --git a/transformers/models/fnet/modeling_fnet.py b/transformers/models/fnet/modeling_fnet.py
deleted file mode 100644
index 5724faee56cf857ae8f0deeb9bcc08e1766c4137..0000000000000000000000000000000000000000
--- a/transformers/models/fnet/modeling_fnet.py
+++ /dev/null
@@ -1,1182 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Google Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch FNet model."""
-
-import warnings
-from dataclasses import dataclass
-from functools import partial
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...utils import is_scipy_available
-
-
-if is_scipy_available():
- from scipy import linalg
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutput,
- BaseModelOutputWithPooling,
- MaskedLMOutput,
- ModelOutput,
- MultipleChoiceModelOutput,
- NextSentencePredictorOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import apply_chunking_to_forward
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_fnet import FNetConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "google/fnet-base"
-_CONFIG_FOR_DOC = "FNetConfig"
-
-
-from ..deprecated._archive_maps import FNET_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Adapted from https://github.com/google-research/google-research/blob/master/f_net/fourier.py
-def _two_dim_matmul(x, matrix_dim_one, matrix_dim_two):
- """Applies 2D matrix multiplication to 3D input arrays."""
- seq_length = x.shape[1]
- matrix_dim_one = matrix_dim_one[:seq_length, :seq_length]
- x = x.type(torch.complex64)
- return torch.einsum("bij,jk,ni->bnk", x, matrix_dim_two, matrix_dim_one)
-
-
-# # Adapted from https://github.com/google-research/google-research/blob/master/f_net/fourier.py
-def two_dim_matmul(x, matrix_dim_one, matrix_dim_two):
- return _two_dim_matmul(x, matrix_dim_one, matrix_dim_two)
-
-
-# Adapted from https://github.com/google-research/google-research/blob/master/f_net/fourier.py
-def fftn(x):
- """
- Applies n-dimensional Fast Fourier Transform (FFT) to input array.
-
- Args:
- x: Input n-dimensional array.
-
- Returns:
- n-dimensional Fourier transform of input n-dimensional array.
- """
- out = x
- for axis in reversed(range(x.ndim)[1:]): # We don't need to apply FFT to last axis
- out = torch.fft.fft(out, axis=axis)
- return out
-
-
-class FNetEmbeddings(nn.Module):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
- self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
-
- # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
- # any TensorFlow checkpoint file
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- # NOTE: This is the project layer and will be needed. The original code allows for different embedding and different model dimensions.
- self.projection = nn.Linear(config.hidden_size, config.hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
-
- self.register_buffer(
- "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
- )
-
- def forward(self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None):
- if input_ids is not None:
- input_shape = input_ids.size()
- else:
- input_shape = inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, :seq_length]
-
- # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
- # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
- # issue #5664
- if token_type_ids is None:
- if hasattr(self, "token_type_ids"):
- buffered_token_type_ids = self.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
-
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
- token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
- embeddings = inputs_embeds + token_type_embeddings
-
- position_embeddings = self.position_embeddings(position_ids)
- embeddings += position_embeddings
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.projection(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-class FNetBasicFourierTransform(nn.Module):
- def __init__(self, config):
- super().__init__()
- self._init_fourier_transform(config)
-
- def _init_fourier_transform(self, config):
- if not config.use_tpu_fourier_optimizations:
- self.fourier_transform = partial(torch.fft.fftn, dim=(1, 2))
- elif config.max_position_embeddings <= 4096:
- if is_scipy_available():
- self.register_buffer(
- "dft_mat_hidden", torch.tensor(linalg.dft(config.hidden_size), dtype=torch.complex64)
- )
- self.register_buffer(
- "dft_mat_seq", torch.tensor(linalg.dft(config.tpu_short_seq_length), dtype=torch.complex64)
- )
- self.fourier_transform = partial(
- two_dim_matmul, matrix_dim_one=self.dft_mat_seq, matrix_dim_two=self.dft_mat_hidden
- )
- else:
- logging.warning(
- "SciPy is needed for DFT matrix calculation and is not found. Using TPU optimized fast fourier"
- " transform instead."
- )
- self.fourier_transform = fftn
- else:
- self.fourier_transform = fftn
-
- def forward(self, hidden_states):
- # NOTE: We do not use torch.vmap as it is not integrated into PyTorch stable versions.
- # Interested users can modify the code to use vmap from the nightly versions, getting the vmap from here:
- # https://pytorch.org/docs/master/generated/torch.vmap.html. Note that fourier transform methods will need
- # change accordingly.
-
- outputs = self.fourier_transform(hidden_states).real
- return (outputs,)
-
-
-class FNetBasicOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.LayerNorm(input_tensor + hidden_states)
- return hidden_states
-
-
-class FNetFourierTransform(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.self = FNetBasicFourierTransform(config)
- self.output = FNetBasicOutput(config)
-
- def forward(self, hidden_states):
- self_outputs = self.self(hidden_states)
- fourier_output = self.output(self_outputs[0], hidden_states)
- outputs = (fourier_output,)
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->FNet
-class FNetIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->FNet
-class FNetOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class FNetLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1 # The dimension which has the sequence length
- self.fourier = FNetFourierTransform(config)
- self.intermediate = FNetIntermediate(config)
- self.output = FNetOutput(config)
-
- def forward(self, hidden_states):
- self_fourier_outputs = self.fourier(hidden_states)
- fourier_output = self_fourier_outputs[0]
-
- layer_output = apply_chunking_to_forward(
- self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, fourier_output
- )
-
- outputs = (layer_output,)
-
- return outputs
-
- def feed_forward_chunk(self, fourier_output):
- intermediate_output = self.intermediate(fourier_output)
- layer_output = self.output(intermediate_output, fourier_output)
- return layer_output
-
-
-class FNetEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([FNetLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(self, hidden_states, output_hidden_states=False, return_dict=True):
- all_hidden_states = () if output_hidden_states else None
-
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(layer_module.__call__, hidden_states)
- else:
- layer_outputs = layer_module(hidden_states)
-
- hidden_states = layer_outputs[0]
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
-
- return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->FNet
-class FNetPooler(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.activation = nn.Tanh()
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- # We "pool" the model by simply taking the hidden state corresponding
- # to the first token.
- first_token_tensor = hidden_states[:, 0]
- pooled_output = self.dense(first_token_tensor)
- pooled_output = self.activation(pooled_output)
- return pooled_output
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPredictionHeadTransform with Bert->FNet
-class FNetPredictionHeadTransform(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- if isinstance(config.hidden_act, str):
- self.transform_act_fn = ACT2FN[config.hidden_act]
- else:
- self.transform_act_fn = config.hidden_act
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.transform_act_fn(hidden_states)
- hidden_states = self.LayerNorm(hidden_states)
- return hidden_states
-
-
-class FNetLMPredictionHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.transform = FNetPredictionHeadTransform(config)
-
- # The output weights are the same as the input embeddings, but there is
- # an output-only bias for each token.
- self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
-
- self.bias = nn.Parameter(torch.zeros(config.vocab_size))
- self.decoder.bias = self.bias
-
- def forward(self, hidden_states):
- hidden_states = self.transform(hidden_states)
- hidden_states = self.decoder(hidden_states)
- return hidden_states
-
- def _tie_weights(self):
- # To tie those two weights if they get disconnected (on TPU or when the bias is resized)
- self.bias = self.decoder.bias
-
-
-class FNetOnlyMLMHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.predictions = FNetLMPredictionHead(config)
-
- def forward(self, sequence_output):
- prediction_scores = self.predictions(sequence_output)
- return prediction_scores
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOnlyNSPHead with Bert->FNet
-class FNetOnlyNSPHead(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.seq_relationship = nn.Linear(config.hidden_size, 2)
-
- def forward(self, pooled_output):
- seq_relationship_score = self.seq_relationship(pooled_output)
- return seq_relationship_score
-
-
-# Copied from transformers.models.bert.modeling_bert.BertPreTrainingHeads with Bert->FNet
-class FNetPreTrainingHeads(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.predictions = FNetLMPredictionHead(config)
- self.seq_relationship = nn.Linear(config.hidden_size, 2)
-
- def forward(self, sequence_output, pooled_output):
- prediction_scores = self.predictions(sequence_output)
- seq_relationship_score = self.seq_relationship(pooled_output)
- return prediction_scores, seq_relationship_score
-
-
-class FNetPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = FNetConfig
- base_model_prefix = "fnet"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- # NOTE: Original code uses same initialization as weights for biases as well.
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-@dataclass
-class FNetForPreTrainingOutput(ModelOutput):
- """
- Output type of [`FNetForPreTraining`].
-
- Args:
- loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
- Total loss as the sum of the masked language modeling loss and the next sequence prediction
- (classification) loss.
- prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
- seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`):
- Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
- before SoftMax).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
- plus the initial embedding outputs.
- """
-
- loss: Optional[torch.FloatTensor] = None
- prediction_logits: torch.FloatTensor = None
- seq_relationship_logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-
-
-FNET_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`FNetConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FNET_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare FNet Model transformer outputting raw hidden-states without any specific head on top.",
- FNET_START_DOCSTRING,
-)
-class FNetModel(FNetPreTrainedModel):
- """
-
- The model can behave as an encoder, following the architecture described in [FNet: Mixing Tokens with Fourier
- Transforms](https://arxiv.org/abs/2105.03824) by James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, Santiago Ontanon.
-
- """
-
- def __init__(self, config, add_pooling_layer=True):
- super().__init__(config)
- self.config = config
-
- self.embeddings = FNetEmbeddings(config)
- self.encoder = FNetEncoder(config)
-
- self.pooler = FNetPooler(config) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
- @add_start_docstrings_to_model_forward(FNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[tuple, BaseModelOutput]:
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = input_ids.size()
- batch_size, seq_length = input_shape
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- batch_size, seq_length = input_shape
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if (
- self.config.use_tpu_fourier_optimizations
- and seq_length <= 4096
- and self.config.tpu_short_seq_length != seq_length
- ):
- raise ValueError(
- "The `tpu_short_seq_length` in FNetConfig should be set equal to the sequence length being passed to"
- " the model when using TPU optimizations."
- )
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if token_type_ids is None:
- if hasattr(self.embeddings, "token_type_ids"):
- buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
- buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
- token_type_ids = buffered_token_type_ids_expanded
- else:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- )
- encoder_outputs = self.encoder(
- embedding_output,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
-
- pooler_output = self.pooler(sequence_output) if self.pooler is not None else None
-
- if not return_dict:
- return (sequence_output, pooler_output) + encoder_outputs[1:]
-
- return BaseModelOutputWithPooling(
- last_hidden_state=sequence_output,
- pooler_output=pooler_output,
- hidden_states=encoder_outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- FNet Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `next
- sentence prediction (classification)` head.
- """,
- FNET_START_DOCSTRING,
-)
-class FNetForPreTraining(FNetPreTrainedModel):
- _tied_weights_keys = ["cls.predictions.decoder.bias", "cls.predictions.decoder.weight"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.fnet = FNetModel(config)
- self.cls = FNetPreTrainingHeads(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.cls.predictions.decoder
-
- def set_output_embeddings(self, new_embeddings):
- self.cls.predictions.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(FNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=FNetForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- next_sentence_label: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, FNetForPreTrainingOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
- (see `input_ids` docstring) Indices should be in `[0, 1]`:
-
- - 0 indicates sequence B is a continuation of sequence A,
- - 1 indicates sequence B is a random sequence.
- kwargs (`Dict[str, any]`, optional, defaults to *{}*):
- Used to hide legacy arguments that have been deprecated.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, FNetForPreTraining
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("google/fnet-base")
- >>> model = FNetForPreTraining.from_pretrained("google/fnet-base")
- >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
- >>> outputs = model(**inputs)
- >>> prediction_logits = outputs.prediction_logits
- >>> seq_relationship_logits = outputs.seq_relationship_logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.fnet(
- input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output, pooled_output = outputs[:2]
- prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
-
- total_loss = None
- if labels is not None and next_sentence_label is not None:
- loss_fct = CrossEntropyLoss()
- masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
- next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
- total_loss = masked_lm_loss + next_sentence_loss
-
- if not return_dict:
- output = (prediction_scores, seq_relationship_score) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return FNetForPreTrainingOutput(
- loss=total_loss,
- prediction_logits=prediction_scores,
- seq_relationship_logits=seq_relationship_score,
- hidden_states=outputs.hidden_states,
- )
-
-
-@add_start_docstrings("""FNet Model with a `language modeling` head on top.""", FNET_START_DOCSTRING)
-class FNetForMaskedLM(FNetPreTrainedModel):
- _tied_weights_keys = ["cls.predictions.decoder.bias", "cls.predictions.decoder.weight"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.fnet = FNetModel(config)
- self.cls = FNetOnlyMLMHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.cls.predictions.decoder
-
- def set_output_embeddings(self, new_embeddings):
- self.cls.predictions.decoder = new_embeddings
-
- @add_start_docstrings_to_model_forward(FNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.fnet(
- input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- prediction_scores = self.cls(sequence_output)
-
- masked_lm_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss() # -100 index = padding token
- masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_scores,) + outputs[2:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return MaskedLMOutput(loss=masked_lm_loss, logits=prediction_scores, hidden_states=outputs.hidden_states)
-
-
-@add_start_docstrings(
- """FNet Model with a `next sentence prediction (classification)` head on top.""",
- FNET_START_DOCSTRING,
-)
-class FNetForNextSentencePrediction(FNetPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.fnet = FNetModel(config)
- self.cls = FNetOnlyNSPHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=NextSentencePredictorOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- **kwargs,
- ) -> Union[Tuple, NextSentencePredictorOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
- (see `input_ids` docstring). Indices should be in `[0, 1]`:
-
- - 0 indicates sequence B is a continuation of sequence A,
- - 1 indicates sequence B is a random sequence.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, FNetForNextSentencePrediction
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("google/fnet-base")
- >>> model = FNetForNextSentencePrediction.from_pretrained("google/fnet-base")
- >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
- >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
- >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
- >>> outputs = model(**encoding, labels=torch.LongTensor([1]))
- >>> logits = outputs.logits
- >>> assert logits[0, 0] < logits[0, 1] # next sentence was random
- ```"""
-
- if "next_sentence_label" in kwargs:
- warnings.warn(
- "The `next_sentence_label` argument is deprecated and will be removed in a future version, use"
- " `labels` instead.",
- FutureWarning,
- )
- labels = kwargs.pop("next_sentence_label")
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.fnet(
- input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- seq_relationship_scores = self.cls(pooled_output)
-
- next_sentence_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1))
-
- if not return_dict:
- output = (seq_relationship_scores,) + outputs[2:]
- return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output
-
- return NextSentencePredictorOutput(
- loss=next_sentence_loss,
- logits=seq_relationship_scores,
- hidden_states=outputs.hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- FNet Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
- output) e.g. for GLUE tasks.
- """,
- FNET_START_DOCSTRING,
-)
-class FNetForSequenceClassification(FNetPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.fnet = FNetModel(config)
-
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.fnet(
- input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
-
-
-@add_start_docstrings(
- """
- FNet Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- FNET_START_DOCSTRING,
-)
-class FNetForMultipleChoice(FNetPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.fnet = FNetModel(config)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.classifier = nn.Linear(config.hidden_size, 1)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FNET_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
- position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
- inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.fnet(
- input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- pooled_output = self.dropout(pooled_output)
- logits = self.classifier(pooled_output)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(loss=loss, logits=reshaped_logits, hidden_states=outputs.hidden_states)
-
-
-@add_start_docstrings(
- """
- FNet Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- FNET_START_DOCSTRING,
-)
-class FNetForTokenClassification(FNetPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.fnet = FNetModel(config)
-
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.fnet(
- input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output)
- logits = self.classifier(sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- # Only keep active parts of the loss
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
-
-
-@add_start_docstrings(
- """
- FNet Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- FNET_START_DOCSTRING,
-)
-class FNetForQuestionAnswering(FNetPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.num_labels = config.num_labels
-
- self.fnet = FNetModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.fnet(
- input_ids,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss, start_logits=start_logits, end_logits=end_logits, hidden_states=outputs.hidden_states
- )
diff --git a/transformers/models/fnet/tokenization_fnet.py b/transformers/models/fnet/tokenization_fnet.py
deleted file mode 100644
index a38114eb6d01ae6bee3c48e82513b7085b865ed2..0000000000000000000000000000000000000000
--- a/transformers/models/fnet/tokenization_fnet.py
+++ /dev/null
@@ -1,338 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Google Research, Google AI, Google Brain and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Tokenization classes for FNet model."""
-
-import os
-import unicodedata
-from shutil import copyfile
-from typing import Any, Dict, List, Optional, Tuple
-
-import sentencepiece as spm
-
-from ...tokenization_utils import AddedToken, PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}
-
-
-SPIECE_UNDERLINE = "▁"
-
-
-class FNetTokenizer(PreTrainedTokenizer):
- """
- Construct an FNet tokenizer. Adapted from [`AlbertTokenizer`]. Based on
- [SentencePiece](https://github.com/google/sentencepiece). This tokenizer inherits from [`PreTrainedTokenizer`]
- which contains most of the main methods. Users should refer to this superclass for more information regarding those
- methods.
-
- Args:
- vocab_file (`str`):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- do_lower_case (`bool`, *optional*, defaults to `False`):
- Whether or not to lowercase the input when tokenizing.
- remove_space (`bool`, *optional*, defaults to `True`):
- Whether or not to strip the text when tokenizing (removing excess spaces before and after the string).
- keep_accents (`bool`, *optional*, defaults to `True`):
- Whether or not to keep accents when tokenizing.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- sp_model_kwargs (`dict`, *optional*):
- Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
- SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
- to set:
-
- - `enable_sampling`: Enable subword regularization.
- - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
-
- - `nbest_size = {0,1}`: No sampling is performed.
- - `nbest_size > 1`: samples from the nbest_size results.
- - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
- using forward-filtering-and-backward-sampling algorithm.
- - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
- BPE-dropout.
-
- Attributes:
- sp_model (`SentencePieceProcessor`):
- The *SentencePiece* processor that is used for every conversion (string, tokens and IDs).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "token_type_ids"]
-
- def __init__(
- self,
- vocab_file,
- do_lower_case=False,
- remove_space=True,
- keep_accents=True,
- unk_token="",
- sep_token="[SEP]",
- pad_token="",
- cls_token="[CLS]",
- mask_token="[MASK]",
- sp_model_kwargs: Optional[Dict[str, Any]] = None,
- **kwargs,
- ) -> None:
- # Mask token behave like a normal word, i.e. include the space before it and
- # is included in the raw text, there should be a match in a non-normalized sentence.
- mask_token = AddedToken(mask_token, lstrip=True, special=True) if isinstance(mask_token, str) else mask_token
- cls_token = AddedToken(cls_token, special=True) if isinstance(cls_token, str) else cls_token
- sep_token = AddedToken(sep_token, special=True) if isinstance(sep_token, str) else sep_token
- mask_token = AddedToken(mask_token, special=True) if isinstance(mask_token, str) else mask_token
- self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
-
- self.do_lower_case = do_lower_case
- self.remove_space = remove_space
- self.keep_accents = keep_accents
- self.vocab_file = vocab_file
-
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(vocab_file)
-
- super().__init__(
- do_lower_case=do_lower_case,
- remove_space=remove_space,
- keep_accents=keep_accents,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- sp_model_kwargs=self.sp_model_kwargs,
- **kwargs,
- )
-
- @property
- def vocab_size(self):
- return len(self.sp_model)
-
- def get_vocab(self):
- vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
- vocab.update(self.added_tokens_encoder)
- return vocab
-
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sp_model"] = None
- return state
-
- def __setstate__(self, d):
- self.__dict__ = d
-
- # for backward compatibility
- if not hasattr(self, "sp_model_kwargs"):
- self.sp_model_kwargs = {}
-
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(self.vocab_file)
-
- def preprocess_text(self, inputs):
- if self.remove_space:
- outputs = " ".join(inputs.strip().split())
- else:
- outputs = inputs
- outputs = outputs.replace("``", '"').replace("''", '"')
-
- if not self.keep_accents:
- outputs = unicodedata.normalize("NFKD", outputs)
- outputs = "".join([c for c in outputs if not unicodedata.combining(c)])
- if self.do_lower_case:
- outputs = outputs.lower()
-
- return outputs
-
- def _tokenize(self, text: str) -> List[str]:
- """Tokenize a string."""
- text = self.preprocess_text(text)
- pieces = self.sp_model.encode(text, out_type=str)
- new_pieces = []
- for piece in pieces:
- if len(piece) > 1 and piece[-1] == str(",") and piece[-2].isdigit():
- cur_pieces = self.sp_model.EncodeAsPieces(piece[:-1].replace(SPIECE_UNDERLINE, ""))
- if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
- if len(cur_pieces[0]) == 1:
- cur_pieces = cur_pieces[1:]
- else:
- cur_pieces[0] = cur_pieces[0][1:]
- cur_pieces.append(piece[-1])
- new_pieces.extend(cur_pieces)
- else:
- new_pieces.append(piece)
-
- return new_pieces
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.sp_model.PieceToId(token)
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.sp_model.IdToPiece(index)
-
- # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.convert_tokens_to_string
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- current_sub_tokens = []
- out_string = ""
- prev_is_special = False
- for token in tokens:
- # make sure that special tokens are not decoded using sentencepiece model
- if token in self.all_special_tokens:
- if not prev_is_special:
- out_string += " "
- out_string += self.sp_model.decode(current_sub_tokens) + token
- prev_is_special = True
- current_sub_tokens = []
- else:
- current_sub_tokens.append(token)
- prev_is_special = False
- out_string += self.sp_model.decode(current_sub_tokens)
- return out_string.strip()
-
- def _decode(
- self,
- token_ids: List[int],
- skip_special_tokens: bool = False,
- clean_up_tokenization_spaces: bool = None,
- spaces_between_special_tokens: bool = False,
- **kwargs,
- ) -> str:
- text = super()._decode(
- token_ids=token_ids,
- skip_special_tokens=skip_special_tokens,
- clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- spaces_between_special_tokens=spaces_between_special_tokens,
- **kwargs,
- )
- # Mimic the behavior of the Rust tokenizer:
- # No space after
- if not spaces_between_special_tokens:
- text = text.replace(" ", "")
- return text
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. An FNet sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return cls + token_ids_0 + sep
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. An FNet sequence
- pair mask has the following format: :
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
-
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
- elif not os.path.isfile(self.vocab_file):
- with open(out_vocab_file, "wb") as fi:
- content_spiece_model = self.sp_model.serialized_model_proto()
- fi.write(content_spiece_model)
-
- return (out_vocab_file,)
diff --git a/transformers/models/fnet/tokenization_fnet_fast.py b/transformers/models/fnet/tokenization_fnet_fast.py
deleted file mode 100644
index f279ad9ca7d0e2547761222cb2f3023510160c98..0000000000000000000000000000000000000000
--- a/transformers/models/fnet/tokenization_fnet_fast.py
+++ /dev/null
@@ -1,187 +0,0 @@
-# coding=utf-8
-# Copyright 2021 Google AI, Google Brain and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Tokenization classes for FNet model."""
-
-
-import os
-from shutil import copyfile
-from typing import List, Optional, Tuple
-
-from ...tokenization_utils import AddedToken
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import is_sentencepiece_available, logging
-
-
-if is_sentencepiece_available():
- from .tokenization_fnet import FNetTokenizer
-else:
- FNetTokenizer = None
-
-logger = logging.get_logger(__name__)
-VOCAB_FILES_NAMES = {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}
-
-
-SPIECE_UNDERLINE = "▁"
-
-
-class FNetTokenizerFast(PreTrainedTokenizerFast):
- """
- Construct a "fast" FNetTokenizer (backed by HuggingFace's *tokenizers* library). Adapted from
- [`AlbertTokenizerFast`]. Based on
- [Unigram](https://huggingface.co/docs/tokenizers/python/latest/components.html?highlight=unigram#models). This
- tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods
-
- Args:
- vocab_file (`str`):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- do_lower_case (`bool`, *optional*, defaults to `False`):
- Whether or not to lowercase the input when tokenizing.
- remove_space (`bool`, *optional*, defaults to `True`):
- Whether or not to strip the text when tokenizing (removing excess spaces before and after the string).
- keep_accents (`bool`, *optional*, defaults to `True`):
- Whether or not to keep accents when tokenizing.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `"[SEP]"`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `"[CLS]"`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `"[MASK]"`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "token_type_ids"]
- slow_tokenizer_class = FNetTokenizer
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- do_lower_case=False,
- remove_space=True,
- keep_accents=True,
- unk_token="",
- sep_token="[SEP]",
- pad_token="",
- cls_token="[CLS]",
- mask_token="[MASK]",
- **kwargs,
- ):
- # Mask token behave like a normal word, i.e. include the space before it and
- # is included in the raw text, there should be a match in a non-normalized sentence.
- mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token
- cls_token = AddedToken(cls_token, lstrip=False, rstrip=False) if isinstance(cls_token, str) else cls_token
- sep_token = AddedToken(sep_token, lstrip=False, rstrip=False) if isinstance(sep_token, str) else sep_token
-
- super().__init__(
- vocab_file,
- tokenizer_file=tokenizer_file,
- do_lower_case=do_lower_case,
- remove_space=remove_space,
- keep_accents=keep_accents,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- **kwargs,
- )
-
- self.do_lower_case = do_lower_case
- self.remove_space = remove_space
- self.keep_accents = keep_accents
- self.vocab_file = vocab_file
-
- @property
- def can_save_slow_tokenizer(self) -> bool:
- return os.path.isfile(self.vocab_file) if self.vocab_file else False
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. An FNet sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: list of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return cls + token_ids_0 + sep
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An FNet
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- if token_ids_1 is None, only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of ids.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
-
- if token_ids_1 is None:
- return len(cls + token_ids_0 + sep) * [0]
- return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
-
- return (out_vocab_file,)
diff --git a/transformers/models/focalnet/__init__.py b/transformers/models/focalnet/__init__.py
deleted file mode 100644
index b043a006f9376609c774e84f5376323f48f2cae7..0000000000000000000000000000000000000000
--- a/transformers/models/focalnet/__init__.py
+++ /dev/null
@@ -1,59 +0,0 @@
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-# rely on isort to merge the imports
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {"configuration_focalnet": ["FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "FocalNetConfig"]}
-
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_focalnet"] = [
- "FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST",
- "FocalNetForImageClassification",
- "FocalNetForMaskedImageModeling",
- "FocalNetBackbone",
- "FocalNetModel",
- "FocalNetPreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_focalnet import FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FocalNetConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_focalnet import (
- FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST,
- FocalNetBackbone,
- FocalNetForImageClassification,
- FocalNetForMaskedImageModeling,
- FocalNetModel,
- FocalNetPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/focalnet/__pycache__/__init__.cpython-310.pyc b/transformers/models/focalnet/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index a34841c0766b2364280cc6439e2bb81925d52a30..0000000000000000000000000000000000000000
Binary files a/transformers/models/focalnet/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/focalnet/__pycache__/configuration_focalnet.cpython-310.pyc b/transformers/models/focalnet/__pycache__/configuration_focalnet.cpython-310.pyc
deleted file mode 100644
index 20bbe1b56cb26da48b71f4778e2d6d19e0d192ad..0000000000000000000000000000000000000000
Binary files a/transformers/models/focalnet/__pycache__/configuration_focalnet.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/focalnet/__pycache__/convert_focalnet_to_hf_format.cpython-310.pyc b/transformers/models/focalnet/__pycache__/convert_focalnet_to_hf_format.cpython-310.pyc
deleted file mode 100644
index 78c457f87f52ad4b10d2ae7d22e13ec42320b3c3..0000000000000000000000000000000000000000
Binary files a/transformers/models/focalnet/__pycache__/convert_focalnet_to_hf_format.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/focalnet/__pycache__/modeling_focalnet.cpython-310.pyc b/transformers/models/focalnet/__pycache__/modeling_focalnet.cpython-310.pyc
deleted file mode 100644
index 024cc6d6aa0560a5ba93ca3b86cc69465920e448..0000000000000000000000000000000000000000
Binary files a/transformers/models/focalnet/__pycache__/modeling_focalnet.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/focalnet/configuration_focalnet.py b/transformers/models/focalnet/configuration_focalnet.py
deleted file mode 100644
index 7f590b9c2c00a40dafaf54f6ae2131316b6674bb..0000000000000000000000000000000000000000
--- a/transformers/models/focalnet/configuration_focalnet.py
+++ /dev/null
@@ -1,164 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" FocalNet model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class FocalNetConfig(BackboneConfigMixin, PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FocalNetModel`]. It is used to instantiate a
- FocalNet model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the FocalNet
- [microsoft/focalnet-tiny](https://huggingface.co/microsoft/focalnet-tiny) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 4):
- The size (resolution) of each patch in the embeddings layer.
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- embed_dim (`int`, *optional*, defaults to 96):
- Dimensionality of patch embedding.
- use_conv_embed (`bool`, *optional*, defaults to `False`):
- Whether to use convolutional embedding. The authors noted that using convolutional embedding usually
- improve the performance, but it's not used by default.
- hidden_sizes (`List[int]`, *optional*, defaults to `[192, 384, 768, 768]`):
- Dimensionality (hidden size) at each stage.
- depths (`list(int)`, *optional*, defaults to `[2, 2, 6, 2]`):
- Depth (number of layers) of each stage in the encoder.
- focal_levels (`list(int)`, *optional*, defaults to `[2, 2, 2, 2]`):
- Number of focal levels in each layer of the respective stages in the encoder.
- focal_windows (`list(int)`, *optional*, defaults to `[3, 3, 3, 3]`):
- Focal window size in each layer of the respective stages in the encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder. If string, `"gelu"`, `"relu"`,
- `"selu"` and `"gelu_new"` are supported.
- mlp_ratio (`float`, *optional*, defaults to 4.0):
- Ratio of MLP hidden dimensionality to embedding dimensionality.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings and encoder.
- drop_path_rate (`float`, *optional*, defaults to 0.1):
- Stochastic depth rate.
- use_layerscale (`bool`, *optional*, defaults to `False`):
- Whether to use layer scale in the encoder.
- layerscale_value (`float`, *optional*, defaults to 0.0001):
- The initial value of the layer scale.
- use_post_layernorm (`bool`, *optional*, defaults to `False`):
- Whether to use post layer normalization in the encoder.
- use_post_layernorm_in_modulation (`bool`, *optional*, defaults to `False`):
- Whether to use post layer normalization in the modulation layer.
- normalize_modulator (`bool`, *optional*, defaults to `False`):
- Whether to normalize the modulator.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the layer normalization layers.
- encoder_stride (`int`, *optional*, defaults to 32):
- Factor to increase the spatial resolution by in the decoder head for masked image modeling.
- out_features (`List[str]`, *optional*):
- If used as backbone, list of features to output. Can be any of `"stem"`, `"stage1"`, `"stage2"`, etc.
- (depending on how many stages the model has). If unset and `out_indices` is set, will default to the
- corresponding stages. If unset and `out_indices` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
- out_indices (`List[int]`, *optional*):
- If used as backbone, list of indices of features to output. Can be any of 0, 1, 2, etc. (depending on how
- many stages the model has). If unset and `out_features` is set, will default to the corresponding stages.
- If unset and `out_features` is unset, will default to the last stage. Must be in the
- same order as defined in the `stage_names` attribute.
-
- Example:
-
- ```python
- >>> from transformers import FocalNetConfig, FocalNetModel
-
- >>> # Initializing a FocalNet microsoft/focalnet-tiny style configuration
- >>> configuration = FocalNetConfig()
-
- >>> # Initializing a model (with random weights) from the microsoft/focalnet-tiny style configuration
- >>> model = FocalNetModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "focalnet"
-
- def __init__(
- self,
- image_size=224,
- patch_size=4,
- num_channels=3,
- embed_dim=96,
- use_conv_embed=False,
- hidden_sizes=[192, 384, 768, 768],
- depths=[2, 2, 6, 2],
- focal_levels=[2, 2, 2, 2],
- focal_windows=[3, 3, 3, 3],
- hidden_act="gelu",
- mlp_ratio=4.0,
- hidden_dropout_prob=0.0,
- drop_path_rate=0.1,
- use_layerscale=False,
- layerscale_value=1e-4,
- use_post_layernorm=False,
- use_post_layernorm_in_modulation=False,
- normalize_modulator=False,
- initializer_range=0.02,
- layer_norm_eps=1e-5,
- encoder_stride=32,
- out_features=None,
- out_indices=None,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.embed_dim = embed_dim
- self.use_conv_embed = use_conv_embed
- self.hidden_sizes = hidden_sizes
- self.depths = depths
- self.focal_levels = focal_levels
- self.focal_windows = focal_windows
- self.hidden_act = hidden_act
- self.mlp_ratio = mlp_ratio
- self.hidden_dropout_prob = hidden_dropout_prob
- self.drop_path_rate = drop_path_rate
- self.use_layerscale = use_layerscale
- self.layerscale_value = layerscale_value
- self.use_post_layernorm = use_post_layernorm
- self.use_post_layernorm_in_modulation = use_post_layernorm_in_modulation
- self.normalize_modulator = normalize_modulator
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.encoder_stride = encoder_stride
- self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
- self._out_features, self._out_indices = get_aligned_output_features_output_indices(
- out_features=out_features, out_indices=out_indices, stage_names=self.stage_names
- )
diff --git a/transformers/models/focalnet/convert_focalnet_to_hf_format.py b/transformers/models/focalnet/convert_focalnet_to_hf_format.py
deleted file mode 100644
index 4aed15928062976c5f9589e2e6896e4e028b4eea..0000000000000000000000000000000000000000
--- a/transformers/models/focalnet/convert_focalnet_to_hf_format.py
+++ /dev/null
@@ -1,237 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert FocalNet checkpoints from the original repository. URL: https://github.com/microsoft/FocalNet/tree/main"""
-
-import argparse
-import json
-
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-from torchvision import transforms
-
-from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
-from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling
-
-
-def get_focalnet_config(model_name):
- depths = [2, 2, 6, 2] if "tiny" in model_name else [2, 2, 18, 2]
- use_conv_embed = True if "large" in model_name or "huge" in model_name else False
- use_post_layernorm = True if "large" in model_name or "huge" in model_name else False
- use_layerscale = True if "large" in model_name or "huge" in model_name else False
-
- if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
- if "fl3" in model_name:
- focal_levels = [3, 3, 3, 3]
- focal_windows = [5, 5, 5, 5]
- elif "fl4" in model_name:
- focal_levels = [4, 4, 4, 4]
- focal_windows = [3, 3, 3, 3]
-
- if "tiny" in model_name or "small" in model_name or "base" in model_name:
- focal_windows = [3, 3, 3, 3]
- if "lrf" in model_name:
- focal_levels = [3, 3, 3, 3]
- else:
- focal_levels = [2, 2, 2, 2]
-
- if "tiny" in model_name:
- embed_dim = 96
- elif "small" in model_name:
- embed_dim = 96
- elif "base" in model_name:
- embed_dim = 128
- elif "large" in model_name:
- embed_dim = 192
- elif "xlarge" in model_name:
- embed_dim = 256
- elif "huge" in model_name:
- embed_dim = 352
-
- # set label information
- repo_id = "huggingface/label-files"
- if "large" in model_name or "huge" in model_name:
- filename = "imagenet-22k-id2label.json"
- else:
- filename = "imagenet-1k-id2label.json"
-
- id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
- id2label = {int(k): v for k, v in id2label.items()}
- label2id = {v: k for k, v in id2label.items()}
-
- config = FocalNetConfig(
- embed_dim=embed_dim,
- depths=depths,
- focal_levels=focal_levels,
- focal_windows=focal_windows,
- use_conv_embed=use_conv_embed,
- id2label=id2label,
- label2id=label2id,
- use_post_layernorm=use_post_layernorm,
- use_layerscale=use_layerscale,
- )
-
- return config
-
-
-def rename_key(name):
- if "patch_embed.proj" in name:
- name = name.replace("patch_embed.proj", "embeddings.patch_embeddings.projection")
- if "patch_embed.norm" in name:
- name = name.replace("patch_embed.norm", "embeddings.norm")
- if "layers" in name:
- name = "encoder." + name
- if "encoder.layers" in name:
- name = name.replace("encoder.layers", "encoder.stages")
- if "downsample.proj" in name:
- name = name.replace("downsample.proj", "downsample.projection")
- if "blocks" in name:
- name = name.replace("blocks", "layers")
- if "modulation.f.weight" in name or "modulation.f.bias" in name:
- name = name.replace("modulation.f", "modulation.projection_in")
- if "modulation.h.weight" in name or "modulation.h.bias" in name:
- name = name.replace("modulation.h", "modulation.projection_context")
- if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
- name = name.replace("modulation.proj", "modulation.projection_out")
-
- if name == "norm.weight":
- name = "layernorm.weight"
- if name == "norm.bias":
- name = "layernorm.bias"
-
- if "head" in name:
- name = name.replace("head", "classifier")
- else:
- name = "focalnet." + name
-
- return name
-
-
-def convert_focalnet_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub=False):
- # fmt: off
- model_name_to_url = {
- "focalnet-tiny": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth",
- "focalnet-tiny-lrf": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth",
- "focalnet-small": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth",
- "focalnet-small-lrf": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth",
- "focalnet-base": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth",
- "focalnet-base-lrf": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth",
- "focalnet-large-lrf-fl3": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth",
- "focalnet-large-lrf-fl4": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth",
- "focalnet-xlarge-lrf-fl3": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth",
- "focalnet-xlarge-lrf-fl4": "https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth",
- }
- # fmt: on
-
- checkpoint_url = model_name_to_url[model_name]
- print("Checkpoint URL: ", checkpoint_url)
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["model"]
-
- # rename keys
- for key in state_dict.copy().keys():
- val = state_dict.pop(key)
- state_dict[rename_key(key)] = val
-
- config = get_focalnet_config(model_name)
- model = FocalNetForImageClassification(config)
- model.eval()
-
- # load state dict
- model.load_state_dict(state_dict)
-
- # verify conversion
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-
- processor = BitImageProcessor(
- do_resize=True,
- size={"shortest_edge": 256},
- resample=PILImageResampling.BILINEAR,
- do_center_crop=True,
- crop_size=224,
- do_normalize=True,
- image_mean=IMAGENET_DEFAULT_MEAN,
- image_std=IMAGENET_DEFAULT_STD,
- )
- image = Image.open(requests.get(url, stream=True).raw)
- inputs = processor(images=image, return_tensors="pt")
-
- image_transforms = transforms.Compose(
- [
- transforms.Resize(256),
- transforms.CenterCrop(224),
- transforms.ToTensor(),
- transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
- ]
- )
-
- original_pixel_values = image_transforms(image).unsqueeze(0)
-
- # verify pixel_values
- assert torch.allclose(inputs.pixel_values, original_pixel_values, atol=1e-4)
-
- outputs = model(**inputs)
-
- predicted_class_idx = outputs.logits.argmax(-1).item()
- print("Predicted class:", model.config.id2label[predicted_class_idx])
-
- print("First values of logits:", outputs.logits[0, :3])
-
- if model_name == "focalnet-tiny":
- expected_slice = torch.tensor([0.2166, -0.4368, 0.2191])
- elif model_name == "focalnet-tiny-lrf":
- expected_slice = torch.tensor([1.1669, 0.0125, -0.1695])
- elif model_name == "focalnet-small":
- expected_slice = torch.tensor([0.4917, -0.0430, 0.1341])
- elif model_name == "focalnet-small-lrf":
- expected_slice = torch.tensor([-0.2588, -0.5342, -0.2331])
- elif model_name == "focalnet-base":
- expected_slice = torch.tensor([-0.1655, -0.4090, -0.1730])
- elif model_name == "focalnet-base-lrf":
- expected_slice = torch.tensor([0.5306, -0.0483, -0.3928])
- assert torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4)
- print("Looks ok!")
-
- if pytorch_dump_folder_path is not None:
- print(f"Saving model and processor of {model_name} to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- print(f"Pushing model and processor of {model_name} to the hub...")
- model.push_to_hub(f"{model_name}")
- processor.push_to_hub(f"{model_name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="focalnet-tiny",
- type=str,
- help="Name of the FocalNet model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- help="Whether to push the model and processor to the hub.",
- )
-
- args = parser.parse_args()
- convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/focalnet/modeling_focalnet.py b/transformers/models/focalnet/modeling_focalnet.py
deleted file mode 100644
index a452f4171d1b6a65671a2592a76f74c6c8d07ebe..0000000000000000000000000000000000000000
--- a/transformers/models/focalnet/modeling_focalnet.py
+++ /dev/null
@@ -1,1032 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Microsoft Research and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch FocalNet model."""
-
-
-import collections.abc
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import BackboneOutput
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from ...utils.backbone_utils import BackboneMixin
-from .configuration_focalnet import FocalNetConfig
-
-
-logger = logging.get_logger(__name__)
-
-# General docstring
-_CONFIG_FOR_DOC = "FocalNetConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "microsoft/focalnet-tiny"
-_EXPECTED_OUTPUT_SHAPE = [1, 49, 768]
-
-# Image classification docstring
-_IMAGE_CLASS_CHECKPOINT = "microsoft/focalnet-tiny"
-_IMAGE_CLASS_EXPECTED_OUTPUT = "tabby, tabby cat"
-
-
-from ..deprecated._archive_maps import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-class FocalNetEncoderOutput(ModelOutput):
- """
- FocalNet encoder's outputs, with potential hidden states.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- last_hidden_state: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class FocalNetModelOutput(ModelOutput):
- """
- FocalNet model's outputs that also contains a pooling of the last hidden states.
-
- Args:
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*, returned when `add_pooling_layer=True` is passed):
- Average pooling of the last layer hidden-state.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- last_hidden_state: torch.FloatTensor = None
- pooler_output: Optional[torch.FloatTensor] = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class FocalNetMaskedImageModelingOutput(ModelOutput):
- """
- FocalNet masked image model outputs.
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `bool_masked_pos` is provided):
- Masked image modeling (MLM) loss.
- reconstruction (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Reconstructed pixel values.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- loss: Optional[torch.FloatTensor] = None
- reconstruction: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-
-
-@dataclass
-class FocalNetImageClassifierOutput(ModelOutput):
- """
- FocalNet outputs for image classification.
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
- Classification (or regression if config.num_labels==1) loss.
- logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
- Classification (or regression if config.num_labels==1) scores (before SoftMax).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
- shape `(batch_size, hidden_size, height, width)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
- include the spatial dimensions.
- """
-
- loss: Optional[torch.FloatTensor] = None
- logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- reshaped_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-
-
-class FocalNetEmbeddings(nn.Module):
- """
- Construct the patch embeddings and layernorm. Optionally, also the mask token.
- """
-
- def __init__(self, config, use_mask_token=False):
- super().__init__()
-
- self.patch_embeddings = FocalNetPatchEmbeddings(
- config=config,
- image_size=config.image_size,
- patch_size=config.patch_size,
- num_channels=config.num_channels,
- embed_dim=config.embed_dim,
- use_conv_embed=config.use_conv_embed,
- is_stem=True,
- )
- self.patch_grid = self.patch_embeddings.grid_size
- self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim)) if use_mask_token else None
-
- self.norm = nn.LayerNorm(config.embed_dim, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(
- self, pixel_values: Optional[torch.FloatTensor], bool_masked_pos: Optional[torch.BoolTensor] = None
- ) -> Tuple[torch.Tensor]:
- embeddings, output_dimensions = self.patch_embeddings(pixel_values)
- embeddings = self.norm(embeddings)
- batch_size, seq_len, _ = embeddings.size()
-
- if bool_masked_pos is not None:
- mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
- # replace the masked visual tokens by mask_tokens
- mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
- embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
-
- embeddings = self.dropout(embeddings)
- return embeddings, output_dimensions
-
-
-class FocalNetPatchEmbeddings(nn.Module):
- def __init__(
- self,
- config,
- image_size,
- patch_size,
- num_channels,
- embed_dim,
- add_norm=False,
- use_conv_embed=False,
- is_stem=False,
- ):
- super().__init__()
- image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
- patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
- num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.num_patches = num_patches
- self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])
-
- if use_conv_embed:
- # if we choose to use conv embedding, then we treat the stem and non-stem differently
- if is_stem:
- kernel_size = 7
- padding = 2
- stride = 4
- else:
- kernel_size = 3
- padding = 1
- stride = 2
- self.projection = nn.Conv2d(
- num_channels, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
- )
- else:
- self.projection = nn.Conv2d(num_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
-
- if add_norm:
- self.norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
- else:
- self.norm = None
-
- def maybe_pad(self, pixel_values, height, width):
- if width % self.patch_size[1] != 0:
- pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
- pixel_values = nn.functional.pad(pixel_values, pad_values)
- if height % self.patch_size[0] != 0:
- pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
- pixel_values = nn.functional.pad(pixel_values, pad_values)
- return pixel_values
-
- def forward(self, pixel_values: Optional[torch.FloatTensor]) -> Tuple[torch.Tensor, Tuple[int]]:
- _, num_channels, height, width = pixel_values.shape
- if num_channels != self.num_channels:
- raise ValueError(
- "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
- )
- # pad the input to be divisible by self.patch_size, if needed
- pixel_values = self.maybe_pad(pixel_values, height, width)
- embeddings = self.projection(pixel_values)
- _, _, height, width = embeddings.shape
- output_dimensions = (height, width)
- embeddings = embeddings.flatten(2).transpose(1, 2)
-
- if self.norm is not None:
- embeddings = self.norm(embeddings)
-
- return embeddings, output_dimensions
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->FocalNet
-class FocalNetDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-class FocalNetModulation(nn.Module):
- def __init__(self, config, index, dim, focal_factor=2, bias=True, projection_dropout=0.0):
- super().__init__()
-
- self.dim = dim
- self.focal_window = config.focal_windows[index]
- self.focal_level = config.focal_levels[index]
- self.focal_factor = focal_factor
- self.use_post_layernorm_in_modulation = config.use_post_layernorm_in_modulation
- self.normalize_modulator = config.normalize_modulator
-
- self.projection_in = nn.Linear(dim, 2 * dim + (self.focal_level + 1), bias=bias)
- self.projection_context = nn.Conv2d(dim, dim, kernel_size=1, stride=1, bias=bias)
-
- self.activation = nn.GELU()
- self.projection_out = nn.Linear(dim, dim)
- self.projection_dropout = nn.Dropout(projection_dropout)
- self.focal_layers = nn.ModuleList()
-
- self.kernel_sizes = []
- for k in range(self.focal_level):
- kernel_size = self.focal_factor * k + self.focal_window
- self.focal_layers.append(
- nn.Sequential(
- nn.Conv2d(
- dim, dim, kernel_size=kernel_size, stride=1, groups=dim, padding=kernel_size // 2, bias=False
- ),
- nn.GELU(),
- )
- )
- self.kernel_sizes.append(kernel_size)
- if self.use_post_layernorm_in_modulation:
- self.layernorm = nn.LayerNorm(dim, eps=config.layer_norm_eps)
-
- def forward(self, hidden_state):
- """
- Args:
- hidden_state:
- Input features with shape of (batch_size, height, width, num_channels)
- """
- num_channels = hidden_state.shape[-1]
-
- # pre linear projection
- x = self.projection_in(hidden_state).permute(0, 3, 1, 2).contiguous()
- q, ctx, self.gates = torch.split(x, (num_channels, num_channels, self.focal_level + 1), 1)
-
- # context aggreation
- ctx_all = 0
- for level in range(self.focal_level):
- ctx = self.focal_layers[level](ctx)
- ctx_all = ctx_all + ctx * self.gates[:, level : level + 1]
- ctx_global = self.activation(ctx.mean(2, keepdim=True).mean(3, keepdim=True))
- ctx_all = ctx_all + ctx_global * self.gates[:, self.focal_level :]
-
- # normalize context
- if self.normalize_modulator:
- ctx_all = ctx_all / (self.focal_level + 1)
-
- # focal modulation
- self.modulator = self.projection_context(ctx_all)
- x_out = q * self.modulator
- x_out = x_out.permute(0, 2, 3, 1).contiguous()
- if self.use_post_layernorm_in_modulation:
- x_out = self.layernorm(x_out)
-
- # post linear porjection
- x_out = self.projection_out(x_out)
- x_out = self.projection_dropout(x_out)
- return x_out
-
-
-class FocalNetMlp(nn.Module):
- def __init__(self, config, in_features, hidden_features=None, out_features=None, drop=0.0):
- super().__init__()
- out_features = out_features or in_features
- hidden_features = hidden_features or in_features
- self.fc1 = nn.Linear(in_features, hidden_features)
- self.activation = ACT2FN[config.hidden_act]
- self.fc2 = nn.Linear(hidden_features, out_features)
- self.drop = nn.Dropout(drop)
-
- def forward(self, hidden_state):
- hidden_state = self.fc1(hidden_state)
- hidden_state = self.activation(hidden_state)
- hidden_state = self.drop(hidden_state)
- hidden_state = self.fc2(hidden_state)
- hidden_state = self.drop(hidden_state)
- return hidden_state
-
-
-class FocalNetLayer(nn.Module):
- r"""Focal Modulation Network layer (block).
-
- Args:
- config (`FocalNetConfig`):
- Model config.
- index (`int`):
- Layer index.
- dim (`int`):
- Number of input channels.
- input_resolution (`Tuple[int]`):
- Input resulotion.
- drop_path (`float`, *optional*, defaults to 0.0):
- Stochastic depth rate.
- """
-
- def __init__(self, config, index, dim, input_resolution, drop_path=0.0):
- super().__init__()
-
- self.config = config
-
- # layer-specific attributes
- self.dim = dim
- self.input_resolution = input_resolution
-
- # general attributes
- self.drop = config.hidden_dropout_prob
- self.use_post_layernorm = config.use_post_layernorm
-
- self.norm1 = nn.LayerNorm(dim, eps=config.layer_norm_eps)
- self.modulation = FocalNetModulation(
- config=config,
- index=index,
- dim=dim,
- projection_dropout=self.drop,
- )
-
- self.drop_path = FocalNetDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
- self.norm2 = nn.LayerNorm(dim, eps=config.layer_norm_eps)
- mlp_hidden_dim = int(dim * config.mlp_ratio)
- self.mlp = FocalNetMlp(config=config, in_features=dim, hidden_features=mlp_hidden_dim, drop=self.drop)
-
- self.gamma_1 = 1.0
- self.gamma_2 = 1.0
- if config.use_layerscale:
- self.gamma_1 = nn.Parameter(config.layerscale_value * torch.ones((dim)), requires_grad=True)
- self.gamma_2 = nn.Parameter(config.layerscale_value * torch.ones((dim)), requires_grad=True)
-
- def forward(self, hidden_state, input_dimensions):
- height, width = input_dimensions
- batch_size, _, num_channels = hidden_state.shape
- shortcut = hidden_state
-
- # Focal Modulation
- hidden_state = hidden_state if self.use_post_layernorm else self.norm1(hidden_state)
- hidden_state = hidden_state.view(batch_size, height, width, num_channels)
- hidden_state = self.modulation(hidden_state).view(batch_size, height * width, num_channels)
- hidden_state = hidden_state if not self.use_post_layernorm else self.norm1(hidden_state)
-
- # FFN
- hidden_state = shortcut + self.drop_path(self.gamma_1 * hidden_state)
- hidden_state = hidden_state + self.drop_path(
- self.gamma_2
- * (self.norm2(self.mlp(hidden_state)) if self.use_post_layernorm else self.mlp(self.norm2(hidden_state)))
- )
-
- return hidden_state
-
-
-class FocalNetStage(nn.Module):
- def __init__(self, config, index, input_resolution):
- super().__init__()
-
- self.config = config
- self.num_stages = len(config.depths)
-
- embed_dim = [config.embed_dim * (2**i) for i in range(self.num_stages)]
- dim = embed_dim[index]
- out_dim = embed_dim[index + 1] if (index < self.num_stages - 1) else None
- downsample = FocalNetPatchEmbeddings if (index < self.num_stages - 1) else None
-
- # stochastic depth decay rule
- dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
- drop_path = dpr[sum(config.depths[:index]) : sum(config.depths[: index + 1])]
-
- self.layers = nn.ModuleList(
- [
- FocalNetLayer(
- config=config,
- index=index,
- dim=dim,
- input_resolution=input_resolution,
- drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
- )
- for i in range(config.depths[index])
- ]
- )
-
- if downsample is not None:
- self.downsample = downsample(
- config=config,
- image_size=input_resolution,
- patch_size=2,
- num_channels=dim,
- embed_dim=out_dim,
- add_norm=True,
- use_conv_embed=config.use_conv_embed,
- is_stem=False,
- )
- else:
- self.downsample = None
-
- self.pointing = False
-
- def forward(self, hidden_states: torch.Tensor, input_dimensions: Tuple[int, int]) -> Tuple[torch.Tensor]:
- height, width = input_dimensions
- for layer_module in self.layers:
- hidden_states = layer_module(hidden_states, input_dimensions)
-
- hidden_states_before_downsampling = hidden_states
- if self.downsample is not None:
- height, width = input_dimensions
- hidden_states = hidden_states.transpose(1, 2).reshape(
- hidden_states_before_downsampling.shape[0], -1, height, width
- )
- hidden_states, output_dimensions = self.downsample(hidden_states)
-
- else:
- output_dimensions = (height, width, height, width)
-
- stage_outputs = (hidden_states, hidden_states_before_downsampling, output_dimensions)
-
- return stage_outputs
-
-
-class FocalNetEncoder(nn.Module):
- def __init__(self, config, grid_size):
- super().__init__()
- self.num_stages = len(config.depths)
- self.config = config
-
- self.stages = nn.ModuleList(
- [
- FocalNetStage(
- config=config,
- index=i_layer,
- input_resolution=(grid_size[0] // (2**i_layer), grid_size[1] // (2**i_layer)),
- )
- for i_layer in range(self.num_stages)
- ]
- )
-
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- input_dimensions: Tuple[int, int],
- output_hidden_states: Optional[bool] = False,
- output_hidden_states_before_downsampling: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple, FocalNetEncoderOutput]:
- all_hidden_states = () if output_hidden_states else None
- all_reshaped_hidden_states = () if output_hidden_states else None
-
- if output_hidden_states:
- batch_size, _, hidden_size = hidden_states.shape
- # rearrange b (h w) c -> b c h w
- reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
- reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
-
- for i, stage_module in enumerate(self.stages):
- if self.gradient_checkpointing and self.training:
- stage_outputs = self._gradient_checkpointing_func(
- stage_module.__call__,
- hidden_states,
- input_dimensions,
- )
- else:
- stage_outputs = stage_module(hidden_states, input_dimensions)
-
- hidden_states = stage_outputs[0]
- hidden_states_before_downsampling = stage_outputs[1]
- output_dimensions = stage_outputs[2]
-
- input_dimensions = (output_dimensions[-2], output_dimensions[-1])
-
- if output_hidden_states and output_hidden_states_before_downsampling:
- batch_size, _, hidden_size = hidden_states_before_downsampling.shape
- # rearrange b (h w) c -> b c h w
- # here we use the original (not downsampled) height and width
- reshaped_hidden_state = hidden_states_before_downsampling.view(
- batch_size, *(output_dimensions[0], output_dimensions[1]), hidden_size
- )
- reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states_before_downsampling,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
- elif output_hidden_states and not output_hidden_states_before_downsampling:
- batch_size, _, hidden_size = hidden_states.shape
- # rearrange b (h w) c -> b c h w
- reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
- reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
- all_hidden_states += (hidden_states,)
- all_reshaped_hidden_states += (reshaped_hidden_state,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
-
- return FocalNetEncoderOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- reshaped_hidden_states=all_reshaped_hidden_states,
- )
-
-
-# Copied from transformers.models.swin.modeling_swin.SwinPreTrainedModel with Swin->FocalNet,swin->focalnet
-class FocalNetPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = FocalNetConfig
- base_model_prefix = "focalnet"
- main_input_name = "pixel_values"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-FOCALNET_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`FocalNetConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FOCALNET_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`AutoImageProcessor.__call__`] for details.
-
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare FocalNet Model outputting raw hidden-states without any specific head on top.",
- FOCALNET_START_DOCSTRING,
-)
-class FocalNetModel(FocalNetPreTrainedModel):
- def __init__(self, config, add_pooling_layer=True, use_mask_token=False):
- super().__init__(config)
- self.config = config
- self.num_stages = len(config.depths)
- self.num_features = int(config.embed_dim * 2 ** (self.num_stages - 1))
-
- self.embeddings = FocalNetEmbeddings(config, use_mask_token=use_mask_token)
- self.encoder = FocalNetEncoder(config, self.embeddings.patch_grid)
-
- self.layernorm = nn.LayerNorm(self.num_features, eps=config.layer_norm_eps)
- self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.patch_embeddings
-
- @add_start_docstrings_to_model_forward(FOCALNET_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=FocalNetModelOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, FocalNetModelOutput]:
- r"""
- bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
- """
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- embedding_output, input_dimensions = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
-
- encoder_outputs = self.encoder(
- embedding_output,
- input_dimensions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = encoder_outputs[0]
- sequence_output = self.layernorm(sequence_output)
-
- pooled_output = None
- if self.pooler is not None:
- pooled_output = self.pooler(sequence_output.transpose(1, 2))
- pooled_output = torch.flatten(pooled_output, 1)
-
- if not return_dict:
- output = (sequence_output, pooled_output) + encoder_outputs[1:]
-
- return output
-
- return FocalNetModelOutput(
- last_hidden_state=sequence_output,
- pooler_output=pooled_output,
- hidden_states=encoder_outputs.hidden_states,
- reshaped_hidden_states=encoder_outputs.reshaped_hidden_states,
- )
-
-
-@add_start_docstrings(
- """FocalNet Model with a decoder on top for masked image modeling.
-
- This follows the same implementation as in [SimMIM](https://arxiv.org/abs/2111.09886).
-
-
-
- Note that we provide a script to pre-train this model on custom data in our [examples
- directory](https://github.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining).
-
-
- """,
- FOCALNET_START_DOCSTRING,
-)
-class FocalNetForMaskedImageModeling(FocalNetPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.focalnet = FocalNetModel(config, add_pooling_layer=False, use_mask_token=True)
-
- self.num_stages = len(config.depths)
- num_features = int(config.embed_dim * 2 ** (self.num_stages - 1))
- self.decoder = nn.Sequential(
- nn.Conv2d(
- in_channels=num_features, out_channels=config.encoder_stride**2 * config.num_channels, kernel_size=1
- ),
- nn.PixelShuffle(config.encoder_stride),
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FOCALNET_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=FocalNetMaskedImageModelingOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- bool_masked_pos: Optional[torch.BoolTensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, FocalNetMaskedImageModelingOutput]:
- r"""
- bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
- Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
-
- Returns:
-
- Examples:
- ```python
- >>> from transformers import AutoImageProcessor, FocalNetConfig, FocalNetForMaskedImageModeling
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/focalnet-base-simmim-window6-192")
- >>> config = FocalNetConfig()
- >>> model = FocalNetForMaskedImageModeling(config)
-
- >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2
- >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
- >>> # create random boolean mask of shape (batch_size, num_patches)
- >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
-
- >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
- >>> loss, reconstructed_pixel_values = outputs.loss, outputs.logits
- >>> list(reconstructed_pixel_values.shape)
- [1, 3, 192, 192]
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.focalnet(
- pixel_values,
- bool_masked_pos=bool_masked_pos,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- # Reshape to (batch_size, num_channels, height, width)
- sequence_output = sequence_output.transpose(1, 2)
- batch_size, num_channels, sequence_length = sequence_output.shape
- height = width = math.floor(sequence_length**0.5)
- sequence_output = sequence_output.reshape(batch_size, num_channels, height, width)
-
- # Reconstruct pixel values
- reconstructed_pixel_values = self.decoder(sequence_output)
-
- masked_im_loss = None
- if bool_masked_pos is not None:
- size = self.config.image_size // self.config.patch_size
- bool_masked_pos = bool_masked_pos.reshape(-1, size, size)
- mask = (
- bool_masked_pos.repeat_interleave(self.config.patch_size, 1)
- .repeat_interleave(self.config.patch_size, 2)
- .unsqueeze(1)
- .contiguous()
- )
- reconstruction_loss = nn.functional.l1_loss(pixel_values, reconstructed_pixel_values, reduction="none")
- masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / self.config.num_channels
-
- if not return_dict:
- output = (reconstructed_pixel_values,) + outputs[2:]
- return ((masked_im_loss,) + output) if masked_im_loss is not None else output
-
- return FocalNetMaskedImageModelingOutput(
- loss=masked_im_loss,
- reconstruction=reconstructed_pixel_values,
- hidden_states=outputs.hidden_states,
- reshaped_hidden_states=outputs.reshaped_hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- FocalNet Model with an image classification head on top (a linear layer on top of the pooled output) e.g. for
- ImageNet.
- """,
- FOCALNET_START_DOCSTRING,
-)
-class FocalNetForImageClassification(FocalNetPreTrainedModel):
- # Copied from transformers.models.swin.modeling_swin.SwinForImageClassification.__init__ with Swin->FocalNet, swin->focalnet
- def __init__(self, config):
- super().__init__(config)
-
- self.num_labels = config.num_labels
- self.focalnet = FocalNetModel(config)
-
- # Classifier head
- self.classifier = (
- nn.Linear(self.focalnet.num_features, config.num_labels) if config.num_labels > 0 else nn.Identity()
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FOCALNET_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_IMAGE_CLASS_CHECKPOINT,
- output_type=FocalNetImageClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,
- )
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, FocalNetImageClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.focalnet(
- pixel_values,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- pooled_output = outputs[1]
-
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return FocalNetImageClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- reshaped_hidden_states=outputs.reshaped_hidden_states,
- )
-
-
-@add_start_docstrings(
- """
- FocalNet backbone, to be used with frameworks like X-Decoder.
- """,
- FOCALNET_START_DOCSTRING,
-)
-class FocalNetBackbone(FocalNetPreTrainedModel, BackboneMixin):
- def __init__(self, config: FocalNetConfig):
- super().__init__(config)
- super()._init_backbone(config)
-
- self.num_features = [config.embed_dim] + config.hidden_sizes
- self.focalnet = FocalNetModel(config)
-
- # initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FOCALNET_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.Tensor,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> BackboneOutput:
- """
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, AutoBackbone
- >>> import torch
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> processor = AutoImageProcessor.from_pretrained("microsoft/focalnet-tiny-lrf")
- >>> model = AutoBackbone.from_pretrained("microsoft/focalnet-tiny-lrf")
-
- >>> inputs = processor(image, return_tensors="pt")
- >>> outputs = model(**inputs)
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- outputs = self.focalnet(pixel_values, output_hidden_states=True, return_dict=True)
-
- hidden_states = outputs.reshaped_hidden_states
-
- feature_maps = ()
- for idx, stage in enumerate(self.stage_names):
- if stage in self.out_features:
- feature_maps += (hidden_states[idx],)
-
- if not return_dict:
- output = (feature_maps,)
- if output_hidden_states:
- output += (outputs.hidden_states,)
- return output
-
- return BackboneOutput(
- feature_maps=feature_maps,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=None,
- )
diff --git a/transformers/models/fsmt/__init__.py b/transformers/models/fsmt/__init__.py
deleted file mode 100644
index 65aba047469da14c6b25523fba31432e823ec47d..0000000000000000000000000000000000000000
--- a/transformers/models/fsmt/__init__.py
+++ /dev/null
@@ -1,49 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {
- "configuration_fsmt": ["FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP", "FSMTConfig"],
- "tokenization_fsmt": ["FSMTTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_fsmt"] = ["FSMTForConditionalGeneration", "FSMTModel", "PretrainedFSMTModel"]
-
-
-if TYPE_CHECKING:
- from .configuration_fsmt import FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP, FSMTConfig
- from .tokenization_fsmt import FSMTTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_fsmt import FSMTForConditionalGeneration, FSMTModel, PretrainedFSMTModel
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/fsmt/__pycache__/__init__.cpython-310.pyc b/transformers/models/fsmt/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index dd6bc7aa91fd15ba1271f7ac144420b7d4929364..0000000000000000000000000000000000000000
Binary files a/transformers/models/fsmt/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fsmt/__pycache__/configuration_fsmt.cpython-310.pyc b/transformers/models/fsmt/__pycache__/configuration_fsmt.cpython-310.pyc
deleted file mode 100644
index e32c47dafd613b913285bd3d6159513d56e14db1..0000000000000000000000000000000000000000
Binary files a/transformers/models/fsmt/__pycache__/configuration_fsmt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fsmt/__pycache__/convert_fsmt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/fsmt/__pycache__/convert_fsmt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 44077ab5bbdc96b604206b3bad6ccca4dce195c5..0000000000000000000000000000000000000000
Binary files a/transformers/models/fsmt/__pycache__/convert_fsmt_original_pytorch_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fsmt/__pycache__/modeling_fsmt.cpython-310.pyc b/transformers/models/fsmt/__pycache__/modeling_fsmt.cpython-310.pyc
deleted file mode 100644
index 17283672716f556834bb4e8bc7bc0c2af8b5f0bd..0000000000000000000000000000000000000000
Binary files a/transformers/models/fsmt/__pycache__/modeling_fsmt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fsmt/__pycache__/tokenization_fsmt.cpython-310.pyc b/transformers/models/fsmt/__pycache__/tokenization_fsmt.cpython-310.pyc
deleted file mode 100644
index 1763718e4c389b655969dc650f84cc38accb8f69..0000000000000000000000000000000000000000
Binary files a/transformers/models/fsmt/__pycache__/tokenization_fsmt.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fsmt/configuration_fsmt.py b/transformers/models/fsmt/configuration_fsmt.py
deleted file mode 100644
index 68abe47c019abaae981eb0beedcdbb7c755dff2e..0000000000000000000000000000000000000000
--- a/transformers/models/fsmt/configuration_fsmt.py
+++ /dev/null
@@ -1,219 +0,0 @@
-# coding=utf-8
-# Copyright 2019-present, Facebook, Inc and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" FSMT configuration"""
-
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class DecoderConfig(PretrainedConfig):
- r"""
- Configuration class for FSMT's decoder specific things. note: this is a private helper class
- """
-
- model_type = "fsmt_decoder"
-
- def __init__(self, vocab_size=0, bos_token_id=0):
- super().__init__()
- self.vocab_size = vocab_size
- self.bos_token_id = bos_token_id
-
-
-class FSMTConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FSMTModel`]. It is used to instantiate a FSMT
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the FSMT
- [facebook/wmt19-en-ru](https://huggingface.co/facebook/wmt19-en-ru) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- langs (`List[str]`):
- A list with source language and target_language (e.g., ['en', 'ru']).
- src_vocab_size (`int`):
- Vocabulary size of the encoder. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed to the forward method in the encoder.
- tgt_vocab_size (`int`):
- Vocabulary size of the decoder. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed to the forward method in the decoder.
- d_model (`int`, *optional*, defaults to 1024):
- Dimensionality of the layers and the pooler layer.
- encoder_layers (`int`, *optional*, defaults to 12):
- Number of encoder layers.
- decoder_layers (`int`, *optional*, defaults to 12):
- Number of decoder layers.
- encoder_attention_heads (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- decoder_attention_heads (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer decoder.
- decoder_ffn_dim (`int`, *optional*, defaults to 4096):
- Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
- encoder_ffn_dim (`int`, *optional*, defaults to 4096):
- Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
- activation_function (`str` or `Callable`, *optional*, defaults to `"relu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- activation_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for activations inside the fully connected layer.
- max_position_embeddings (`int`, *optional*, defaults to 1024):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- init_std (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- scale_embedding (`bool`, *optional*, defaults to `True`):
- Scale embeddings by diving by sqrt(d_model).
- bos_token_id (`int`, *optional*, defaults to 0)
- Beginning of stream token id.
- pad_token_id (`int`, *optional*, defaults to 1)
- Padding token id.
- eos_token_id (`int`, *optional*, defaults to 2)
- End of stream token id.
- decoder_start_token_id (`int`, *optional*):
- This model starts decoding with `eos_token_id`
- encoder_layerdrop (`float`, *optional*, defaults to 0.0):
- Google "layerdrop arxiv", as its not explainable in one line.
- decoder_layerdrop (`float`, *optional*, defaults to 0.0):
- Google "layerdrop arxiv", as its not explainable in one line.
- is_encoder_decoder (`bool`, *optional*, defaults to `True`):
- Whether this is an encoder/decoder model.
- tie_word_embeddings (`bool`, *optional*, defaults to `False`):
- Whether to tie input and output embeddings.
- num_beams (`int`, *optional*, defaults to 5)
- Number of beams for beam search that will be used by default in the `generate` method of the model. 1 means
- no beam search.
- length_penalty (`float`, *optional*, defaults to 1)
- Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent to
- the sequence length, which in turn is used to divide the score of the sequence. Since the score is the log
- likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences, while
- `length_penalty` < 0.0 encourages shorter sequences.
- early_stopping (`bool`, *optional*, defaults to `False`)
- Flag that will be used by default in the `generate` method of the model. Whether to stop the beam search
- when at least `num_beams` sentences are finished per batch or not.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
- forced_eos_token_id (`int`, *optional*, defaults to 2):
- The id of the token to force as the last generated token when `max_length` is reached. Usually set to
- `eos_token_id`.
-
- Examples:
-
- ```python
- >>> from transformers import FSMTConfig, FSMTModel
-
- >>> # Initializing a FSMT facebook/wmt19-en-ru style configuration
- >>> config = FSMTConfig()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = FSMTModel(config)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "fsmt"
- attribute_map = {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}
-
- # update the defaults from config file
- def __init__(
- self,
- langs=["en", "de"],
- src_vocab_size=42024,
- tgt_vocab_size=42024,
- activation_function="relu",
- d_model=1024,
- max_length=200,
- max_position_embeddings=1024,
- encoder_ffn_dim=4096,
- encoder_layers=12,
- encoder_attention_heads=16,
- encoder_layerdrop=0.0,
- decoder_ffn_dim=4096,
- decoder_layers=12,
- decoder_attention_heads=16,
- decoder_layerdrop=0.0,
- attention_dropout=0.0,
- dropout=0.1,
- activation_dropout=0.0,
- init_std=0.02,
- decoder_start_token_id=2,
- is_encoder_decoder=True,
- scale_embedding=True,
- tie_word_embeddings=False,
- num_beams=5,
- length_penalty=1.0,
- early_stopping=False,
- use_cache=True,
- pad_token_id=1,
- bos_token_id=0,
- eos_token_id=2,
- forced_eos_token_id=2,
- **common_kwargs,
- ):
- self.langs = langs
- self.src_vocab_size = src_vocab_size
- self.tgt_vocab_size = tgt_vocab_size
- self.d_model = d_model # encoder_embed_dim and decoder_embed_dim
-
- self.encoder_ffn_dim = encoder_ffn_dim
- self.encoder_layers = self.num_hidden_layers = encoder_layers
- self.encoder_attention_heads = encoder_attention_heads
- self.encoder_layerdrop = encoder_layerdrop
- self.decoder_layerdrop = decoder_layerdrop
- self.decoder_ffn_dim = decoder_ffn_dim
- self.decoder_layers = decoder_layers
- self.decoder_attention_heads = decoder_attention_heads
- self.max_position_embeddings = max_position_embeddings
- self.init_std = init_std # Normal(0, this parameter)
- self.activation_function = activation_function
-
- self.decoder = DecoderConfig(vocab_size=tgt_vocab_size, bos_token_id=eos_token_id)
- if "decoder" in common_kwargs:
- del common_kwargs["decoder"]
-
- self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True
-
- # 3 Types of Dropout
- self.attention_dropout = attention_dropout
- self.activation_dropout = activation_dropout
- self.dropout = dropout
-
- self.use_cache = use_cache
- super().__init__(
- pad_token_id=pad_token_id,
- bos_token_id=bos_token_id,
- eos_token_id=eos_token_id,
- decoder_start_token_id=decoder_start_token_id,
- is_encoder_decoder=is_encoder_decoder,
- tie_word_embeddings=tie_word_embeddings,
- forced_eos_token_id=forced_eos_token_id,
- max_length=max_length,
- num_beams=num_beams,
- length_penalty=length_penalty,
- early_stopping=early_stopping,
- **common_kwargs,
- )
diff --git a/transformers/models/fsmt/convert_fsmt_original_pytorch_checkpoint_to_pytorch.py b/transformers/models/fsmt/convert_fsmt_original_pytorch_checkpoint_to_pytorch.py
deleted file mode 100644
index ef2764f0ed10bace714f42f5f74ea6d9a147c613..0000000000000000000000000000000000000000
--- a/transformers/models/fsmt/convert_fsmt_original_pytorch_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,280 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Note: if you intend to run this script make sure you look under scripts/fsmt/
-# to locate the appropriate script to do the work correctly. There is a set of scripts to:
-# - download and prepare data and run the conversion script
-# - perform eval to get the best hparam into the config
-# - generate model_cards - useful if you have multiple models from the same paper
-
-import argparse
-import json
-import os
-import re
-from collections import OrderedDict
-from os.path import basename, dirname
-
-import fairseq
-import torch
-from fairseq import hub_utils
-from fairseq.data.dictionary import Dictionary
-
-from transformers import FSMTConfig, FSMTForConditionalGeneration
-from transformers.models.fsmt.tokenization_fsmt import VOCAB_FILES_NAMES
-from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE
-from transformers.utils import WEIGHTS_NAME, logging
-
-
-logging.set_verbosity_warning()
-
-json_indent = 2
-
-# based on the results of a search on a range of `num_beams`, `length_penalty` and `early_stopping`
-# values against wmt19 test data to obtain the best BLEU scores, we will use the following defaults:
-#
-# * `num_beams`: 5 (higher scores better, but requires more memory/is slower, can be adjusted by users)
-# * `early_stopping`: `False` consistently scored better
-# * `length_penalty` varied, so will assign the best one depending on the model
-best_score_hparams = {
- # fairseq:
- "wmt19-ru-en": {"length_penalty": 1.1},
- "wmt19-en-ru": {"length_penalty": 1.15},
- "wmt19-en-de": {"length_penalty": 1.0},
- "wmt19-de-en": {"length_penalty": 1.1},
- # allenai:
- "wmt16-en-de-dist-12-1": {"length_penalty": 0.6},
- "wmt16-en-de-dist-6-1": {"length_penalty": 0.6},
- "wmt16-en-de-12-1": {"length_penalty": 0.8},
- "wmt19-de-en-6-6-base": {"length_penalty": 0.6},
- "wmt19-de-en-6-6-big": {"length_penalty": 0.6},
-}
-
-# this remaps the different models to their organization names
-org_names = {}
-for m in ["wmt19-ru-en", "wmt19-en-ru", "wmt19-en-de", "wmt19-de-en"]:
- org_names[m] = "facebook"
-for m in [
- "wmt16-en-de-dist-12-1",
- "wmt16-en-de-dist-6-1",
- "wmt16-en-de-12-1",
- "wmt19-de-en-6-6-base",
- "wmt19-de-en-6-6-big",
-]:
- org_names[m] = "allenai"
-
-
-def rewrite_dict_keys(d):
- # (1) remove word breaking symbol, (2) add word ending symbol where the word is not broken up,
- # e.g.: d = {'le@@': 5, 'tt@@': 6, 'er': 7} => {'le': 5, 'tt': 6, 'er': 7}
- d2 = dict((re.sub(r"@@$", "", k), v) if k.endswith("@@") else (re.sub(r"$", "", k), v) for k, v in d.items())
- keep_keys = " ".split()
- # restore the special tokens
- for k in keep_keys:
- del d2[f"{k}"]
- d2[k] = d[k] # restore
- return d2
-
-
-def convert_fsmt_checkpoint_to_pytorch(fsmt_checkpoint_path, pytorch_dump_folder_path):
- # prep
- assert os.path.exists(fsmt_checkpoint_path)
- os.makedirs(pytorch_dump_folder_path, exist_ok=True)
- print(f"Writing results to {pytorch_dump_folder_path}")
-
- # handle various types of models
-
- checkpoint_file = basename(fsmt_checkpoint_path)
- fsmt_folder_path = dirname(fsmt_checkpoint_path)
-
- cls = fairseq.model_parallel.models.transformer.ModelParallelTransformerModel
- models = cls.hub_models()
- kwargs = {"bpe": "fastbpe", "tokenizer": "moses"}
- data_name_or_path = "."
- # note: since the model dump is old, fairseq has upgraded its model some
- # time later, and it does a whole lot of rewrites and splits on the saved
- # weights, therefore we can't use torch.load() directly on the model file.
- # see: upgrade_state_dict(state_dict) in fairseq_model.py
- print(f"using checkpoint {checkpoint_file}")
- chkpt = hub_utils.from_pretrained(
- fsmt_folder_path, checkpoint_file, data_name_or_path, archive_map=models, **kwargs
- )
-
- args = vars(chkpt["args"]["model"])
-
- src_lang = args["source_lang"]
- tgt_lang = args["target_lang"]
-
- data_root = dirname(pytorch_dump_folder_path)
- model_dir = basename(pytorch_dump_folder_path)
-
- # dicts
- src_dict_file = os.path.join(fsmt_folder_path, f"dict.{src_lang}.txt")
- tgt_dict_file = os.path.join(fsmt_folder_path, f"dict.{tgt_lang}.txt")
-
- src_dict = Dictionary.load(src_dict_file)
- src_vocab = rewrite_dict_keys(src_dict.indices)
- src_vocab_size = len(src_vocab)
- src_vocab_file = os.path.join(pytorch_dump_folder_path, "vocab-src.json")
- print(f"Generating {src_vocab_file} of {src_vocab_size} of {src_lang} records")
- with open(src_vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(src_vocab, ensure_ascii=False, indent=json_indent))
-
- # detect whether this is a do_lower_case situation, which can be derived by checking whether we
- # have at least one uppercase letter in the source vocab
- do_lower_case = True
- for k in src_vocab.keys():
- if not k.islower():
- do_lower_case = False
- break
-
- tgt_dict = Dictionary.load(tgt_dict_file)
- tgt_vocab = rewrite_dict_keys(tgt_dict.indices)
- tgt_vocab_size = len(tgt_vocab)
- tgt_vocab_file = os.path.join(pytorch_dump_folder_path, "vocab-tgt.json")
- print(f"Generating {tgt_vocab_file} of {tgt_vocab_size} of {tgt_lang} records")
- with open(tgt_vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(tgt_vocab, ensure_ascii=False, indent=json_indent))
-
- # merges_file (bpecodes)
- merges_file = os.path.join(pytorch_dump_folder_path, VOCAB_FILES_NAMES["merges_file"])
- for fn in ["bpecodes", "code"]: # older fairseq called the merges file "code"
- fsmt_merges_file = os.path.join(fsmt_folder_path, fn)
- if os.path.exists(fsmt_merges_file):
- break
- with open(fsmt_merges_file, encoding="utf-8") as fin:
- merges = fin.read()
- merges = re.sub(r" \d+$", "", merges, 0, re.M) # remove frequency number
- print(f"Generating {merges_file}")
- with open(merges_file, "w", encoding="utf-8") as fout:
- fout.write(merges)
-
- # model config
- fsmt_model_config_file = os.path.join(pytorch_dump_folder_path, "config.json")
-
- # validate bpe/tokenizer config, as currently it's hardcoded to moses+fastbpe -
- # may have to modify the tokenizer if a different type is used by a future model
- assert args["bpe"] == "fastbpe", f"need to extend tokenizer to support bpe={args['bpe']}"
- assert args["tokenizer"] == "moses", f"need to extend tokenizer to support bpe={args['tokenizer']}"
-
- model_conf = {
- "architectures": ["FSMTForConditionalGeneration"],
- "model_type": "fsmt",
- "activation_dropout": args["activation_dropout"],
- "activation_function": "relu",
- "attention_dropout": args["attention_dropout"],
- "d_model": args["decoder_embed_dim"],
- "dropout": args["dropout"],
- "init_std": 0.02,
- "max_position_embeddings": args["max_source_positions"],
- "num_hidden_layers": args["encoder_layers"],
- "src_vocab_size": src_vocab_size,
- "tgt_vocab_size": tgt_vocab_size,
- "langs": [src_lang, tgt_lang],
- "encoder_attention_heads": args["encoder_attention_heads"],
- "encoder_ffn_dim": args["encoder_ffn_embed_dim"],
- "encoder_layerdrop": args["encoder_layerdrop"],
- "encoder_layers": args["encoder_layers"],
- "decoder_attention_heads": args["decoder_attention_heads"],
- "decoder_ffn_dim": args["decoder_ffn_embed_dim"],
- "decoder_layerdrop": args["decoder_layerdrop"],
- "decoder_layers": args["decoder_layers"],
- "bos_token_id": 0,
- "pad_token_id": 1,
- "eos_token_id": 2,
- "is_encoder_decoder": True,
- "scale_embedding": not args["no_scale_embedding"],
- "tie_word_embeddings": args["share_all_embeddings"],
- }
-
- # good hparam defaults to start with
- model_conf["num_beams"] = 5
- model_conf["early_stopping"] = False
- if model_dir in best_score_hparams and "length_penalty" in best_score_hparams[model_dir]:
- model_conf["length_penalty"] = best_score_hparams[model_dir]["length_penalty"]
- else:
- model_conf["length_penalty"] = 1.0
-
- print(f"Generating {fsmt_model_config_file}")
- with open(fsmt_model_config_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(model_conf, ensure_ascii=False, indent=json_indent))
-
- # tokenizer config
- fsmt_tokenizer_config_file = os.path.join(pytorch_dump_folder_path, TOKENIZER_CONFIG_FILE)
-
- tokenizer_conf = {
- "langs": [src_lang, tgt_lang],
- "model_max_length": 1024,
- "do_lower_case": do_lower_case,
- }
-
- print(f"Generating {fsmt_tokenizer_config_file}")
- with open(fsmt_tokenizer_config_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(tokenizer_conf, ensure_ascii=False, indent=json_indent))
-
- # model
- model = chkpt["models"][0]
- model_state_dict = model.state_dict()
-
- # rename keys to start with 'model.'
- model_state_dict = OrderedDict(("model." + k, v) for k, v in model_state_dict.items())
-
- # remove unneeded keys
- ignore_keys = [
- "model.model",
- "model.encoder.version",
- "model.decoder.version",
- "model.encoder_embed_tokens.weight",
- "model.decoder_embed_tokens.weight",
- "model.encoder.embed_positions._float_tensor",
- "model.decoder.embed_positions._float_tensor",
- ]
- for k in ignore_keys:
- model_state_dict.pop(k, None)
-
- config = FSMTConfig.from_pretrained(pytorch_dump_folder_path)
- model_new = FSMTForConditionalGeneration(config)
-
- # check that it loads ok
- model_new.load_state_dict(model_state_dict, strict=False)
-
- # save
- pytorch_weights_dump_path = os.path.join(pytorch_dump_folder_path, WEIGHTS_NAME)
- print(f"Generating {pytorch_weights_dump_path}")
- torch.save(model_state_dict, pytorch_weights_dump_path)
-
- print("Conversion is done!")
- print("\nLast step is to upload the files to s3")
- print(f"cd {data_root}")
- print(f"transformers-cli upload {model_dir}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--fsmt_checkpoint_path",
- default=None,
- type=str,
- required=True,
- help=(
- "Path to the official PyTorch checkpoint file which is expected to reside in the dump dir with dicts,"
- " bpecodes, etc."
- ),
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
- args = parser.parse_args()
- convert_fsmt_checkpoint_to_pytorch(args.fsmt_checkpoint_path, args.pytorch_dump_folder_path)
diff --git a/transformers/models/fsmt/modeling_fsmt.py b/transformers/models/fsmt/modeling_fsmt.py
deleted file mode 100644
index 4c180c52678b82b80c5a1aa43f42292556c8f1e4..0000000000000000000000000000000000000000
--- a/transformers/models/fsmt/modeling_fsmt.py
+++ /dev/null
@@ -1,1386 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-# Original implementation: https://github.com/pytorch/fairseq/tree/master/examples/wmt19
-# Authors:
-# - @alexeib Alexei Baevski
-# - @edunov Sergey Edunov
-# - @michaelauli Michael Auli
-# - @myleott Myle Ott
-# - @nng555 Nathan Ng
-# - David Grangier
-# - Kyra Yee
-#
-# Paper: Facebook FAIR's WMT19 News Translation Task Submission https://arxiv.org/abs/1907.06616
-#
-"""PyTorch Fairseq model, ported from https://github.com/pytorch/fairseq/tree/master/examples/wmt19"""
-
-import math
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import torch
-from torch import Tensor, nn
-from torch.nn import CrossEntropyLoss, LayerNorm
-
-from ...activations import ACT2FN
-from ...integrations.deepspeed import is_deepspeed_zero3_enabled
-from ...modeling_outputs import (
- BaseModelOutput,
- BaseModelOutputWithPastAndCrossAttentions,
- Seq2SeqLMOutput,
- Seq2SeqModelOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- add_code_sample_docstrings,
- add_end_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_fsmt import FSMTConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "facebook/wmt19-ru-en"
-_CONFIG_FOR_DOC = "FSMTConfig"
-
-# See all FSMT models at https://huggingface.co/models?filter=fsmt
-
-# Porting notes:
-# this one is modeled after BartModel*
-#
-# Currently only translation (fairseq also has weights for LM)
-#
-# fairseq provides weights for ru-en, en-ru and de-en, en-de pairs. All have been ported.
-# - ru-en, en-ru use asymmetric vocab
-# - de-en, en-de use a merged single vocab (but the code works as if they are separate)
-#
-# Differences with Bart:
-# - not using bos token
-# - 2 separate vocabs (src and target)
-# - embed weights aren't tied
-# - uses a model Ensemble (but that part isn't ported/implemented yet) - so we
-# aren't getting as good of a BLEU score
-# - uses a projection layer at the end of the decoder
-# - doesn't use final_logits_bias
-# - beam search: stops as soon as num_beams == len(hypos) (whereas transformers
-# is not satisfied there and will continue searching until the next cycles
-# aren't promising something better), comparing BLEU scores - the transformers
-# algorithm is slightly superior, therefore using the latter. But if you want
-# to match fairseq outputs, you need to pass ``early_stopping=True`` to ``generate()``.
-#
-# SinusoidalPositionalEmbedding is slightly different from Bart's - generates
-# different embeddings. This implementation is copied verbatim from fairseq with
-# some small changes to make it work here.
-#
-# Other changes:
-# - doesn't support use_cache as Bart's version does
-#
-#
-# FSMTConfig changes with BartConfig
-#
-# Differences with BART:
-# - src/tgt vocabs aren't shared
-# - token embeddings aren't shared
-# - needs a language pair
-# - scale_embedding are True
-#
-# some unused args were removed too
-#
-#
-# TODO:
-# - port model ensemble (fs uses 4 model checkpoints)
-# - solve beam search discrepancies
-# docstyle-ignore
-
-"""
-
-Here is how to compare BLEU scores against fairseq implementation:
-
-# en-ru
-
-export PAIR=en-ru
-export DATA_DIR=data/$PAIR
-export SAVE_DIR=data/$PAIR
-export BS=8
-export NUM_BEAMS=50
-mkdir -p $DATA_DIR
-sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
-sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
-echo $PAIR
-PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
-
-# (fairseq BLEU: 36.4 http://matrix.statmt.org/matrix/output/1914?score_id=37605)
-
-
-# ru-en
-
-export PAIR=ru-en
-export DATA_DIR=data/$PAIR
-export SAVE_DIR=data/$PAIR
-export BS=8
-export NUM_BEAMS=50
-mkdir -p $DATA_DIR
-sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
-sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
-PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
-
-
-# (fairseq BLEU: 41.3 http://matrix.statmt.org/matrix/output/1907?run_id=6937)
-
-
-# de-en
-
-export PAIR=de-en
-export DATA_DIR=data/$PAIR
-export SAVE_DIR=data/$PAIR
-export BS=8
-export NUM_BEAMS=50
-mkdir -p $DATA_DIR
-sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
-sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
-echo $PAIR
-PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
-
-# (fairseq BLEU: 42.3 http://matrix.statmt.org/matrix/output/1902?run_id=6750)
-
-
-
-# en-de
-
-export PAIR=en-de
-export DATA_DIR=data/$PAIR
-export SAVE_DIR=data/$PAIR
-export BS=8
-mkdir -p $DATA_DIR
-sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
-sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
-echo $PAIR
-PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
-
-# (fairseq BLEU: 43.1 http://matrix.statmt.org/matrix/output/1909?run_id=6862)
-
-"""
-
-
-FSMT_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`FSMTConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-
-"""
-FSMT_GENERATION_EXAMPLE = r"""
- Translation example::
-
- ```python
- >>> from transformers import AutoTokenizer, FSMTForConditionalGeneration
-
- >>> mname = "facebook/wmt19-ru-en"
- >>> model = FSMTForConditionalGeneration.from_pretrained(mname)
- >>> tokenizer = AutoTokenizer.from_pretrained(mname)
-
- >>> src_text = "Машинное обучение - это здорово, не так ли?"
- >>> input_ids = tokenizer(src_text, return_tensors="pt").input_ids
- >>> outputs = model.generate(input_ids, num_beams=5, num_return_sequences=3)
- >>> tokenizer.decode(outputs[0], skip_special_tokens=True)
- "Machine learning is great, isn't it?"
- ```
-
-"""
-
-FSMT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`FSTMTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
- Indices of decoder input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are decoder input IDs?](../glossary#decoder-input-ids)
-
- FSMT uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values`
- is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`).
- decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
- Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
- be used by default.
- head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*):
- Mask to nullify selected heads of the attention modules in the encoder. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- decoder_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
- Mask to nullify selected heads of the attention modules in the decoder. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
- Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0,
- 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- encoder_outputs (`Tuple(torch.FloatTensor)`, *optional*):
- Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)
- `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` is a sequence of hidden-states at
- the output of the last layer of the encoder. Used in the cross-attention of the decoder.
- past_key_values (`Tuple(torch.FloatTensor)` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden-states of the attention blocks. Can be used to speed up decoding.
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded
- representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be
- input (see `past_key_values`). This is useful if you want more control over how to convert
- `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix.
-
- If `decoder_input_ids` and `decoder_inputs_embeds` are both unset, `decoder_inputs_embeds` takes the value
- of `inputs_embeds`.
- use_cache (`bool`, *optional*, defaults to `True`):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-def invert_mask(attention_mask):
- """Turns 1->0, 0->1, False->True, True-> False"""
- assert attention_mask.dim() == 2
- return attention_mask.eq(0)
-
-
-def triu_onnx(x, diagonal=0):
- l = x.shape[0]
- arange = torch.arange(l, device=x.device)
- mask = arange.expand(l, l)
- arange = arange.unsqueeze(-1)
- if diagonal:
- arange = arange + diagonal
- mask = mask >= arange
- return x.masked_fill(mask == 0, 0)
-
-
-def _prepare_fsmt_decoder_inputs(
- config,
- input_ids,
- decoder_input_ids=None,
- decoder_padding_mask=None,
- causal_mask_dtype=torch.float32,
-):
- """
- Prepare masks that ignore padding tokens in the decoder and a causal mask for the decoder if none are provided.
- This mimics the default behavior in fairseq. To override it pass in masks. Note: this is not called during
- generation
- """
- pad_token_id = config.pad_token_id
- if decoder_input_ids is None:
- decoder_input_ids = shift_tokens_right(input_ids, pad_token_id)
- bsz, tgt_len = decoder_input_ids.size()
- if decoder_padding_mask is None:
- decoder_padding_mask = make_padding_mask(decoder_input_ids, pad_token_id)
- else:
- decoder_padding_mask = invert_mask(decoder_padding_mask)
- causal_mask = triu_onnx(fill_with_neg_inf(torch.zeros(tgt_len, tgt_len, dtype=causal_mask_dtype)), 1).to(
- device=decoder_input_ids.device
- )
- return decoder_input_ids, decoder_padding_mask, causal_mask
-
-
-class PretrainedFSMTModel(PreTrainedModel):
- config_class = FSMTConfig
- base_model_prefix = "model"
-
- def _init_weights(self, module):
- std = self.config.init_std
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, SinusoidalPositionalEmbedding):
- pass
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
- @property
- def dummy_inputs(self):
- pad_token = self.config.pad_token_id
- input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device)
- dummy_inputs = {
- "attention_mask": input_ids.ne(pad_token),
- "input_ids": input_ids,
- }
- return dummy_inputs
-
-
-def _make_linear_from_emb(emb):
- vocab_size, emb_size = emb.weight.shape
- lin_layer = nn.Linear(vocab_size, emb_size, bias=False)
- lin_layer.weight.data = emb.weight.data
- return lin_layer
-
-
-# Helper Functions, mostly for making masks
-def _check_shapes(shape_1, shape2):
- if shape_1 != shape2:
- raise AssertionError(f"shape mismatch: {shape_1} != {shape2}")
-
-
-def shift_tokens_right(input_ids, pad_token_id):
- """Shift input ids one token to the right, and wrap the last non pad token (usually )."""
-
- # replace possible -100 values in labels by `pad_token_id`
- input_ids.masked_fill_(input_ids == -100, pad_token_id)
-
- prev_output_tokens = input_ids.clone()
- index_of_eos = (input_ids.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
- prev_output_tokens[:, 0] = input_ids.gather(1, index_of_eos).squeeze()
- prev_output_tokens[:, 1:] = input_ids[:, :-1]
- return prev_output_tokens
-
-
-def make_padding_mask(input_ids, padding_idx=1):
- """True for pad tokens"""
- padding_mask = input_ids.eq(padding_idx)
- if not padding_mask.any():
- padding_mask = None
- return padding_mask
-
-
-# Helper Modules
-
-
-class EncoderLayer(nn.Module):
- def __init__(self, config: FSMTConfig):
- super().__init__()
- self.embed_dim = config.d_model
- self.self_attn = Attention(self.embed_dim, config.encoder_attention_heads, dropout=config.attention_dropout)
- self.self_attn_layer_norm = LayerNorm(self.embed_dim)
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
- self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
- self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = LayerNorm(self.embed_dim)
-
- def forward(self, x, encoder_padding_mask, layer_head_mask, output_attentions=False):
- """
- Args:
- x (`torch.Tensor`): input to the layer of shape *(seq_len, batch, embed_dim)*
- encoder_padding_mask (`torch.ByteTensor`): binary ByteTensor of shape
- *(batch, src_len)* where padding elements are indicated by `1`.
- for t_tgt, t_src is excluded (or masked out), =0 means it is
- included in attention
- layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
- *(config.encoder_attention_heads,)*.
-
- Returns:
- encoded output of shape *(seq_len, batch, embed_dim)*
- """
- residual = x
- x, attn_weights = self.self_attn(
- query=x,
- key=x,
- key_padding_mask=encoder_padding_mask,
- layer_head_mask=layer_head_mask,
- output_attentions=output_attentions,
- )
- x = nn.functional.dropout(x, p=self.dropout, training=self.training)
- x = residual + x
- x = self.self_attn_layer_norm(x)
-
- residual = x
- x = self.activation_fn(self.fc1(x))
- x = nn.functional.dropout(x, p=self.activation_dropout, training=self.training)
- x = self.fc2(x)
- x = nn.functional.dropout(x, p=self.dropout, training=self.training)
- x = residual + x
- x = self.final_layer_norm(x)
- return x, attn_weights
-
-
-class FSMTEncoder(nn.Module):
- """
- Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a [`EncoderLayer`].
-
- Args:
- config: FSMTConfig
- """
-
- def __init__(self, config: FSMTConfig, embed_tokens):
- super().__init__()
- self.dropout = config.dropout
- self.layerdrop = config.encoder_layerdrop
- self.padding_idx = embed_tokens.padding_idx
- self.embed_tokens = embed_tokens
- embed_dim = embed_tokens.embedding_dim
- self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
- self.embed_positions = SinusoidalPositionalEmbedding(
- config.max_position_embeddings + self.padding_idx + 1, embed_dim, self.padding_idx
- )
- self.layers = nn.ModuleList([EncoderLayer(config) for _ in range(config.encoder_layers)]) # type: List[EncoderLayer]
-
- def forward(
- self,
- input_ids: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- inputs_embeds: torch.Tensor = None,
- head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- """
- Args:
- input_ids (`torch.LongTensor`): tokens in the source language of shape
- *(batch, src_len)*
- attention_mask (`torch.LongTensor`): indicating which indices are padding tokens
- inputs_embeds (`torch.FloatTensor`):
- embedding vectors of shape *(batch, src_len, embed_dim)*
- head_mask (`torch.Tensor` of shape `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- Returns:
- BaseModelOutput or Tuple comprised of:
-
- - **x** (`torch.Tensor`): the last encoder layer's output of shape *(src_len, batch, embed_dim)*
- - **encoder_states** (`Tuple(torch.FloatTensor`)): all intermediate hidden states of shape *(src_len,
- batch, embed_dim)*. Only populated if *output_hidden_states:* is True.
- - **all_attentions** (`Tuple(torch.FloatTensor`)): Attention weights for each layer.
- During training might not be of length n_layers because of layer dropout.
- """
- # check attention mask and invert
- if attention_mask is not None:
- attention_mask = invert_mask(attention_mask)
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
- embed_pos = self.embed_positions(input_ids)
- elif inputs_embeds is not None:
- inputs_embeds = inputs_embeds * self.embed_scale
-
- # We assume zeros hidden states correspond to padding tokens
- # and create `position_ids` where inputs_embeds[:, :, 0] == 0
- position_ids = inputs_embeds[:, :, 0].masked_fill(
- inputs_embeds[:, :, 0].eq(0), self.embed_positions.padding_idx
- )
-
- embed_pos = self.embed_positions(position_ids)
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- x = inputs_embeds + embed_pos
- x = nn.functional.dropout(x, p=self.dropout, training=self.training)
-
- # B x T x C -> T x B x C
- x = x.transpose(0, 1)
-
- encoder_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
- # check if head_mask has a correct number of layers specified if desired
- if head_mask is not None:
- assert head_mask.size()[0] == (
- len(self.layers)
- ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
- for idx, encoder_layer in enumerate(self.layers):
- if output_hidden_states:
- x = x.transpose(0, 1) # T x B x C -> B x T x C
- encoder_states += (x,)
- x = x.transpose(0, 1) # B x T x C -> T x B x C
- # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
- dropout_probability = torch.rand([])
- if self.training and (dropout_probability < self.layerdrop): # skip the layer
- attn = None
- else:
- x, attn = encoder_layer(
- x,
- attention_mask,
- layer_head_mask=(head_mask[idx] if head_mask is not None else None),
- output_attentions=output_attentions,
- )
-
- if output_attentions:
- all_attentions = all_attentions + (attn,)
-
- # T x B x C -> B x T x C
- x = x.transpose(0, 1)
-
- if output_hidden_states:
- encoder_states += (x,)
-
- if not return_dict:
- return tuple(v for v in [x, encoder_states, all_attentions] if v is not None)
- return BaseModelOutput(last_hidden_state=x, hidden_states=encoder_states, attentions=all_attentions)
-
-
-class DecoderLayer(nn.Module):
- def __init__(self, config: FSMTConfig):
- super().__init__()
- self.embed_dim = config.d_model
-
- self.self_attn = Attention(
- embed_dim=self.embed_dim,
- num_heads=config.decoder_attention_heads,
- dropout=config.attention_dropout,
- )
- self.dropout = config.dropout
- self.activation_fn = ACT2FN[config.activation_function]
- self.activation_dropout = config.activation_dropout
-
- self.self_attn_layer_norm = LayerNorm(self.embed_dim)
- self.encoder_attn = Attention(
- self.embed_dim,
- config.decoder_attention_heads,
- dropout=config.attention_dropout,
- encoder_decoder_attention=True,
- )
- self.encoder_attn_layer_norm = LayerNorm(self.embed_dim)
- self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
- self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
- self.final_layer_norm = LayerNorm(self.embed_dim)
-
- def forward(
- self,
- x,
- encoder_hidden_states,
- encoder_attn_mask=None,
- layer_state=None,
- causal_mask=None,
- layer_head_mask=None,
- cross_attn_layer_head_mask=None,
- decoder_padding_mask=None,
- output_attentions=False,
- ):
- residual = x
-
- if layer_state is None:
- layer_state = {}
-
- # Self Attention
- x, self_attn_weights = self.self_attn(
- query=x,
- key=x,
- layer_state=layer_state, # adds keys to layer state
- key_padding_mask=decoder_padding_mask,
- attn_mask=causal_mask,
- layer_head_mask=layer_head_mask,
- output_attentions=output_attentions,
- )
- x = nn.functional.dropout(x, p=self.dropout, training=self.training)
- x = residual + x
- x = self.self_attn_layer_norm(x)
-
- # Cross attention
- residual = x
- assert self.encoder_attn.cache_key != self.self_attn.cache_key
- x, cross_attn_weights = self.encoder_attn(
- query=x,
- key=encoder_hidden_states,
- key_padding_mask=encoder_attn_mask,
- layer_state=layer_state, # mutates layer state
- layer_head_mask=cross_attn_layer_head_mask,
- output_attentions=output_attentions,
- )
- x = nn.functional.dropout(x, p=self.dropout, training=self.training)
- x = residual + x
- x = self.encoder_attn_layer_norm(x)
-
- # Fully Connected
- residual = x
- x = self.activation_fn(self.fc1(x))
- x = nn.functional.dropout(x, p=self.activation_dropout, training=self.training)
- x = self.fc2(x)
- x = nn.functional.dropout(x, p=self.dropout, training=self.training)
- x = residual + x
- x = self.final_layer_norm(x)
- return (
- x,
- self_attn_weights,
- layer_state,
- cross_attn_weights,
- ) # layer_state = cache for decoding
-
-
-class FSMTDecoder(nn.Module):
- """
- Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DecoderLayer`]
-
- Args:
- config: FSMTConfig
- embed_tokens (nn.Embedding): output embedding
- """
-
- def __init__(self, config: FSMTConfig, embed_tokens: nn.Embedding):
- super().__init__()
- self.dropout = config.dropout
- self.layerdrop = config.decoder_layerdrop
- self.padding_idx = embed_tokens.padding_idx
- self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
- self.embed_tokens = embed_tokens
- embed_dim = embed_tokens.embedding_dim
- self.embed_positions = SinusoidalPositionalEmbedding(
- config.max_position_embeddings + self.padding_idx + 1, embed_dim, self.padding_idx
- )
- self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.decoder_layers)]) # type: List[DecoderLayer]
-
- if is_deepspeed_zero3_enabled():
- import deepspeed
-
- with deepspeed.zero.GatheredParameters(self.embed_tokens.weight, modifier_rank=None):
- embed_tokens_weight_shape = self.embed_tokens.weight.shape
- else:
- embed_tokens_weight_shape = self.embed_tokens.weight.shape
- self.output_projection = nn.Linear(embed_tokens_weight_shape[1], embed_tokens_weight_shape[0], bias=False)
- self.output_projection.weight = self.embed_tokens.weight
-
- def forward(
- self,
- input_ids: torch.Tensor,
- encoder_hidden_states: torch.Tensor,
- encoder_padding_mask: torch.Tensor,
- decoder_padding_mask: torch.Tensor,
- decoder_causal_mask: torch.Tensor,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- cross_attn_head_mask: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- use_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- """
- Includes several features from "Jointly Learning to Align and Translate with Transformer Models" (Garg et al.,
- EMNLP 2019).
-
- Args:
- input_ids (`torch.LongTensor` of shape `(batch, tgt_len)`):
- previous decoder outputs for teacher forcing
- encoder_hidden_states: output from the encoder, used for
- encoder-side attention
- encoder_padding_mask: for ignoring pad tokens
- past_key_values (dict or None): dictionary used for storing state during generation
- head_mask (`torch.Tensor` of shape `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- cross_attn_head_mask (`torch.Tensor` of shape `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- Returns:
- BaseModelOutputWithPast or tuple:
-
- - the decoder's features of shape *(batch, tgt_len, embed_dim)*
- - the cache
- - hidden states
- - attentions
- """
- # check attention mask and invert
- if encoder_padding_mask is not None:
- encoder_padding_mask = invert_mask(encoder_padding_mask)
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
- elif input_ids is not None:
- # embed positions
- positions = self.embed_positions(input_ids)
- if use_cache:
- input_ids = input_ids[:, -1:]
- positions = positions[:, -1:] # happens after we embed them
- x = self.embed_tokens(input_ids) * self.embed_scale
- elif inputs_embeds is not None:
- # We assume zeros hidden states correspond to padding tokens
- # and create `position_ids` where inputs_embeds[:, :, 0] == 0
- position_ids = inputs_embeds[:, :, 0].masked_fill(
- inputs_embeds[:, :, 0].eq(0), self.embed_positions.padding_idx
- )
- positions = self.embed_positions(position_ids)
- x = inputs_embeds * self.embed_scale
- else:
- raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
-
- x += positions
- x = nn.functional.dropout(x, p=self.dropout, training=self.training)
-
- # Convert to FSMT output format: (BS, seq_len, model_dim) -> (seq_len, BS, model_dim)
- x = x.transpose(0, 1)
- encoder_hidden_states = encoder_hidden_states.transpose(0, 1)
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- all_cross_attns = () if output_attentions else None
- next_decoder_cache = []
-
- # check if head_mask has a correct number of layers specified if desired
- for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
- if attn_mask is not None:
- assert attn_mask.size()[0] == (len(self.layers)), (
- f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
- f" {head_mask.size()[0]}."
- )
- for idx, decoder_layer in enumerate(self.layers):
- # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
- if output_hidden_states:
- x = x.transpose(0, 1)
- all_hidden_states += (x,)
- x = x.transpose(0, 1)
- if self.training:
- dropout_probability = torch.rand([])
- if dropout_probability < self.layerdrop:
- continue
-
- layer_state = past_key_values[idx] if past_key_values is not None else None
-
- x, layer_self_attn, layer_past, layer_cross_attn = decoder_layer(
- x,
- encoder_hidden_states,
- encoder_attn_mask=encoder_padding_mask,
- decoder_padding_mask=decoder_padding_mask,
- layer_state=layer_state,
- causal_mask=decoder_causal_mask,
- layer_head_mask=(head_mask[idx] if head_mask is not None else None),
- cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None),
- output_attentions=output_attentions,
- )
-
- if use_cache:
- next_decoder_cache.append(layer_past.copy())
-
- if output_attentions:
- all_self_attns += (layer_self_attn,)
- all_cross_attns += (layer_cross_attn,)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- x = x.transpose(0, 1)
- all_hidden_states += (x,)
- x = x.transpose(0, 1)
-
- # Convert to standard output format: (seq_len, BS, model_dim) -> (BS, seq_len, model_dim)
- x = x.transpose(0, 1)
- encoder_hidden_states = encoder_hidden_states.transpose(0, 1)
-
- x = self.output_projection(x)
-
- next_cache = next_decoder_cache if use_cache else None
-
- if not return_dict:
- return tuple(
- v for v in [x, next_cache, all_hidden_states, all_self_attns, all_cross_attns] if v is not None
- )
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=x,
- past_key_values=next_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- cross_attentions=all_cross_attns,
- )
-
-
-def _reorder_buffer(attn_cache, new_order):
- for k, input_buffer_k in attn_cache.items():
- if input_buffer_k is not None:
- attn_cache[k] = input_buffer_k.index_select(0, new_order)
- return attn_cache
-
-
-class Attention(nn.Module):
- """Multi-headed attention from 'Attention Is All You Need' paper"""
-
- def __init__(
- self,
- embed_dim,
- num_heads,
- dropout=0.0,
- bias=True,
- encoder_decoder_attention=False, # otherwise self_attention
- ):
- super().__init__()
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.dropout = dropout
- self.head_dim = embed_dim // num_heads
- assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
- self.scaling = self.head_dim**-0.5
-
- self.encoder_decoder_attention = encoder_decoder_attention
- self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.cache_key = "encoder_decoder" if self.encoder_decoder_attention else "self"
-
- def _shape(self, tensor, seq_len, bsz):
- return tensor.contiguous().view(seq_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
-
- def forward(
- self,
- query,
- key: Optional[Tensor],
- key_padding_mask: Optional[Tensor] = None,
- layer_state: Optional[Dict[str, Optional[Tensor]]] = None,
- attn_mask: Optional[Tensor] = None,
- layer_head_mask: Optional[Tensor] = None,
- output_attentions=False,
- ) -> Tuple[Tensor, Optional[Tensor]]:
- """Input shape: Time(SeqLen) x Batch x Channel"""
- static_kv: bool = self.encoder_decoder_attention
- tgt_len, bsz, embed_dim = query.size()
- assert embed_dim == self.embed_dim
- assert list(query.size()) == [tgt_len, bsz, embed_dim]
- # get here for encoder decoder cause of static_kv
- if layer_state is not None: # reuse k,v and encoder_padding_mask
- saved_state = layer_state.get(self.cache_key, {})
- if "prev_key" in saved_state and static_kv:
- # previous time steps are cached - no need to recompute key and value if they are static
- key = None
- else:
- saved_state = None
- layer_state = {}
-
- q = self.q_proj(query) * self.scaling
- if static_kv:
- if key is None:
- k = v = None
- else:
- k = self.k_proj(key)
- v = self.v_proj(key)
- else:
- k = self.k_proj(query)
- v = self.v_proj(query)
-
- q = self._shape(q, tgt_len, bsz)
- if k is not None:
- k = self._shape(k, -1, bsz)
- if v is not None:
- v = self._shape(v, -1, bsz)
-
- if saved_state is not None:
- k, v, key_padding_mask = self._use_saved_state(k, v, saved_state, key_padding_mask, static_kv, bsz)
-
- # Update cache
- layer_state[self.cache_key] = {
- "prev_key": k.view(bsz, self.num_heads, -1, self.head_dim),
- "prev_value": v.view(bsz, self.num_heads, -1, self.head_dim),
- "prev_key_padding_mask": key_padding_mask if not static_kv else None,
- }
-
- assert k is not None
- src_len = k.size(1)
- attn_weights = torch.bmm(q, k.transpose(1, 2))
- assert attn_weights.size() == (bsz * self.num_heads, tgt_len, src_len)
-
- if attn_mask is not None:
- attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attn_mask
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- # This is part of a workaround to get around fork/join parallelism not supporting Optional types.
- if key_padding_mask is not None and key_padding_mask.dim() == 0:
- key_padding_mask = None
- assert key_padding_mask is None or key_padding_mask.size()[:2] == (
- bsz,
- src_len,
- )
-
- if key_padding_mask is not None: # don't attend to padding symbols
- attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
- reshaped = key_padding_mask.unsqueeze(1).unsqueeze(2)
- attn_weights = attn_weights.masked_fill(reshaped, torch.finfo(attn_weights.dtype).min)
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if layer_head_mask is not None:
- assert layer_head_mask.size() == (
- self.num_heads,
- ), f"Head mask for a single layer should be of size {(self.num_heads,)}, but is {layer_head_mask.size()}"
- attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- if output_attentions:
- # make sure that attn_weights are included in graph
- attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
- attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(
- attn_weights,
- p=self.dropout,
- training=self.training,
- )
-
- assert v is not None
- attn_output = torch.bmm(attn_probs, v)
- assert attn_output.size() == (bsz * self.num_heads, tgt_len, self.head_dim)
- attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped
-
- def _use_saved_state(self, k, v, saved_state, key_padding_mask, static_kv, bsz):
- # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
- if "prev_key" in saved_state:
- _prev_key = saved_state["prev_key"]
- assert _prev_key is not None
- prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
- if static_kv:
- k = prev_key
- else:
- assert k is not None
- k = torch.cat([prev_key, k], dim=1)
- if "prev_value" in saved_state:
- _prev_value = saved_state["prev_value"]
- assert _prev_value is not None
- prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
- if static_kv:
- v = prev_value
- else:
- assert v is not None
- v = torch.cat([prev_value, v], dim=1)
- assert k is not None and v is not None
- prev_key_padding_mask: Optional[Tensor] = saved_state.get("prev_key_padding_mask", None)
- if prev_key_padding_mask is not None:
- if static_kv:
- new_key_padding_mask = prev_key_padding_mask
- else:
- new_key_padding_mask = torch.cat([prev_key_padding_mask, key_padding_mask], dim=1)
- else:
- new_key_padding_mask = key_padding_mask
- return k, v, new_key_padding_mask
-
-
-def fill_with_neg_inf(t):
- """FP16-compatible function that fills a input_ids with -inf."""
- return t.float().fill_(torch.finfo(t.dtype).min).type_as(t)
-
-
-# Public API
-def _get_shape(t):
- return getattr(t, "shape", None)
-
-
-@add_start_docstrings(
- "The bare FSMT Model outputting raw hidden-states without any specific head on top.",
- FSMT_START_DOCSTRING,
-)
-class FSMTModel(PretrainedFSMTModel):
- _tied_weights_keys = ["decoder.embed_tokens.weight", "decoder.output_projection.weight"]
-
- def __init__(self, config: FSMTConfig):
- super().__init__(config)
-
- padding_idx = config.pad_token_id
- encoder_embed_tokens = nn.Embedding(config.src_vocab_size, config.d_model, padding_idx)
- decoder_embed_tokens = nn.Embedding(config.tgt_vocab_size, config.d_model, padding_idx)
-
- self.encoder = FSMTEncoder(config, encoder_embed_tokens)
- self.decoder = FSMTDecoder(config, decoder_embed_tokens)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_encoder(self):
- return self.encoder
-
- def get_decoder(self):
- return self.decoder
-
- def _tie_weights(self):
- if self.config.tie_word_embeddings:
- self._tie_or_clone_weights(self.decoder.embed_tokens, self.get_input_embeddings())
- self._tie_or_clone_weights(self.decoder.output_projection, self.get_input_embeddings())
-
- @add_start_docstrings_to_model_forward(FSMT_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=Seq2SeqModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: torch.LongTensor,
- attention_mask: Optional[torch.Tensor] = None,
- decoder_input_ids: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.BoolTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- decoder_head_mask: Optional[torch.Tensor] = None,
- cross_attn_head_mask: Optional[torch.Tensor] = None,
- encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
- past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], Seq2SeqModelOutput]:
- if decoder_input_ids is None:
- use_cache = False
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # make masks if user doesn't supply
- if not use_cache and input_ids is not None:
- decoder_input_ids, decoder_padding_mask, causal_mask = _prepare_fsmt_decoder_inputs(
- self.config,
- input_ids,
- decoder_input_ids=decoder_input_ids,
- decoder_padding_mask=decoder_attention_mask,
- causal_mask_dtype=self.decoder.embed_tokens.weight.dtype,
- )
- else:
- decoder_padding_mask, causal_mask = None, None
-
- if decoder_input_ids is None and decoder_inputs_embeds is None:
- raise ValueError("Make sure that `decoder_input_ids` or `decoder_inputs_embeds` are passed.")
-
- if encoder_outputs is None:
- encoder_outputs = self.encoder(
- input_ids=input_ids,
- attention_mask=attention_mask,
- inputs_embeds=inputs_embeds,
- head_mask=head_mask,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=False
- elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
- encoder_outputs = BaseModelOutput(
- last_hidden_state=encoder_outputs[0],
- hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
- attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
- )
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- decoder_outputs = self.decoder(
- decoder_input_ids,
- encoder_outputs[0],
- attention_mask,
- decoder_padding_mask,
- decoder_causal_mask=causal_mask,
- inputs_embeds=decoder_inputs_embeds,
- head_mask=decoder_head_mask,
- cross_attn_head_mask=cross_attn_head_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- return decoder_outputs + encoder_outputs
-
- return Seq2SeqModelOutput(
- last_hidden_state=decoder_outputs.last_hidden_state,
- past_key_values=decoder_outputs.past_key_values,
- decoder_hidden_states=decoder_outputs.hidden_states,
- decoder_attentions=decoder_outputs.attentions,
- cross_attentions=decoder_outputs.cross_attentions,
- encoder_last_hidden_state=encoder_outputs.last_hidden_state,
- encoder_hidden_states=encoder_outputs.hidden_states,
- encoder_attentions=encoder_outputs.attentions,
- )
-
- def get_input_embeddings(self):
- return self.encoder.embed_tokens
-
- def set_input_embeddings(self, value):
- self.encoder.embed_tokens = value
-
- def get_output_embeddings(self):
- return self.decoder.embed_tokens
-
- def set_output_embeddings(self, value):
- self.decoder.embed_tokens = value
-
-
-@add_start_docstrings(
- "The FSMT Model with a language modeling head. Can be used for summarization.", FSMT_START_DOCSTRING
-)
-class FSMTForConditionalGeneration(PretrainedFSMTModel):
- base_model_prefix = "model"
- _tied_weights_keys = ["decoder.embed_tokens.weight", "decoder.output_projection.weight"]
-
- def __init__(self, config: FSMTConfig):
- super().__init__(config)
- base_model = FSMTModel(config)
- self.model = base_model
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FSMT_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
- @add_end_docstrings(FSMT_GENERATION_EXAMPLE)
- def forward(
- self,
- input_ids: torch.LongTensor,
- attention_mask: Optional[torch.Tensor] = None,
- decoder_input_ids: Optional[torch.LongTensor] = None,
- decoder_attention_mask: Optional[torch.BoolTensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- decoder_head_mask: Optional[torch.Tensor] = None,
- cross_attn_head_mask: Optional[torch.Tensor] = None,
- encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
- past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- decoder_inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- Returns:
-
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if labels is not None:
- use_cache = False
-
- outputs = self.model(
- input_ids,
- inputs_embeds=inputs_embeds,
- attention_mask=attention_mask,
- decoder_input_ids=decoder_input_ids,
- decoder_inputs_embeds=decoder_inputs_embeds,
- encoder_outputs=encoder_outputs,
- decoder_attention_mask=decoder_attention_mask,
- head_mask=head_mask,
- decoder_head_mask=decoder_head_mask,
- cross_attn_head_mask=cross_attn_head_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- lm_logits = outputs[0]
-
- masked_lm_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- # TODO(SS): do we need to ignore pad tokens in labels?
- masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.tgt_vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (lm_logits,) + outputs[1:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return Seq2SeqLMOutput(
- loss=masked_lm_loss,
- logits=lm_logits,
- past_key_values=outputs.past_key_values,
- decoder_hidden_states=outputs.decoder_hidden_states,
- decoder_attentions=outputs.decoder_attentions,
- cross_attentions=outputs.cross_attentions,
- encoder_last_hidden_state=outputs.encoder_last_hidden_state,
- encoder_hidden_states=outputs.encoder_hidden_states,
- encoder_attentions=outputs.encoder_attentions,
- )
-
- def prepare_inputs_for_generation(
- self,
- decoder_input_ids,
- past_key_values=None,
- attention_mask=None,
- head_mask=None,
- decoder_head_mask=None,
- cross_attn_head_mask=None,
- use_cache=None,
- encoder_outputs=None,
- **kwargs,
- ):
- return {
- "input_ids": None, # encoder_outputs is defined. input_ids not needed
- "encoder_outputs": encoder_outputs,
- "past_key_values": past_key_values,
- "decoder_input_ids": decoder_input_ids,
- "attention_mask": attention_mask,
- "head_mask": head_mask,
- "decoder_head_mask": decoder_head_mask,
- "cross_attn_head_mask": cross_attn_head_mask,
- "use_cache": use_cache, # change this to avoid caching (presumably for debugging)
- }
-
- def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
- return shift_tokens_right(labels, self.config.pad_token_id)
-
- @staticmethod
- def _reorder_cache(past_key_values, beam_idx):
- reordered_past = []
- for layer_past in past_key_values:
- # get the correct batch idx from decoder layer's batch dim for cross and self-attn
- layer_past_new = {
- attn_key: _reorder_buffer(attn_cache, beam_idx) for attn_key, attn_cache in layer_past.items()
- }
- reordered_past.append(layer_past_new)
- return reordered_past
-
- def get_encoder(self):
- return self.model.encoder
-
- def get_decoder(self):
- return self.model.decoder
-
- def get_output_embeddings(self):
- return self.model.decoder.embed_tokens
-
- def set_output_embeddings(self, value):
- self.model.decoder.embed_tokens = value
-
-
-class SinusoidalPositionalEmbedding(nn.Embedding):
- """
- This module produces sinusoidal positional embeddings of any length.
-
- We don't want to save the weight of this embedding since it's not trained (deterministic) and it can be huge.
-
- Padding symbols are ignored.
-
- These embeddings get automatically extended in forward if more positions is needed.
- """
-
- def __init__(self, num_positions, embedding_dim, padding_idx):
- self.make_weight(num_positions, embedding_dim, padding_idx)
-
- def make_weight(self, num_positions, embedding_dim, padding_idx):
- weight = self.get_embedding(num_positions, embedding_dim, padding_idx)
- if not hasattr(self, "weight"):
- # in ___init__
- super().__init__(num_positions, embedding_dim, padding_idx, _weight=weight)
- else:
- # in forward put the weights on the correct dtype and device of the param
- weight = weight.to(dtype=self.weight.dtype, device=self.weight.device)
- self.weight = nn.Parameter(weight)
- self.weight.detach_()
- self.weight.requires_grad = False
-
- @staticmethod
- def get_embedding(num_embeddings, embedding_dim, padding_idx):
- """
- Build sinusoidal embeddings.
-
- This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of
- "Attention Is All You Need".
- """
- half_dim = embedding_dim // 2
- emb = math.log(10000) / (half_dim - 1)
- emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb)
- emb = torch.arange(num_embeddings, dtype=torch.int64).float().unsqueeze(1) * emb.unsqueeze(0)
- emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
- if embedding_dim % 2 == 1:
- # zero pad
- emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
- if padding_idx is not None:
- emb[padding_idx, :] = 0
- return emb
-
- @staticmethod
- def make_positions(tensor, padding_idx: int):
- """
- Replace non-padding symbols with their position numbers.
-
- Position numbers begin at padding_idx+1. Padding symbols are ignored.
- """
- # The series of casts and type-conversions here are carefully
- # balanced to both work with ONNX export and XLA. In particular XLA
- # prefers ints, cumsum defaults to output longs, and ONNX doesn't know
- # how to handle the dtype kwarg in cumsum.
- mask = tensor.ne(padding_idx).int()
- return (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + padding_idx
-
- def forward(
- self,
- input,
- incremental_state: Optional[Any] = None,
- timestep: Optional[Tensor] = None,
- ):
- """Input is expected to be of size [bsz x seqlen]."""
- bsz, seq_len = input.shape[:2]
- max_pos = self.padding_idx + 1 + seq_len
- if max_pos > self.weight.size(0):
- # expand embeddings if needed
- self.make_weight(max_pos, self.embedding_dim, self.padding_idx)
- positions = self.make_positions(input, self.padding_idx)
- return super().forward(positions)
diff --git a/transformers/models/fsmt/tokenization_fsmt.py b/transformers/models/fsmt/tokenization_fsmt.py
deleted file mode 100644
index 8b0be1f8be24987259aaee01d3165aa03c9218a9..0000000000000000000000000000000000000000
--- a/transformers/models/fsmt/tokenization_fsmt.py
+++ /dev/null
@@ -1,519 +0,0 @@
-# coding=utf-8
-# Copyright 2019 The Open AI Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for FSMT."""
-
-
-import json
-import os
-import re
-import unicodedata
-from typing import Dict, List, Optional, Tuple
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {
- "src_vocab_file": "vocab-src.json",
- "tgt_vocab_file": "vocab-tgt.json",
- "merges_file": "merges.txt",
-}
-
-
-def get_pairs(word):
- """
- Return set of symbol pairs in a word. word is represented as tuple of symbols (symbols being variable-length
- strings)
- """
- pairs = set()
- prev_char = word[0]
- for char in word[1:]:
- pairs.add((prev_char, char))
- prev_char = char
- return pairs
-
-
-def replace_unicode_punct(text):
- """
- Port of https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/replace-unicode-punctuation.perl
- """
- text = text.replace("，", ",")
- text = re.sub(r"。\s*", ". ", text)
- text = text.replace("、", ",")
- text = text.replace("”", '"')
- text = text.replace("“", '"')
- text = text.replace("∶", ":")
- text = text.replace("：", ":")
- text = text.replace("？", "?")
- text = text.replace("《", '"')
- text = text.replace("》", '"')
- text = text.replace("）", ")")
- text = text.replace("！", "!")
- text = text.replace("（", "(")
- text = text.replace("；", ";")
- text = text.replace("１", "1")
- text = text.replace("」", '"')
- text = text.replace("「", '"')
- text = text.replace("０", "0")
- text = text.replace("３", "3")
- text = text.replace("２", "2")
- text = text.replace("５", "5")
- text = text.replace("６", "6")
- text = text.replace("９", "9")
- text = text.replace("７", "7")
- text = text.replace("８", "8")
- text = text.replace("４", "4")
- text = re.sub(r"．\s*", ". ", text)
- text = text.replace("～", "~")
- text = text.replace("’", "'")
- text = text.replace("…", "...")
- text = text.replace("━", "-")
- text = text.replace("〈", "<")
- text = text.replace("〉", ">")
- text = text.replace("【", "[")
- text = text.replace("】", "]")
- text = text.replace("％", "%")
- return text
-
-
-def remove_non_printing_char(text):
- """
- Port of https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/remove-non-printing-char.perl
- """
- output = []
- for char in text:
- cat = unicodedata.category(char)
- if cat.startswith("C"):
- continue
- output.append(char)
- return "".join(output)
-
-
-# Porting notes:
-# this one is modeled after XLMTokenizer
-#
-# added:
-# - src_vocab_file,
-# - tgt_vocab_file,
-# - langs,
-
-
-class FSMTTokenizer(PreTrainedTokenizer):
- """
- Construct an FAIRSEQ Transformer tokenizer. Based on Byte-Pair Encoding. The tokenization process is the following:
-
- - Moses preprocessing and tokenization.
- - Normalizing all inputs text.
- - The arguments `special_tokens` and the function `set_special_tokens`, can be used to add additional symbols (like
- "__classify__") to a vocabulary.
- - The argument `langs` defines a pair of languages.
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- langs (`List[str]`, *optional*):
- A list of two languages to translate from and to, for instance `["en", "ru"]`.
- src_vocab_file (`str`, *optional*):
- File containing the vocabulary for the source language.
- tgt_vocab_file (`st`, *optional*):
- File containing the vocabulary for the target language.
- merges_file (`str`, *optional*):
- File containing the merges.
- do_lower_case (`bool`, *optional*, defaults to `False`):
- Whether or not to lowercase the input when tokenizing.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
-
-
-
- When building a sequence using special tokens, this is not the token that is used for the beginning of
- sequence. The token used is the `cls_token`.
-
-
-
- sep_token (`str`, *optional*, defaults to `""`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
-
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- langs=None,
- src_vocab_file=None,
- tgt_vocab_file=None,
- merges_file=None,
- do_lower_case=False,
- unk_token="",
- bos_token="",
- sep_token="",
- pad_token="",
- **kwargs,
- ):
- try:
- import sacremoses
- except ImportError:
- raise ImportError(
- "You need to install sacremoses to use XLMTokenizer. "
- "See https://pypi.org/project/sacremoses/ for installation."
- )
-
- self.sm = sacremoses
-
- self.src_vocab_file = src_vocab_file
- self.tgt_vocab_file = tgt_vocab_file
- self.merges_file = merges_file
- self.do_lower_case = do_lower_case
-
- # cache of sm.MosesPunctNormalizer instance
- self.cache_moses_punct_normalizer = {}
- # cache of sm.MosesTokenizer instance
- self.cache_moses_tokenizer = {}
- self.cache_moses_detokenizer = {}
-
- if langs and len(langs) == 2:
- self.src_lang, self.tgt_lang = langs
- else:
- raise ValueError(
- f"arg `langs` needs to be a list of 2 langs, e.g. ['en', 'ru'], but got {langs}. "
- "Usually that means that tokenizer can't find a mapping for the given model path "
- "in PRETRAINED_VOCAB_FILES_MAP, and other maps of this tokenizer."
- )
-
- with open(src_vocab_file, encoding="utf-8") as src_vocab_handle:
- self.encoder = json.load(src_vocab_handle)
- with open(tgt_vocab_file, encoding="utf-8") as tgt_vocab_handle:
- tgt_vocab = json.load(tgt_vocab_handle)
- self.decoder = {v: k for k, v in tgt_vocab.items()}
- with open(merges_file, encoding="utf-8") as merges_handle:
- merges = merges_handle.read().split("\n")[:-1]
- merges = [tuple(merge.split()[:2]) for merge in merges]
- self.bpe_ranks = dict(zip(merges, range(len(merges))))
- self.cache = {}
- super().__init__(
- langs=langs,
- src_vocab_file=src_vocab_file,
- tgt_vocab_file=tgt_vocab_file,
- merges_file=merges_file,
- do_lower_case=do_lower_case,
- unk_token=unk_token,
- bos_token=bos_token,
- sep_token=sep_token,
- pad_token=pad_token,
- **kwargs,
- )
-
- # hack override
- def get_vocab(self) -> Dict[str, int]:
- return self.get_src_vocab()
-
- # hack override
- @property
- def vocab_size(self) -> int:
- return self.src_vocab_size
-
- def moses_punct_norm(self, text, lang):
- if lang not in self.cache_moses_punct_normalizer:
- punct_normalizer = self.sm.MosesPunctNormalizer(lang=lang)
- self.cache_moses_punct_normalizer[lang] = punct_normalizer
- return self.cache_moses_punct_normalizer[lang].normalize(text)
-
- def moses_tokenize(self, text, lang):
- if lang not in self.cache_moses_tokenizer:
- moses_tokenizer = self.sm.MosesTokenizer(lang=lang)
- self.cache_moses_tokenizer[lang] = moses_tokenizer
- return self.cache_moses_tokenizer[lang].tokenize(
- text, aggressive_dash_splits=True, return_str=False, escape=True
- )
-
- def moses_detokenize(self, tokens, lang):
- if lang not in self.cache_moses_detokenizer:
- moses_detokenizer = self.sm.MosesDetokenizer(lang=lang)
- self.cache_moses_detokenizer[lang] = moses_detokenizer
- return self.cache_moses_detokenizer[lang].detokenize(tokens)
-
- def moses_pipeline(self, text, lang):
- text = replace_unicode_punct(text)
- text = self.moses_punct_norm(text, lang)
- text = remove_non_printing_char(text)
- return text
-
- @property
- def src_vocab_size(self):
- return len(self.encoder)
-
- @property
- def tgt_vocab_size(self):
- return len(self.decoder)
-
- def get_src_vocab(self):
- return dict(self.encoder, **self.added_tokens_encoder)
-
- def get_tgt_vocab(self):
- return dict(self.decoder, **self.added_tokens_decoder)
-
- def bpe(self, token):
- word = tuple(token[:-1]) + (token[-1] + "",)
- if token in self.cache:
- return self.cache[token]
- pairs = get_pairs(word)
-
- if not pairs:
- return token + ""
-
- while True:
- bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
- if bigram not in self.bpe_ranks:
- break
- first, second = bigram
- new_word = []
- i = 0
- while i < len(word):
- try:
- j = word.index(first, i)
- except ValueError:
- new_word.extend(word[i:])
- break
- else:
- new_word.extend(word[i:j])
- i = j
-
- if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
- new_word.append(first + second)
- i += 2
- else:
- new_word.append(word[i])
- i += 1
- new_word = tuple(new_word)
- word = new_word
- if len(word) == 1:
- break
- else:
- pairs = get_pairs(word)
- word = " ".join(word)
- if word == "\n ":
- word = "\n"
- self.cache[token] = word
- return word
-
- def _tokenize(self, text, lang="en", bypass_tokenizer=False):
- """
- Tokenize a string given language code using Moses.
-
- Details of tokenization:
-
- - [sacremoses](https://github.com/alvations/sacremoses): port of Moses
- - Install with `pip install sacremoses`
-
- Args:
- - lang: ISO language code (default = 'en') (string). Languages should belong of the model supported
- languages. However, we don't enforce it.
- - bypass_tokenizer: Allow users to preprocess and tokenize the sentences externally (default = False)
- (bool). If True, we only apply BPE.
-
- Returns:
- List of tokens.
- """
- # ignore `lang` which is currently isn't explicitly passed in tokenization_utils.py and always results in lang=en
- # if lang != self.src_lang:
- # raise ValueError(f"Expected lang={self.src_lang}, but got {lang}")
- lang = self.src_lang
-
- if self.do_lower_case:
- text = text.lower()
-
- if bypass_tokenizer:
- text = text.split()
- else:
- text = self.moses_pipeline(text, lang=lang)
- text = self.moses_tokenize(text, lang=lang)
-
- split_tokens = []
- for token in text:
- if token:
- split_tokens.extend(list(self.bpe(token).split(" ")))
-
- return split_tokens
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index, self.unk_token)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
-
- # remove BPE
- tokens = [t.replace(" ", "").replace("", " ") for t in tokens]
- tokens = "".join(tokens).split()
- # detokenize
- text = self.moses_detokenize(tokens, self.tgt_lang)
- return text
-
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A FAIRSEQ Transformer sequence has the following format:
-
- - single sequence: ` X `
- - pair of sequences: ` A B `
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- sep = [self.sep_token_id]
-
- # no bos used in fairseq
- if token_ids_1 is None:
- return token_ids_0 + sep
- return token_ids_0 + sep + token_ids_1 + sep
-
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
- # no bos used in fairseq
- if token_ids_1 is not None:
- return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return ([0] * len(token_ids_0)) + [1]
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A FAIRSEQ
- Transformer sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
-
- Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An
- FAIRSEQ_TRANSFORMER sequence pair mask has the following format:
- """
- sep = [self.sep_token_id]
-
- # no bos used in fairseq
- if token_ids_1 is None:
- return len(token_ids_0 + sep) * [0]
- return len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
-
- src_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["src_vocab_file"]
- )
- tgt_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["tgt_vocab_file"]
- )
- merges_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
- )
-
- with open(src_vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
- with open(tgt_vocab_file, "w", encoding="utf-8") as f:
- tgt_vocab = {v: k for k, v in self.decoder.items()}
- f.write(json.dumps(tgt_vocab, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
- index = 0
- with open(merges_file, "w", encoding="utf-8") as writer:
- for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {merges_file}: BPE merge indices are not consecutive."
- " Please check that the tokenizer is not corrupted!"
- )
- index = token_index
- writer.write(" ".join(bpe_tokens) + "\n")
- index += 1
-
- return src_vocab_file, tgt_vocab_file, merges_file
-
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sm"] = None
- return state
-
- def __setstate__(self, d):
- self.__dict__ = d
-
- try:
- import sacremoses
- except ImportError:
- raise ImportError(
- "You need to install sacremoses to use XLMTokenizer. "
- "See https://pypi.org/project/sacremoses/ for installation."
- )
-
- self.sm = sacremoses
diff --git a/transformers/models/funnel/__init__.py b/transformers/models/funnel/__init__.py
deleted file mode 100644
index 28b9a34290c8264e37ddd3a20e1c6c15e28bcd5c..0000000000000000000000000000000000000000
--- a/transformers/models/funnel/__init__.py
+++ /dev/null
@@ -1,134 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_tf_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_funnel": ["FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP", "FunnelConfig"],
- "convert_funnel_original_tf_checkpoint_to_pytorch": [],
- "tokenization_funnel": ["FunnelTokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_funnel_fast"] = ["FunnelTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_funnel"] = [
- "FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST",
- "FunnelBaseModel",
- "FunnelForMaskedLM",
- "FunnelForMultipleChoice",
- "FunnelForPreTraining",
- "FunnelForQuestionAnswering",
- "FunnelForSequenceClassification",
- "FunnelForTokenClassification",
- "FunnelModel",
- "FunnelPreTrainedModel",
- "load_tf_weights_in_funnel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_funnel"] = [
- "TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFFunnelBaseModel",
- "TFFunnelForMaskedLM",
- "TFFunnelForMultipleChoice",
- "TFFunnelForPreTraining",
- "TFFunnelForQuestionAnswering",
- "TFFunnelForSequenceClassification",
- "TFFunnelForTokenClassification",
- "TFFunnelModel",
- "TFFunnelPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
- from .tokenization_funnel import FunnelTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_funnel_fast import FunnelTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_funnel import (
- FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
- FunnelBaseModel,
- FunnelForMaskedLM,
- FunnelForMultipleChoice,
- FunnelForPreTraining,
- FunnelForQuestionAnswering,
- FunnelForSequenceClassification,
- FunnelForTokenClassification,
- FunnelModel,
- FunnelPreTrainedModel,
- load_tf_weights_in_funnel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_funnel import (
- TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFFunnelBaseModel,
- TFFunnelForMaskedLM,
- TFFunnelForMultipleChoice,
- TFFunnelForPreTraining,
- TFFunnelForQuestionAnswering,
- TFFunnelForSequenceClassification,
- TFFunnelForTokenClassification,
- TFFunnelModel,
- TFFunnelPreTrainedModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/funnel/__pycache__/__init__.cpython-310.pyc b/transformers/models/funnel/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index fb02c95aaee808489b1fcef702cc761689ddd90f..0000000000000000000000000000000000000000
Binary files a/transformers/models/funnel/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/funnel/__pycache__/configuration_funnel.cpython-310.pyc b/transformers/models/funnel/__pycache__/configuration_funnel.cpython-310.pyc
deleted file mode 100644
index a39ab03ce516bb5159bd0d06cca599d68fc1e24b..0000000000000000000000000000000000000000
Binary files a/transformers/models/funnel/__pycache__/configuration_funnel.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/funnel/__pycache__/convert_funnel_original_tf_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/funnel/__pycache__/convert_funnel_original_tf_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 026eaa1d53c02810deab476ec8931b748c12bbcc..0000000000000000000000000000000000000000
Binary files a/transformers/models/funnel/__pycache__/convert_funnel_original_tf_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/funnel/__pycache__/modeling_funnel.cpython-310.pyc b/transformers/models/funnel/__pycache__/modeling_funnel.cpython-310.pyc
deleted file mode 100644
index 9b5e54273c0b093b74359b013c8c0789c06b0711..0000000000000000000000000000000000000000
Binary files a/transformers/models/funnel/__pycache__/modeling_funnel.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/funnel/__pycache__/modeling_tf_funnel.cpython-310.pyc b/transformers/models/funnel/__pycache__/modeling_tf_funnel.cpython-310.pyc
deleted file mode 100644
index a0d71dc0be20cb3662bd68e13466c27db72c37d9..0000000000000000000000000000000000000000
Binary files a/transformers/models/funnel/__pycache__/modeling_tf_funnel.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/funnel/__pycache__/tokenization_funnel.cpython-310.pyc b/transformers/models/funnel/__pycache__/tokenization_funnel.cpython-310.pyc
deleted file mode 100644
index 0969867c531e8d04e107621dfa23c86d9189bfff..0000000000000000000000000000000000000000
Binary files a/transformers/models/funnel/__pycache__/tokenization_funnel.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/funnel/__pycache__/tokenization_funnel_fast.cpython-310.pyc b/transformers/models/funnel/__pycache__/tokenization_funnel_fast.cpython-310.pyc
deleted file mode 100644
index e6d3ee28b4733c4fe0b69dada90ef1618396f1be..0000000000000000000000000000000000000000
Binary files a/transformers/models/funnel/__pycache__/tokenization_funnel_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/funnel/configuration_funnel.py b/transformers/models/funnel/configuration_funnel.py
deleted file mode 100644
index 0b49c22fb4c345fa2e997c5bd5eaa865c680068f..0000000000000000000000000000000000000000
--- a/transformers/models/funnel/configuration_funnel.py
+++ /dev/null
@@ -1,166 +0,0 @@
-# coding=utf-8
-# Copyright 2020, Hugging Face
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Funnel Transformer model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class FunnelConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FunnelModel`] or a [`TFBertModel`]. It is used to
- instantiate a Funnel Transformer model according to the specified arguments, defining the model architecture.
- Instantiating a configuration with the defaults will yield a similar configuration to that of the Funnel
- Transformer [funnel-transformer/small](https://huggingface.co/funnel-transformer/small) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the Funnel transformer. Defines the number of different tokens that can be represented
- by the `inputs_ids` passed when calling [`FunnelModel`] or [`TFFunnelModel`].
- block_sizes (`List[int]`, *optional*, defaults to `[4, 4, 4]`):
- The sizes of the blocks used in the model.
- block_repeats (`List[int]`, *optional*):
- If passed along, each layer of each block is repeated the number of times indicated.
- num_decoder_layers (`int`, *optional*, defaults to 2):
- The number of layers in the decoder (when not using the base model).
- d_model (`int`, *optional*, defaults to 768):
- Dimensionality of the model's hidden states.
- n_head (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- d_head (`int`, *optional*, defaults to 64):
- Dimensionality of the model's heads.
- d_inner (`int`, *optional*, defaults to 3072):
- Inner dimension in the feed-forward blocks.
- hidden_act (`str` or `callable`, *optional*, defaults to `"gelu_new"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- hidden_dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.1):
- The dropout probability for the attention probabilities.
- activation_dropout (`float`, *optional*, defaults to 0.0):
- The dropout probability used between the two layers of the feed-forward blocks.
- initializer_range (`float`, *optional*, defaults to 0.1):
- The upper bound of the *uniform initializer* for initializing all weight matrices in attention layers.
- initializer_std (`float`, *optional*):
- The standard deviation of the *normal initializer* for initializing the embedding matrix and the weight of
- linear layers. Will default to 1 for the embedding matrix and the value given by Xavier initialization for
- linear layers.
- layer_norm_eps (`float`, *optional*, defaults to 1e-09):
- The epsilon used by the layer normalization layers.
- pooling_type (`str`, *optional*, defaults to `"mean"`):
- Possible values are `"mean"` or `"max"`. The way pooling is performed at the beginning of each block.
- attention_type (`str`, *optional*, defaults to `"relative_shift"`):
- Possible values are `"relative_shift"` or `"factorized"`. The former is faster on CPU/GPU while the latter
- is faster on TPU.
- separate_cls (`bool`, *optional*, defaults to `True`):
- Whether or not to separate the cls token when applying pooling.
- truncate_seq (`bool`, *optional*, defaults to `True`):
- When using `separate_cls`, whether or not to truncate the last token when pooling, to avoid getting a
- sequence length that is not a multiple of 2.
- pool_q_only (`bool`, *optional*, defaults to `True`):
- Whether or not to apply the pooling only to the query or to query, key and values for the attention layers.
- """
-
- model_type = "funnel"
- attribute_map = {
- "hidden_size": "d_model",
- "num_attention_heads": "n_head",
- }
-
- def __init__(
- self,
- vocab_size=30522,
- block_sizes=[4, 4, 4],
- block_repeats=None,
- num_decoder_layers=2,
- d_model=768,
- n_head=12,
- d_head=64,
- d_inner=3072,
- hidden_act="gelu_new",
- hidden_dropout=0.1,
- attention_dropout=0.1,
- activation_dropout=0.0,
- initializer_range=0.1,
- initializer_std=None,
- layer_norm_eps=1e-9,
- pooling_type="mean",
- attention_type="relative_shift",
- separate_cls=True,
- truncate_seq=True,
- pool_q_only=True,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.block_sizes = block_sizes
- self.block_repeats = [1] * len(block_sizes) if block_repeats is None else block_repeats
- assert len(block_sizes) == len(
- self.block_repeats
- ), "`block_sizes` and `block_repeats` should have the same length."
- self.num_decoder_layers = num_decoder_layers
- self.d_model = d_model
- self.n_head = n_head
- self.d_head = d_head
- self.d_inner = d_inner
- self.hidden_act = hidden_act
- self.hidden_dropout = hidden_dropout
- self.attention_dropout = attention_dropout
- self.activation_dropout = activation_dropout
- self.initializer_range = initializer_range
- self.initializer_std = initializer_std
- self.layer_norm_eps = layer_norm_eps
- assert pooling_type in [
- "mean",
- "max",
- ], f"Got {pooling_type} for `pooling_type` but only 'mean' and 'max' are supported."
- self.pooling_type = pooling_type
- assert attention_type in [
- "relative_shift",
- "factorized",
- ], f"Got {attention_type} for `attention_type` but only 'relative_shift' and 'factorized' are supported."
- self.attention_type = attention_type
- self.separate_cls = separate_cls
- self.truncate_seq = truncate_seq
- self.pool_q_only = pool_q_only
-
- super().__init__(**kwargs)
-
- @property
- def num_hidden_layers(self):
- return sum(self.block_sizes)
-
- @num_hidden_layers.setter
- def num_hidden_layers(self, value):
- raise NotImplementedError(
- "This model does not support the setting of `num_hidden_layers`. Please set `block_sizes`."
- )
-
- @property
- def num_blocks(self):
- return len(self.block_sizes)
-
- @num_blocks.setter
- def num_blocks(self, value):
- raise NotImplementedError("This model does not support the setting of `num_blocks`. Please set `block_sizes`.")
diff --git a/transformers/models/funnel/convert_funnel_original_tf_checkpoint_to_pytorch.py b/transformers/models/funnel/convert_funnel_original_tf_checkpoint_to_pytorch.py
deleted file mode 100644
index 848101f083582bafa26e58c87aaa612502f3f79c..0000000000000000000000000000000000000000
--- a/transformers/models/funnel/convert_funnel_original_tf_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert Funnel checkpoint."""
-
-
-import argparse
-
-import torch
-
-from transformers import FunnelBaseModel, FunnelConfig, FunnelModel, load_tf_weights_in_funnel
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-
-
-def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path, base_model):
- # Initialise PyTorch model
- config = FunnelConfig.from_json_file(config_file)
- print(f"Building PyTorch model from configuration: {config}")
- model = FunnelBaseModel(config) if base_model else FunnelModel(config)
-
- # Load weights from tf checkpoint
- load_tf_weights_in_funnel(model, config, tf_checkpoint_path)
-
- # Save pytorch-model
- print(f"Save PyTorch model to {pytorch_dump_path}")
- torch.save(model.state_dict(), pytorch_dump_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
- )
- parser.add_argument(
- "--config_file",
- default=None,
- type=str,
- required=True,
- help="The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.",
- )
- parser.add_argument(
- "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
- parser.add_argument(
- "--base_model", action="store_true", help="Whether you want just the base model (no decoder) or not."
- )
- args = parser.parse_args()
- convert_tf_checkpoint_to_pytorch(
- args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.base_model
- )
diff --git a/transformers/models/funnel/modeling_funnel.py b/transformers/models/funnel/modeling_funnel.py
deleted file mode 100644
index ce0c7789487d8fa7f376f383a40958f9aeb3fb37..0000000000000000000000000000000000000000
--- a/transformers/models/funnel/modeling_funnel.py
+++ /dev/null
@@ -1,1599 +0,0 @@
-# coding=utf-8
-# Copyright 2020-present Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Funnel Transformer model."""
-
-import os
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import numpy as np
-import torch
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutput,
- MaskedLMOutput,
- MultipleChoiceModelOutput,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutput,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_funnel import FunnelConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "FunnelConfig"
-_CHECKPOINT_FOR_DOC = "funnel-transformer/small"
-
-
-from ..deprecated._archive_maps import FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-INF = 1e6
-
-
-def load_tf_weights_in_funnel(model, config, tf_checkpoint_path):
- """Load tf checkpoints in a pytorch model."""
- try:
- import re
-
- import numpy as np
- import tensorflow as tf
- except ImportError:
- logger.error(
- "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
- "https://www.tensorflow.org/install/ for installation instructions."
- )
- raise
- tf_path = os.path.abspath(tf_checkpoint_path)
- logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
- # Load weights from TF model
- init_vars = tf.train.list_variables(tf_path)
- names = []
- arrays = []
- for name, shape in init_vars:
- logger.info(f"Loading TF weight {name} with shape {shape}")
- array = tf.train.load_variable(tf_path, name)
- names.append(name)
- arrays.append(array)
-
- _layer_map = {
- "k": "k_head",
- "q": "q_head",
- "v": "v_head",
- "o": "post_proj",
- "layer_1": "linear_1",
- "layer_2": "linear_2",
- "rel_attn": "attention",
- "ff": "ffn",
- "kernel": "weight",
- "gamma": "weight",
- "beta": "bias",
- "lookup_table": "weight",
- "word_embedding": "word_embeddings",
- "input": "embeddings",
- }
-
- for name, array in zip(names, arrays):
- name = name.split("/")
- # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
- # which are not required for using pretrained model
- if any(
- n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
- for n in name
- ):
- logger.info(f"Skipping {'/'.join(name)}")
- continue
- if name[0] == "generator":
- continue
- pointer = model
- skipped = False
- for m_name in name[1:]:
- if not isinstance(pointer, FunnelPositionwiseFFN) and re.fullmatch(r"layer_\d+", m_name):
- layer_index = int(re.search(r"layer_(\d+)", m_name).groups()[0])
- if layer_index < config.num_hidden_layers:
- block_idx = 0
- while layer_index >= config.block_sizes[block_idx]:
- layer_index -= config.block_sizes[block_idx]
- block_idx += 1
- pointer = pointer.blocks[block_idx][layer_index]
- else:
- layer_index -= config.num_hidden_layers
- pointer = pointer.layers[layer_index]
- elif m_name == "r" and isinstance(pointer, FunnelRelMultiheadAttention):
- pointer = pointer.r_kernel
- break
- elif m_name in _layer_map:
- pointer = getattr(pointer, _layer_map[m_name])
- else:
- try:
- pointer = getattr(pointer, m_name)
- except AttributeError:
- print(f"Skipping {'/'.join(name)}", array.shape)
- skipped = True
- break
- if not skipped:
- if len(pointer.shape) != len(array.shape):
- array = array.reshape(pointer.shape)
- if m_name == "kernel":
- array = np.transpose(array)
- pointer.data = torch.from_numpy(array)
-
- return model
-
-
-class FunnelEmbeddings(nn.Module):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
- self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout)
-
- def forward(
- self, input_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None
- ) -> torch.Tensor:
- if inputs_embeds is None:
- inputs_embeds = self.word_embeddings(input_ids)
- embeddings = self.layer_norm(inputs_embeds)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-class FunnelAttentionStructure(nn.Module):
- """
- Contains helpers for `FunnelRelMultiheadAttention `.
- """
-
- cls_token_type_id: int = 2
-
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__()
- self.config = config
- self.sin_dropout = nn.Dropout(config.hidden_dropout)
- self.cos_dropout = nn.Dropout(config.hidden_dropout)
- # Track where we are at in terms of pooling from the original input, e.g., by how much the sequence length was
- # divided.
- self.pooling_mult = None
-
- def init_attention_inputs(
- self,
- inputs_embeds: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- ) -> Tuple[torch.Tensor]:
- """Returns the attention inputs associated to the inputs of the model."""
- # inputs_embeds has shape batch_size x seq_len x d_model
- # attention_mask and token_type_ids have shape batch_size x seq_len
- self.pooling_mult = 1
- self.seq_len = seq_len = inputs_embeds.size(1)
- position_embeds = self.get_position_embeds(seq_len, inputs_embeds.dtype, inputs_embeds.device)
- token_type_mat = self.token_type_ids_to_mat(token_type_ids) if token_type_ids is not None else None
- cls_mask = (
- nn.functional.pad(inputs_embeds.new_ones([seq_len - 1, seq_len - 1]), (1, 0, 1, 0))
- if self.config.separate_cls
- else None
- )
- return (position_embeds, token_type_mat, attention_mask, cls_mask)
-
- def token_type_ids_to_mat(self, token_type_ids: torch.Tensor) -> torch.Tensor:
- """Convert `token_type_ids` to `token_type_mat`."""
- token_type_mat = token_type_ids[:, :, None] == token_type_ids[:, None]
- # Treat as in the same segment as both A & B
- cls_ids = token_type_ids == self.cls_token_type_id
- cls_mat = cls_ids[:, :, None] | cls_ids[:, None]
- return cls_mat | token_type_mat
-
- def get_position_embeds(
- self, seq_len: int, dtype: torch.dtype, device: torch.device
- ) -> Union[Tuple[torch.Tensor], List[List[torch.Tensor]]]:
- """
- Create and cache inputs related to relative position encoding. Those are very different depending on whether we
- are using the factorized or the relative shift attention:
-
- For the factorized attention, it returns the matrices (phi, pi, psi, omega) used in the paper, appendix A.2.2,
- final formula.
-
- For the relative shift attention, it returns all possible vectors R used in the paper, appendix A.2.1, final
- formula.
-
- Paper link: https://arxiv.org/abs/2006.03236
- """
- d_model = self.config.d_model
- if self.config.attention_type == "factorized":
- # Notations from the paper, appending A.2.2, final formula.
- # We need to create and return the matrices phi, psi, pi and omega.
- pos_seq = torch.arange(0, seq_len, 1.0, dtype=torch.int64, device=device).to(dtype)
- freq_seq = torch.arange(0, d_model // 2, 1.0, dtype=torch.int64, device=device).to(dtype)
- inv_freq = 1 / (10000 ** (freq_seq / (d_model // 2)))
- sinusoid = pos_seq[:, None] * inv_freq[None]
- sin_embed = torch.sin(sinusoid)
- sin_embed_d = self.sin_dropout(sin_embed)
- cos_embed = torch.cos(sinusoid)
- cos_embed_d = self.cos_dropout(cos_embed)
- # This is different from the formula on the paper...
- phi = torch.cat([sin_embed_d, sin_embed_d], dim=-1)
- psi = torch.cat([cos_embed, sin_embed], dim=-1)
- pi = torch.cat([cos_embed_d, cos_embed_d], dim=-1)
- omega = torch.cat([-sin_embed, cos_embed], dim=-1)
- return (phi, pi, psi, omega)
- else:
- # Notations from the paper, appending A.2.1, final formula.
- # We need to create and return all the possible vectors R for all blocks and shifts.
- freq_seq = torch.arange(0, d_model // 2, 1.0, dtype=torch.int64, device=device).to(dtype)
- inv_freq = 1 / (10000 ** (freq_seq / (d_model // 2)))
- # Maximum relative positions for the first input
- rel_pos_id = torch.arange(-seq_len * 2, seq_len * 2, 1.0, dtype=torch.int64, device=device).to(dtype)
- zero_offset = seq_len * 2
- sinusoid = rel_pos_id[:, None] * inv_freq[None]
- sin_embed = self.sin_dropout(torch.sin(sinusoid))
- cos_embed = self.cos_dropout(torch.cos(sinusoid))
- pos_embed = torch.cat([sin_embed, cos_embed], dim=-1)
-
- pos = torch.arange(0, seq_len, dtype=torch.int64, device=device).to(dtype)
- pooled_pos = pos
- position_embeds_list = []
- for block_index in range(0, self.config.num_blocks):
- # For each block with block_index > 0, we need two types position embeddings:
- # - Attention(pooled-q, unpooled-kv)
- # - Attention(pooled-q, pooled-kv)
- # For block_index = 0 we only need the second one and leave the first one as None.
-
- # First type
- if block_index == 0:
- position_embeds_pooling = None
- else:
- pooled_pos = self.stride_pool_pos(pos, block_index)
-
- # construct rel_pos_id
- stride = 2 ** (block_index - 1)
- rel_pos = self.relative_pos(pos, stride, pooled_pos, shift=2)
- rel_pos = rel_pos[:, None] + zero_offset
- rel_pos = rel_pos.expand(rel_pos.size(0), d_model)
- position_embeds_pooling = torch.gather(pos_embed, 0, rel_pos)
-
- # Second type
- pos = pooled_pos
- stride = 2**block_index
- rel_pos = self.relative_pos(pos, stride)
-
- rel_pos = rel_pos[:, None] + zero_offset
- rel_pos = rel_pos.expand(rel_pos.size(0), d_model)
- position_embeds_no_pooling = torch.gather(pos_embed, 0, rel_pos)
-
- position_embeds_list.append([position_embeds_no_pooling, position_embeds_pooling])
- return position_embeds_list
-
- def stride_pool_pos(self, pos_id: torch.Tensor, block_index: int):
- """
- Pool `pos_id` while keeping the cls token separate (if `config.separate_cls=True`).
- """
- if self.config.separate_cls:
- # Under separate , we treat the as the first token in
- # the previous block of the 1st real block. Since the 1st real
- # block always has position 1, the position of the previous block
- # will be at `1 - 2 ** block_index`.
- cls_pos = pos_id.new_tensor([-(2**block_index) + 1])
- pooled_pos_id = pos_id[1:-1] if self.config.truncate_seq else pos_id[1:]
- return torch.cat([cls_pos, pooled_pos_id[::2]], 0)
- else:
- return pos_id[::2]
-
- def relative_pos(self, pos: torch.Tensor, stride: int, pooled_pos=None, shift: int = 1) -> torch.Tensor:
- """
- Build the relative positional vector between `pos` and `pooled_pos`.
- """
- if pooled_pos is None:
- pooled_pos = pos
-
- ref_point = pooled_pos[0] - pos[0]
- num_remove = shift * len(pooled_pos)
- max_dist = ref_point + num_remove * stride
- min_dist = pooled_pos[0] - pos[-1]
-
- return torch.arange(max_dist, min_dist - 1, -stride, dtype=torch.long, device=pos.device)
-
- def stride_pool(
- self,
- tensor: Union[torch.Tensor, Tuple[torch.Tensor], List[torch.Tensor]],
- axis: Union[int, Tuple[int], List[int]],
- ) -> torch.Tensor:
- """
- Perform pooling by stride slicing the tensor along the given axis.
- """
- if tensor is None:
- return None
-
- # Do the stride pool recursively if axis is a list or a tuple of ints.
- if isinstance(axis, (list, tuple)):
- for ax in axis:
- tensor = self.stride_pool(tensor, ax)
- return tensor
-
- # Do the stride pool recursively if tensor is a list or tuple of tensors.
- if isinstance(tensor, (tuple, list)):
- return type(tensor)(self.stride_pool(x, axis) for x in tensor)
-
- # Deal with negative axis
- axis %= tensor.ndim
-
- axis_slice = (
- slice(None, -1, 2) if self.config.separate_cls and self.config.truncate_seq else slice(None, None, 2)
- )
- enc_slice = [slice(None)] * axis + [axis_slice]
- if self.config.separate_cls:
- cls_slice = [slice(None)] * axis + [slice(None, 1)]
- tensor = torch.cat([tensor[cls_slice], tensor], axis=axis)
- return tensor[enc_slice]
-
- def pool_tensor(
- self, tensor: Union[torch.Tensor, Tuple[torch.Tensor], List[torch.Tensor]], mode: str = "mean", stride: int = 2
- ) -> torch.Tensor:
- """Apply 1D pooling to a tensor of size [B x T (x H)]."""
- if tensor is None:
- return None
-
- # Do the pool recursively if tensor is a list or tuple of tensors.
- if isinstance(tensor, (tuple, list)):
- return type(tensor)(self.pool_tensor(tensor, mode=mode, stride=stride) for x in tensor)
-
- if self.config.separate_cls:
- suffix = tensor[:, :-1] if self.config.truncate_seq else tensor
- tensor = torch.cat([tensor[:, :1], suffix], dim=1)
-
- ndim = tensor.ndim
- if ndim == 2:
- tensor = tensor[:, None, :, None]
- elif ndim == 3:
- tensor = tensor[:, None, :, :]
- # Stride is applied on the second-to-last dimension.
- stride = (stride, 1)
-
- if mode == "mean":
- tensor = nn.functional.avg_pool2d(tensor, stride, stride=stride, ceil_mode=True)
- elif mode == "max":
- tensor = nn.functional.max_pool2d(tensor, stride, stride=stride, ceil_mode=True)
- elif mode == "min":
- tensor = -nn.functional.max_pool2d(-tensor, stride, stride=stride, ceil_mode=True)
- else:
- raise NotImplementedError("The supported modes are 'mean', 'max' and 'min'.")
-
- if ndim == 2:
- return tensor[:, 0, :, 0]
- elif ndim == 3:
- return tensor[:, 0]
- return tensor
-
- def pre_attention_pooling(
- self, output, attention_inputs: Tuple[torch.Tensor]
- ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]:
- """Pool `output` and the proper parts of `attention_inputs` before the attention layer."""
- position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
- if self.config.pool_q_only:
- if self.config.attention_type == "factorized":
- position_embeds = self.stride_pool(position_embeds[:2], 0) + position_embeds[2:]
- token_type_mat = self.stride_pool(token_type_mat, 1)
- cls_mask = self.stride_pool(cls_mask, 0)
- output = self.pool_tensor(output, mode=self.config.pooling_type)
- else:
- self.pooling_mult *= 2
- if self.config.attention_type == "factorized":
- position_embeds = self.stride_pool(position_embeds, 0)
- token_type_mat = self.stride_pool(token_type_mat, [1, 2])
- cls_mask = self.stride_pool(cls_mask, [1, 2])
- attention_mask = self.pool_tensor(attention_mask, mode="min")
- output = self.pool_tensor(output, mode=self.config.pooling_type)
- attention_inputs = (position_embeds, token_type_mat, attention_mask, cls_mask)
- return output, attention_inputs
-
- def post_attention_pooling(self, attention_inputs: Tuple[torch.Tensor]) -> Tuple[torch.Tensor]:
- """Pool the proper parts of `attention_inputs` after the attention layer."""
- position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
- if self.config.pool_q_only:
- self.pooling_mult *= 2
- if self.config.attention_type == "factorized":
- position_embeds = position_embeds[:2] + self.stride_pool(position_embeds[2:], 0)
- token_type_mat = self.stride_pool(token_type_mat, 2)
- cls_mask = self.stride_pool(cls_mask, 1)
- attention_mask = self.pool_tensor(attention_mask, mode="min")
- attention_inputs = (position_embeds, token_type_mat, attention_mask, cls_mask)
- return attention_inputs
-
-
-def _relative_shift_gather(positional_attn: torch.Tensor, context_len: int, shift: int) -> torch.Tensor:
- batch_size, n_head, seq_len, max_rel_len = positional_attn.shape
- # max_rel_len = 2 * context_len + shift -1 is the numbers of possible relative positions i-j
-
- # What's next is the same as doing the following gather, which might be clearer code but less efficient.
- # idxs = context_len + torch.arange(0, context_len).unsqueeze(0) - torch.arange(0, seq_len).unsqueeze(1)
- # # matrix of context_len + i-j
- # return positional_attn.gather(3, idxs.expand([batch_size, n_head, context_len, context_len]))
-
- positional_attn = torch.reshape(positional_attn, [batch_size, n_head, max_rel_len, seq_len])
- positional_attn = positional_attn[:, :, shift:, :]
- positional_attn = torch.reshape(positional_attn, [batch_size, n_head, seq_len, max_rel_len - shift])
- positional_attn = positional_attn[..., :context_len]
- return positional_attn
-
-
-class FunnelRelMultiheadAttention(nn.Module):
- def __init__(self, config: FunnelConfig, block_index: int) -> None:
- super().__init__()
- self.config = config
- self.block_index = block_index
- d_model, n_head, d_head = config.d_model, config.n_head, config.d_head
-
- self.hidden_dropout = nn.Dropout(config.hidden_dropout)
- self.attention_dropout = nn.Dropout(config.attention_dropout)
-
- self.q_head = nn.Linear(d_model, n_head * d_head, bias=False)
- self.k_head = nn.Linear(d_model, n_head * d_head)
- self.v_head = nn.Linear(d_model, n_head * d_head)
-
- self.r_w_bias = nn.Parameter(torch.zeros([n_head, d_head]))
- self.r_r_bias = nn.Parameter(torch.zeros([n_head, d_head]))
- self.r_kernel = nn.Parameter(torch.zeros([d_model, n_head, d_head]))
- self.r_s_bias = nn.Parameter(torch.zeros([n_head, d_head]))
- self.seg_embed = nn.Parameter(torch.zeros([2, n_head, d_head]))
-
- self.post_proj = nn.Linear(n_head * d_head, d_model)
- self.layer_norm = nn.LayerNorm(d_model, eps=config.layer_norm_eps)
- self.scale = 1.0 / (d_head**0.5)
-
- def relative_positional_attention(self, position_embeds, q_head, context_len, cls_mask=None):
- """Relative attention score for the positional encodings"""
- # q_head has shape batch_size x sea_len x n_head x d_head
- if self.config.attention_type == "factorized":
- # Notations from the paper, appending A.2.2, final formula (https://arxiv.org/abs/2006.03236)
- # phi and pi have shape seq_len x d_model, psi and omega have shape context_len x d_model
- phi, pi, psi, omega = position_embeds
- # Shape n_head x d_head
- u = self.r_r_bias * self.scale
- # Shape d_model x n_head x d_head
- w_r = self.r_kernel
-
- # Shape batch_size x sea_len x n_head x d_model
- q_r_attention = torch.einsum("binh,dnh->bind", q_head + u, w_r)
- q_r_attention_1 = q_r_attention * phi[:, None]
- q_r_attention_2 = q_r_attention * pi[:, None]
-
- # Shape batch_size x n_head x seq_len x context_len
- positional_attn = torch.einsum("bind,jd->bnij", q_r_attention_1, psi) + torch.einsum(
- "bind,jd->bnij", q_r_attention_2, omega
- )
- else:
- shift = 2 if q_head.shape[1] != context_len else 1
- # Notations from the paper, appending A.2.1, final formula (https://arxiv.org/abs/2006.03236)
- # Grab the proper positional encoding, shape max_rel_len x d_model
- r = position_embeds[self.block_index][shift - 1]
- # Shape n_head x d_head
- v = self.r_r_bias * self.scale
- # Shape d_model x n_head x d_head
- w_r = self.r_kernel
-
- # Shape max_rel_len x n_head x d_model
- r_head = torch.einsum("td,dnh->tnh", r, w_r)
- # Shape batch_size x n_head x seq_len x max_rel_len
- positional_attn = torch.einsum("binh,tnh->bnit", q_head + v, r_head)
- # Shape batch_size x n_head x seq_len x context_len
- positional_attn = _relative_shift_gather(positional_attn, context_len, shift)
-
- if cls_mask is not None:
- positional_attn *= cls_mask
- return positional_attn
-
- def relative_token_type_attention(self, token_type_mat, q_head, cls_mask=None):
- """Relative attention score for the token_type_ids"""
- if token_type_mat is None:
- return 0
- batch_size, seq_len, context_len = token_type_mat.shape
- # q_head has shape batch_size x seq_len x n_head x d_head
- # Shape n_head x d_head
- r_s_bias = self.r_s_bias * self.scale
-
- # Shape batch_size x n_head x seq_len x 2
- token_type_bias = torch.einsum("bind,snd->bnis", q_head + r_s_bias, self.seg_embed)
- # Shape batch_size x n_head x seq_len x context_len
- token_type_mat = token_type_mat[:, None].expand([batch_size, q_head.shape[2], seq_len, context_len])
- # Shapes batch_size x n_head x seq_len
- diff_token_type, same_token_type = torch.split(token_type_bias, 1, dim=-1)
- # Shape batch_size x n_head x seq_len x context_len
- token_type_attn = torch.where(
- token_type_mat, same_token_type.expand(token_type_mat.shape), diff_token_type.expand(token_type_mat.shape)
- )
-
- if cls_mask is not None:
- token_type_attn *= cls_mask
- return token_type_attn
-
- def forward(
- self,
- query: torch.Tensor,
- key: torch.Tensor,
- value: torch.Tensor,
- attention_inputs: Tuple[torch.Tensor],
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, ...]:
- # query has shape batch_size x seq_len x d_model
- # key and value have shapes batch_size x context_len x d_model
- position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
-
- batch_size, seq_len, _ = query.shape
- context_len = key.shape[1]
- n_head, d_head = self.config.n_head, self.config.d_head
-
- # Shape batch_size x seq_len x n_head x d_head
- q_head = self.q_head(query).view(batch_size, seq_len, n_head, d_head)
- # Shapes batch_size x context_len x n_head x d_head
- k_head = self.k_head(key).view(batch_size, context_len, n_head, d_head)
- v_head = self.v_head(value).view(batch_size, context_len, n_head, d_head)
-
- q_head = q_head * self.scale
- # Shape n_head x d_head
- r_w_bias = self.r_w_bias * self.scale
- # Shapes batch_size x n_head x seq_len x context_len
- content_score = torch.einsum("bind,bjnd->bnij", q_head + r_w_bias, k_head)
- positional_attn = self.relative_positional_attention(position_embeds, q_head, context_len, cls_mask)
- token_type_attn = self.relative_token_type_attention(token_type_mat, q_head, cls_mask)
-
- # merge attention scores
- attn_score = content_score + positional_attn + token_type_attn
-
- # precision safe in case of mixed precision training
- dtype = attn_score.dtype
- attn_score = attn_score.float()
- # perform masking
- if attention_mask is not None:
- attn_score = attn_score - INF * (1 - attention_mask[:, None, None].float())
- # attention probability
- attn_prob = torch.softmax(attn_score, dim=-1, dtype=dtype)
- attn_prob = self.attention_dropout(attn_prob)
-
- # attention output, shape batch_size x seq_len x n_head x d_head
- attn_vec = torch.einsum("bnij,bjnd->bind", attn_prob, v_head)
-
- # Shape shape batch_size x seq_len x d_model
- attn_out = self.post_proj(attn_vec.reshape(batch_size, seq_len, n_head * d_head))
- attn_out = self.hidden_dropout(attn_out)
-
- output = self.layer_norm(query + attn_out)
- return (output, attn_prob) if output_attentions else (output,)
-
-
-class FunnelPositionwiseFFN(nn.Module):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__()
- self.linear_1 = nn.Linear(config.d_model, config.d_inner)
- self.activation_function = ACT2FN[config.hidden_act]
- self.activation_dropout = nn.Dropout(config.activation_dropout)
- self.linear_2 = nn.Linear(config.d_inner, config.d_model)
- self.dropout = nn.Dropout(config.hidden_dropout)
- self.layer_norm = nn.LayerNorm(config.d_model, config.layer_norm_eps)
-
- def forward(self, hidden: torch.Tensor) -> torch.Tensor:
- h = self.linear_1(hidden)
- h = self.activation_function(h)
- h = self.activation_dropout(h)
- h = self.linear_2(h)
- h = self.dropout(h)
- return self.layer_norm(hidden + h)
-
-
-class FunnelLayer(nn.Module):
- def __init__(self, config: FunnelConfig, block_index: int) -> None:
- super().__init__()
- self.attention = FunnelRelMultiheadAttention(config, block_index)
- self.ffn = FunnelPositionwiseFFN(config)
-
- def forward(
- self,
- query: torch.Tensor,
- key: torch.Tensor,
- value: torch.Tensor,
- attention_inputs,
- output_attentions: bool = False,
- ) -> Tuple:
- attn = self.attention(query, key, value, attention_inputs, output_attentions=output_attentions)
- output = self.ffn(attn[0])
- return (output, attn[1]) if output_attentions else (output,)
-
-
-class FunnelEncoder(nn.Module):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__()
- self.config = config
- self.attention_structure = FunnelAttentionStructure(config)
- self.blocks = nn.ModuleList(
- [
- nn.ModuleList([FunnelLayer(config, block_index) for _ in range(block_size)])
- for block_index, block_size in enumerate(config.block_sizes)
- ]
- )
-
- def forward(
- self,
- inputs_embeds: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[Tuple, BaseModelOutput]:
- # The pooling is not implemented on long tensors, so we convert this mask.
- attention_mask = attention_mask.type_as(inputs_embeds)
- attention_inputs = self.attention_structure.init_attention_inputs(
- inputs_embeds,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- )
- hidden = inputs_embeds
-
- all_hidden_states = (inputs_embeds,) if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- for block_index, block in enumerate(self.blocks):
- pooling_flag = hidden.size(1) > (2 if self.config.separate_cls else 1)
- pooling_flag = pooling_flag and block_index > 0
- if pooling_flag:
- pooled_hidden, attention_inputs = self.attention_structure.pre_attention_pooling(
- hidden, attention_inputs
- )
- for layer_index, layer in enumerate(block):
- for repeat_index in range(self.config.block_repeats[block_index]):
- do_pooling = (repeat_index == 0) and (layer_index == 0) and pooling_flag
- if do_pooling:
- query = pooled_hidden
- key = value = hidden if self.config.pool_q_only else pooled_hidden
- else:
- query = key = value = hidden
- layer_output = layer(query, key, value, attention_inputs, output_attentions=output_attentions)
- hidden = layer_output[0]
- if do_pooling:
- attention_inputs = self.attention_structure.post_attention_pooling(attention_inputs)
-
- if output_attentions:
- all_attentions = all_attentions + layer_output[1:]
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden,)
-
- if not return_dict:
- return tuple(v for v in [hidden, all_hidden_states, all_attentions] if v is not None)
- return BaseModelOutput(last_hidden_state=hidden, hidden_states=all_hidden_states, attentions=all_attentions)
-
-
-def upsample(
- x: torch.Tensor, stride: int, target_len: int, separate_cls: bool = True, truncate_seq: bool = False
-) -> torch.Tensor:
- """
- Upsample tensor `x` to match `target_len` by repeating the tokens `stride` time on the sequence length dimension.
- """
- if stride == 1:
- return x
- if separate_cls:
- cls = x[:, :1]
- x = x[:, 1:]
- output = torch.repeat_interleave(x, repeats=stride, dim=1)
- if separate_cls:
- if truncate_seq:
- output = nn.functional.pad(output, (0, 0, 0, stride - 1, 0, 0))
- output = output[:, : target_len - 1]
- output = torch.cat([cls, output], dim=1)
- else:
- output = output[:, :target_len]
- return output
-
-
-class FunnelDecoder(nn.Module):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__()
- self.config = config
- self.attention_structure = FunnelAttentionStructure(config)
- self.layers = nn.ModuleList([FunnelLayer(config, 0) for _ in range(config.num_decoder_layers)])
-
- def forward(
- self,
- final_hidden: torch.Tensor,
- first_block_hidden: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ) -> Union[Tuple, BaseModelOutput]:
- upsampled_hidden = upsample(
- final_hidden,
- stride=2 ** (len(self.config.block_sizes) - 1),
- target_len=first_block_hidden.shape[1],
- separate_cls=self.config.separate_cls,
- truncate_seq=self.config.truncate_seq,
- )
-
- hidden = upsampled_hidden + first_block_hidden
- all_hidden_states = (hidden,) if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- attention_inputs = self.attention_structure.init_attention_inputs(
- hidden,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- )
-
- for layer in self.layers:
- layer_output = layer(hidden, hidden, hidden, attention_inputs, output_attentions=output_attentions)
- hidden = layer_output[0]
-
- if output_attentions:
- all_attentions = all_attentions + layer_output[1:]
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden,)
-
- if not return_dict:
- return tuple(v for v in [hidden, all_hidden_states, all_attentions] if v is not None)
- return BaseModelOutput(last_hidden_state=hidden, hidden_states=all_hidden_states, attentions=all_attentions)
-
-
-class FunnelDiscriminatorPredictions(nn.Module):
- """Prediction module for the discriminator, made up of two dense layers."""
-
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__()
- self.config = config
- self.dense = nn.Linear(config.d_model, config.d_model)
- self.dense_prediction = nn.Linear(config.d_model, 1)
-
- def forward(self, discriminator_hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(discriminator_hidden_states)
- hidden_states = ACT2FN[self.config.hidden_act](hidden_states)
- logits = self.dense_prediction(hidden_states).squeeze(-1)
- return logits
-
-
-class FunnelPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = FunnelConfig
- load_tf_weights = load_tf_weights_in_funnel
- base_model_prefix = "funnel"
-
- def _init_weights(self, module):
- classname = module.__class__.__name__
- if classname.find("Linear") != -1:
- if getattr(module, "weight", None) is not None:
- if self.config.initializer_std is None:
- fan_out, fan_in = module.weight.shape
- std = np.sqrt(1.0 / float(fan_in + fan_out))
- else:
- std = self.config.initializer_std
- nn.init.normal_(module.weight, std=std)
- if getattr(module, "bias", None) is not None:
- nn.init.constant_(module.bias, 0.0)
- elif classname == "FunnelRelMultiheadAttention":
- nn.init.uniform_(module.r_w_bias, b=self.config.initializer_range)
- nn.init.uniform_(module.r_r_bias, b=self.config.initializer_range)
- nn.init.uniform_(module.r_kernel, b=self.config.initializer_range)
- nn.init.uniform_(module.r_s_bias, b=self.config.initializer_range)
- nn.init.uniform_(module.seg_embed, b=self.config.initializer_range)
- elif classname == "FunnelEmbeddings":
- std = 1.0 if self.config.initializer_std is None else self.config.initializer_std
- nn.init.normal_(module.word_embeddings.weight, std=std)
- if module.word_embeddings.padding_idx is not None:
- module.word_embeddings.weight.data[module.padding_idx].zero_()
-
-
-class FunnelClassificationHead(nn.Module):
- def __init__(self, config: FunnelConfig, n_labels: int) -> None:
- super().__init__()
- self.linear_hidden = nn.Linear(config.d_model, config.d_model)
- self.dropout = nn.Dropout(config.hidden_dropout)
- self.linear_out = nn.Linear(config.d_model, n_labels)
-
- def forward(self, hidden: torch.Tensor) -> torch.Tensor:
- hidden = self.linear_hidden(hidden)
- hidden = torch.tanh(hidden)
- hidden = self.dropout(hidden)
- return self.linear_out(hidden)
-
-
-@dataclass
-class FunnelForPreTrainingOutput(ModelOutput):
- """
- Output type of [`FunnelForPreTraining`].
-
- Args:
- loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
- Total loss of the ELECTRA-style objective.
- logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
- Prediction scores of the head (scores for each token before SoftMax).
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- logits: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-FUNNEL_START_DOCSTRING = r"""
-
- The Funnel Transformer model was proposed in [Funnel-Transformer: Filtering out Sequential Redundancy for Efficient
- Language Processing](https://arxiv.org/abs/2006.03236) by Zihang Dai, Guokun Lai, Yiming Yang, Quoc V. Le.
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`FunnelConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FUNNEL_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- """
- The base Funnel Transformer Model transformer outputting raw hidden-states without upsampling head (also called
- decoder) or any task-specific head on top.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class FunnelBaseModel(FunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__(config)
-
- self.embeddings = FunnelEmbeddings(config)
- self.encoder = FunnelEncoder(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self) -> nn.Embedding:
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
- self.embeddings.word_embeddings = new_embeddings
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small-base",
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device)
- if token_type_ids is None:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- # TODO: deal with head_mask
- if inputs_embeds is None:
- inputs_embeds = self.embeddings(input_ids)
-
- encoder_outputs = self.encoder(
- inputs_embeds,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- return encoder_outputs
-
-
-@add_start_docstrings(
- "The bare Funnel Transformer Model transformer outputting raw hidden-states without any specific head on top.",
- FUNNEL_START_DOCSTRING,
-)
-class FunnelModel(FunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__(config)
- self.config = config
- self.embeddings = FunnelEmbeddings(config)
- self.encoder = FunnelEncoder(config)
- self.decoder = FunnelDecoder(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self) -> nn.Embedding:
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
- self.embeddings.word_embeddings = new_embeddings
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if attention_mask is None:
- attention_mask = torch.ones(input_shape, device=device)
- if token_type_ids is None:
- token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
-
- # TODO: deal with head_mask
- if inputs_embeds is None:
- inputs_embeds = self.embeddings(input_ids)
-
- encoder_outputs = self.encoder(
- inputs_embeds,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- output_attentions=output_attentions,
- output_hidden_states=True,
- return_dict=return_dict,
- )
-
- decoder_outputs = self.decoder(
- final_hidden=encoder_outputs[0],
- first_block_hidden=encoder_outputs[1][self.config.block_sizes[0]],
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- if not return_dict:
- idx = 0
- outputs = (decoder_outputs[0],)
- if output_hidden_states:
- idx += 1
- outputs = outputs + (encoder_outputs[1] + decoder_outputs[idx],)
- if output_attentions:
- idx += 1
- outputs = outputs + (encoder_outputs[2] + decoder_outputs[idx],)
- return outputs
-
- return BaseModelOutput(
- last_hidden_state=decoder_outputs[0],
- hidden_states=(encoder_outputs.hidden_states + decoder_outputs.hidden_states)
- if output_hidden_states
- else None,
- attentions=(encoder_outputs.attentions + decoder_outputs.attentions) if output_attentions else None,
- )
-
-
-add_start_docstrings(
- """
- Funnel Transformer model with a binary classification head on top as used during pretraining for identifying
- generated tokens.
- """,
- FUNNEL_START_DOCSTRING,
-)
-
-
-class FunnelForPreTraining(FunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__(config)
-
- self.funnel = FunnelModel(config)
- self.discriminator_predictions = FunnelDiscriminatorPredictions(config)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=FunnelForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, FunnelForPreTrainingOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the ELECTRA-style loss. Input should be a sequence of tokens (see `input_ids`
- docstring) Indices should be in `[0, 1]`:
-
- - 0 indicates the token is an original token,
- - 1 indicates the token was replaced.
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoTokenizer, FunnelForPreTraining
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("funnel-transformer/small")
- >>> model = FunnelForPreTraining.from_pretrained("funnel-transformer/small")
-
- >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
- >>> logits = model(**inputs).logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- discriminator_hidden_states = self.funnel(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- discriminator_sequence_output = discriminator_hidden_states[0]
-
- logits = self.discriminator_predictions(discriminator_sequence_output)
-
- loss = None
- if labels is not None:
- loss_fct = nn.BCEWithLogitsLoss()
- if attention_mask is not None:
- active_loss = attention_mask.view(-1, discriminator_sequence_output.shape[1]) == 1
- active_logits = logits.view(-1, discriminator_sequence_output.shape[1])[active_loss]
- active_labels = labels[active_loss]
- loss = loss_fct(active_logits, active_labels.float())
- else:
- loss = loss_fct(logits.view(-1, discriminator_sequence_output.shape[1]), labels.float())
-
- if not return_dict:
- output = (logits,) + discriminator_hidden_states[1:]
- return ((loss,) + output) if loss is not None else output
-
- return FunnelForPreTrainingOutput(
- loss=loss,
- logits=logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
-
-@add_start_docstrings("""Funnel Transformer Model with a `language modeling` head on top.""", FUNNEL_START_DOCSTRING)
-class FunnelForMaskedLM(FunnelPreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__(config)
-
- self.funnel = FunnelModel(config)
- self.lm_head = nn.Linear(config.d_model, config.vocab_size)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self) -> nn.Linear:
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings: nn.Embedding) -> None:
- self.lm_head = new_embeddings
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=MaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- mask="",
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MaskedLMOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.funnel(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden_state = outputs[0]
- prediction_logits = self.lm_head(last_hidden_state)
-
- masked_lm_loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss() # -100 index = padding token
- masked_lm_loss = loss_fct(prediction_logits.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (prediction_logits,) + outputs[1:]
- return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
-
- return MaskedLMOutput(
- loss=masked_lm_loss,
- logits=prediction_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Funnel Transformer Model with a sequence classification/regression head on top (two linear layer on top of the
- first timestep of the last hidden state) e.g. for GLUE tasks.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class FunnelForSequenceClassification(FunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__(config)
- self.num_labels = config.num_labels
- self.config = config
-
- self.funnel = FunnelBaseModel(config)
- self.classifier = FunnelClassificationHead(config, config.num_labels)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small-base",
- output_type=SequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.funnel(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden_state = outputs[0]
- pooled_output = last_hidden_state[:, 0]
- logits = self.classifier(pooled_output)
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(logits, labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Funnel Transformer Model with a multiple choice classification head on top (two linear layer on top of the first
- timestep of the last hidden state, and a softmax) e.g. for RocStories/SWAG tasks.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class FunnelForMultipleChoice(FunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__(config)
-
- self.funnel = FunnelBaseModel(config)
- self.classifier = FunnelClassificationHead(config, 1)
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small-base",
- output_type=MultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, MultipleChoiceModelOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
- num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
- `input_ids` above)
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
-
- input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
- attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
- token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
- inputs_embeds = (
- inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.funnel(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden_state = outputs[0]
- pooled_output = last_hidden_state[:, 0]
- logits = self.classifier(pooled_output)
- reshaped_logits = logits.view(-1, num_choices)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(reshaped_logits, labels)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return MultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Funnel Transformer Model with a token classification head on top (a linear layer on top of the hidden-states
- output) e.g. for Named-Entity-Recognition (NER) tasks.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class FunnelForTokenClassification(FunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.funnel = FunnelModel(config)
- self.dropout = nn.Dropout(config.hidden_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.funnel(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden_state = outputs[0]
- last_hidden_state = self.dropout(last_hidden_state)
- logits = self.classifier(last_hidden_state)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- Funnel Transformer Model with a span classification head on top for extractive question-answering tasks like SQuAD
- (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- FUNNEL_START_DOCSTRING,
-)
-class FunnelForQuestionAnswering(FunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig) -> None:
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.funnel = FunnelModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- start_positions: Optional[torch.Tensor] = None,
- end_positions: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.funnel(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden_state = outputs[0]
-
- logits = self.qa_outputs(last_hidden_state)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[1:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/funnel/modeling_tf_funnel.py b/transformers/models/funnel/modeling_tf_funnel.py
deleted file mode 100644
index b50b96df1c54083d552e94f8d0c5e219a62d3e65..0000000000000000000000000000000000000000
--- a/transformers/models/funnel/modeling_tf_funnel.py
+++ /dev/null
@@ -1,1871 +0,0 @@
-# coding=utf-8
-# Copyright 2020-present Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 Funnel model."""
-
-
-from __future__ import annotations
-
-import warnings
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutput,
- TFMaskedLMOutput,
- TFMultipleChoiceModelOutput,
- TFQuestionAnsweringModelOutput,
- TFSequenceClassifierOutput,
- TFTokenClassifierOutput,
-)
-from ...modeling_tf_utils import (
- TFMaskedLanguageModelingLoss,
- TFModelInputType,
- TFMultipleChoiceLoss,
- TFPreTrainedModel,
- TFQuestionAnsweringLoss,
- TFSequenceClassificationLoss,
- TFTokenClassificationLoss,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_funnel import FunnelConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "FunnelConfig"
-
-
-from ..deprecated._archive_maps import TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-INF = 1e6
-
-
-class TFFunnelEmbeddings(keras.layers.Layer):
- """Construct the embeddings from word, position and token_type embeddings."""
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.hidden_size = config.hidden_size
- self.initializer_std = 1.0 if config.initializer_std is None else config.initializer_std
-
- self.LayerNorm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
- self.dropout = keras.layers.Dropout(rate=config.hidden_dropout)
-
- def build(self, input_shape=None):
- with tf.name_scope("word_embeddings"):
- self.weight = self.add_weight(
- name="weight",
- shape=[self.config.vocab_size, self.hidden_size],
- initializer=get_initializer(initializer_range=self.initializer_std),
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "LayerNorm", None) is not None:
- with tf.name_scope(self.LayerNorm.name):
- self.LayerNorm.build([None, None, self.config.d_model])
-
- def call(self, input_ids=None, inputs_embeds=None, training=False):
- """
- Applies embedding based on inputs tensor.
-
- Returns:
- final_embeddings (`tf.Tensor`): output embedding tensor.
- """
- assert not (input_ids is None and inputs_embeds is None)
- assert not (input_ids is not None and inputs_embeds is not None)
-
- if input_ids is not None:
- check_embeddings_within_bounds(input_ids, self.config.vocab_size)
- inputs_embeds = tf.gather(self.weight, input_ids)
-
- final_embeddings = self.LayerNorm(inputs=inputs_embeds)
- final_embeddings = self.dropout(inputs=final_embeddings, training=training)
-
- return final_embeddings
-
-
-class TFFunnelAttentionStructure:
- """
- Contains helpers for `TFFunnelRelMultiheadAttention `.
- """
-
- cls_token_type_id: int = 2
-
- def __init__(self, config):
- self.d_model = config.d_model
- self.attention_type = config.attention_type
- self.num_blocks = config.num_blocks
- self.separate_cls = config.separate_cls
- self.truncate_seq = config.truncate_seq
- self.pool_q_only = config.pool_q_only
- self.pooling_type = config.pooling_type
-
- self.sin_dropout = keras.layers.Dropout(config.hidden_dropout)
- self.cos_dropout = keras.layers.Dropout(config.hidden_dropout)
- # Track where we are at in terms of pooling from the original input, e.g., by how much the sequence length was
- # divided.
- self.pooling_mult = None
-
- def init_attention_inputs(self, inputs_embeds, attention_mask=None, token_type_ids=None, training=False):
- """Returns the attention inputs associated to the inputs of the model."""
- # inputs_embeds has shape batch_size x seq_len x d_model
- # attention_mask and token_type_ids have shape batch_size x seq_len
- self.pooling_mult = 1
- self.seq_len = seq_len = shape_list(inputs_embeds)[1]
- position_embeds = self.get_position_embeds(seq_len, training=training)
- token_type_mat = self.token_type_ids_to_mat(token_type_ids) if token_type_ids is not None else None
- cls_mask = (
- tf.pad(tf.ones([seq_len - 1, seq_len - 1], dtype=inputs_embeds.dtype), [[1, 0], [1, 0]])
- if self.separate_cls
- else None
- )
- return (position_embeds, token_type_mat, attention_mask, cls_mask)
-
- def token_type_ids_to_mat(self, token_type_ids):
- """Convert `token_type_ids` to `token_type_mat`."""
- token_type_mat = tf.equal(tf.expand_dims(token_type_ids, -1), tf.expand_dims(token_type_ids, -2))
- # Treat as in the same segment as both A & B
- cls_ids = tf.equal(token_type_ids, tf.constant([self.cls_token_type_id], dtype=token_type_ids.dtype))
- cls_mat = tf.logical_or(tf.expand_dims(cls_ids, -1), tf.expand_dims(cls_ids, -2))
- return tf.logical_or(cls_mat, token_type_mat)
-
- def get_position_embeds(self, seq_len, training=False):
- """
- Create and cache inputs related to relative position encoding. Those are very different depending on whether we
- are using the factorized or the relative shift attention:
-
- For the factorized attention, it returns the matrices (phi, pi, psi, omega) used in the paper, appendix A.2.2,
- final formula.
-
- For the relative shift attention, it returns all possible vectors R used in the paper, appendix A.2.1, final
- formula.
-
- Paper link: https://arxiv.org/abs/2006.03236
- """
- if self.attention_type == "factorized":
- # Notations from the paper, appending A.2.2, final formula.
- # We need to create and return the matrices phi, psi, pi and omega.
- pos_seq = tf.range(0, seq_len, 1.0)
- freq_seq = tf.range(0, self.d_model // 2, 1.0)
- inv_freq = 1 / (10000 ** (freq_seq / (self.d_model // 2)))
- sinusoid = tf.einsum("i,d->id", pos_seq, inv_freq)
-
- sin_embed = tf.sin(sinusoid)
- sin_embed_d = self.sin_dropout(sin_embed, training=training)
- cos_embed = tf.cos(sinusoid)
- cos_embed_d = self.cos_dropout(cos_embed, training=training)
- # This is different from the formula on the paper...
- phi = tf.concat([sin_embed_d, sin_embed_d], axis=-1)
- psi = tf.concat([cos_embed, sin_embed], axis=-1)
- pi = tf.concat([cos_embed_d, cos_embed_d], axis=-1)
- omega = tf.concat([-sin_embed, cos_embed], axis=-1)
- return (phi, pi, psi, omega)
- else:
- # Notations from the paper, appending A.2.1, final formula.
- # We need to create and return all the possible vectors R for all blocks and shifts.
- freq_seq = tf.range(0, self.d_model // 2, 1.0)
- inv_freq = 1 / (10000 ** (freq_seq / (self.d_model // 2)))
- # Maximum relative positions for the first input
- rel_pos_id = tf.range(-seq_len * 2, seq_len * 2, 1.0)
- zero_offset = seq_len * tf.constant(2)
- sinusoid = tf.einsum("i,d->id", rel_pos_id, inv_freq)
- sin_embed = self.sin_dropout(tf.sin(sinusoid), training=training)
- cos_embed = self.cos_dropout(tf.cos(sinusoid), training=training)
- pos_embed = tf.concat([sin_embed, cos_embed], axis=-1)
-
- pos = tf.range(0, seq_len)
- pooled_pos = pos
- position_embeds_list = []
- for block_index in range(0, self.num_blocks):
- # For each block with block_index > 0, we need two types position embeddings:
- # - Attention(pooled-q, unpooled-kv)
- # - Attention(pooled-q, pooled-kv)
- # For block_index = 0 we only need the second one and leave the first one as None.
-
- # First type
- position_embeds_pooling = tf.fill([1], value=-1.0)
-
- if block_index != 0:
- pooled_pos = self.stride_pool_pos(pos, block_index)
-
- # construct rel_pos_id
- stride = 2 ** (block_index - 1)
- rel_pos = self.relative_pos(pos, stride, pooled_pos, shift=2)
- # rel_pos = tf.expand_dims(rel_pos,1) + zero_offset
- # rel_pos = tf.broadcast_to(rel_pos, (rel_pos.shape[0], self.d_model))
- rel_pos = tf.cast(rel_pos, dtype=zero_offset.dtype)
- rel_pos = rel_pos + zero_offset
- position_embeds_pooling = tf.gather(pos_embed, rel_pos, axis=0)
-
- # Second type
- pos = pooled_pos
- stride = 2**block_index
- rel_pos = self.relative_pos(pos, stride)
-
- # rel_pos = tf.expand_dims(rel_pos,1) + zero_offset
- # rel_pos = tf.broadcast_to(rel_pos, (rel_pos.shape[0], self.d_model))
- rel_pos = tf.cast(rel_pos, dtype=zero_offset.dtype)
- rel_pos = rel_pos + zero_offset
- tf.debugging.assert_less(rel_pos, tf.shape(pos_embed)[0])
- position_embeds_no_pooling = tf.gather(pos_embed, rel_pos, axis=0)
-
- position_embeds_list.append([position_embeds_no_pooling, position_embeds_pooling])
- return position_embeds_list
-
- def stride_pool_pos(self, pos_id, block_index):
- """
- Pool `pos_id` while keeping the cls token separate (if `self.separate_cls=True`).
- """
- if self.separate_cls:
- # Under separate , we treat the as the first token in
- # the previous block of the 1st real block. Since the 1st real
- # block always has position 1, the position of the previous block
- # will be at `1 - 2 ** block_index`.
- cls_pos = tf.constant([-(2**block_index) + 1], dtype=pos_id.dtype)
- pooled_pos_id = pos_id[1:-1] if self.truncate_seq else pos_id[1:]
- return tf.concat([cls_pos, pooled_pos_id[::2]], 0)
- else:
- return pos_id[::2]
-
- def relative_pos(self, pos, stride, pooled_pos=None, shift=1):
- """
- Build the relative positional vector between `pos` and `pooled_pos`.
- """
- if pooled_pos is None:
- pooled_pos = pos
-
- ref_point = pooled_pos[0] - pos[0]
- num_remove = shift * shape_list(pooled_pos)[0]
- max_dist = ref_point + num_remove * stride
- min_dist = pooled_pos[0] - pos[-1]
-
- return tf.range(max_dist, min_dist - 1, -stride)
-
- def stride_pool(self, tensor, axis):
- """
- Perform pooling by stride slicing the tensor along the given axis.
- """
- if tensor is None:
- return None
-
- # Do the stride pool recursively if axis is a list or a tuple of ints.
- if isinstance(axis, (list, tuple)):
- for ax in axis:
- tensor = self.stride_pool(tensor, ax)
- return tensor
-
- # Do the stride pool recursively if tensor is a list or tuple of tensors.
- if isinstance(tensor, (tuple, list)):
- return type(tensor)(self.stride_pool(x, axis) for x in tensor)
-
- # Deal with negative axis
- axis %= len(shape_list(tensor))
-
- axis_slice = slice(None, -1, 2) if self.separate_cls and self.truncate_seq else slice(None, None, 2)
- enc_slice = [slice(None)] * axis + [axis_slice]
- if self.separate_cls:
- cls_slice = [slice(None)] * axis + [slice(None, 1)]
- tensor = tf.concat([tensor[cls_slice], tensor], axis)
- return tensor[enc_slice]
-
- def pool_tensor(self, tensor, mode="mean", stride=2):
- """Apply 1D pooling to a tensor of size [B x T (x H)]."""
- if tensor is None:
- return None
-
- # Do the pool recursively if tensor is a list or tuple of tensors.
- if isinstance(tensor, (tuple, list)):
- return type(tensor)(self.pool_tensor(tensor, mode=mode, stride=stride) for x in tensor)
-
- if self.separate_cls:
- suffix = tensor[:, :-1] if self.truncate_seq else tensor
- tensor = tf.concat([tensor[:, :1], suffix], axis=1)
-
- ndim = len(shape_list(tensor))
- if ndim == 2:
- tensor = tensor[:, :, None]
-
- if mode == "mean":
- tensor = tf.nn.avg_pool1d(tensor, stride, strides=stride, data_format="NWC", padding="SAME")
- elif mode == "max":
- tensor = tf.nn.max_pool1d(tensor, stride, strides=stride, data_format="NWC", padding="SAME")
- elif mode == "min":
- tensor = -tf.nn.max_pool1d(-tensor, stride, strides=stride, data_format="NWC", padding="SAME")
- else:
- raise NotImplementedError("The supported modes are 'mean', 'max' and 'min'.")
-
- return tf.squeeze(tensor, 2) if ndim == 2 else tensor
-
- def pre_attention_pooling(self, output, attention_inputs):
- """Pool `output` and the proper parts of `attention_inputs` before the attention layer."""
- position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
- if self.pool_q_only:
- if self.attention_type == "factorized":
- position_embeds = self.stride_pool(position_embeds[:2], 0) + position_embeds[2:]
- token_type_mat = self.stride_pool(token_type_mat, 1)
- cls_mask = self.stride_pool(cls_mask, 0)
- output = self.pool_tensor(output, mode=self.pooling_type)
- else:
- self.pooling_mult *= 2
- if self.attention_type == "factorized":
- position_embeds = self.stride_pool(position_embeds, 0)
- token_type_mat = self.stride_pool(token_type_mat, [1, 2])
- cls_mask = self.stride_pool(cls_mask, [1, 2])
- attention_mask = self.pool_tensor(attention_mask, mode="min")
- output = self.pool_tensor(output, mode=self.pooling_type)
- attention_inputs = (position_embeds, token_type_mat, attention_mask, cls_mask)
- return output, attention_inputs
-
- def post_attention_pooling(self, attention_inputs):
- """Pool the proper parts of `attention_inputs` after the attention layer."""
- position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
- if self.pool_q_only:
- self.pooling_mult *= 2
- if self.attention_type == "factorized":
- position_embeds = position_embeds[:2] + self.stride_pool(position_embeds[2:], 0)
- token_type_mat = self.stride_pool(token_type_mat, 2)
- cls_mask = self.stride_pool(cls_mask, 1)
- attention_mask = self.pool_tensor(attention_mask, mode="min")
- attention_inputs = (position_embeds, token_type_mat, attention_mask, cls_mask)
- return attention_inputs
-
-
-def _relative_shift_gather(positional_attn, context_len, shift):
- batch_size, n_head, seq_len, max_rel_len = shape_list(positional_attn)
- # max_rel_len = 2 * context_len + shift -1 is the numbers of possible relative positions i-j
-
- # What's next is the same as doing the following gather in PyTorch, which might be clearer code but less efficient.
- # idxs = context_len + torch.arange(0, context_len).unsqueeze(0) - torch.arange(0, seq_len).unsqueeze(1)
- # # matrix of context_len + i-j
- # return positional_attn.gather(3, idxs.expand([batch_size, n_head, context_len, context_len]))
-
- positional_attn = tf.reshape(positional_attn, [batch_size, n_head, max_rel_len, seq_len])
- positional_attn = positional_attn[:, :, shift:, :]
- positional_attn = tf.reshape(positional_attn, [batch_size, n_head, seq_len, max_rel_len - shift])
- positional_attn = positional_attn[..., :context_len]
- return positional_attn
-
-
-class TFFunnelRelMultiheadAttention(keras.layers.Layer):
- def __init__(self, config, block_index, **kwargs):
- super().__init__(**kwargs)
- self.attention_type = config.attention_type
- self.n_head = n_head = config.n_head
- self.d_head = d_head = config.d_head
- self.d_model = d_model = config.d_model
- self.initializer_range = config.initializer_range
- self.block_index = block_index
-
- self.hidden_dropout = keras.layers.Dropout(config.hidden_dropout)
- self.attention_dropout = keras.layers.Dropout(config.attention_dropout)
-
- initializer = get_initializer(config.initializer_range)
-
- self.q_head = keras.layers.Dense(
- n_head * d_head, use_bias=False, kernel_initializer=initializer, name="q_head"
- )
- self.k_head = keras.layers.Dense(n_head * d_head, kernel_initializer=initializer, name="k_head")
- self.v_head = keras.layers.Dense(n_head * d_head, kernel_initializer=initializer, name="v_head")
-
- self.post_proj = keras.layers.Dense(d_model, kernel_initializer=initializer, name="post_proj")
- self.layer_norm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
- self.scale = 1.0 / (d_head**0.5)
-
- def build(self, input_shape=None):
- n_head, d_head, d_model = self.n_head, self.d_head, self.d_model
- initializer = get_initializer(self.initializer_range)
-
- self.r_w_bias = self.add_weight(
- shape=(n_head, d_head), initializer=initializer, trainable=True, name="r_w_bias"
- )
- self.r_r_bias = self.add_weight(
- shape=(n_head, d_head), initializer=initializer, trainable=True, name="r_r_bias"
- )
- self.r_kernel = self.add_weight(
- shape=(d_model, n_head, d_head), initializer=initializer, trainable=True, name="r_kernel"
- )
- self.r_s_bias = self.add_weight(
- shape=(n_head, d_head), initializer=initializer, trainable=True, name="r_s_bias"
- )
- self.seg_embed = self.add_weight(
- shape=(2, n_head, d_head), initializer=initializer, trainable=True, name="seg_embed"
- )
-
- if self.built:
- return
- self.built = True
- if getattr(self, "q_head", None) is not None:
- with tf.name_scope(self.q_head.name):
- self.q_head.build([None, None, d_model])
- if getattr(self, "k_head", None) is not None:
- with tf.name_scope(self.k_head.name):
- self.k_head.build([None, None, d_model])
- if getattr(self, "v_head", None) is not None:
- with tf.name_scope(self.v_head.name):
- self.v_head.build([None, None, d_model])
- if getattr(self, "post_proj", None) is not None:
- with tf.name_scope(self.post_proj.name):
- self.post_proj.build([None, None, n_head * d_head])
- if getattr(self, "layer_norm", None) is not None:
- with tf.name_scope(self.layer_norm.name):
- self.layer_norm.build([None, None, d_model])
-
- def relative_positional_attention(self, position_embeds, q_head, context_len, cls_mask=None):
- """Relative attention score for the positional encodings"""
- # q_head has shape batch_size x sea_len x n_head x d_head
- if self.attention_type == "factorized":
- # Notations from the paper, appending A.2.2, final formula (https://arxiv.org/abs/2006.03236)
- # phi and pi have shape seq_len x d_model, psi and omega have shape context_len x d_model
- phi, pi, psi, omega = position_embeds
- # Shape n_head x d_head
- u = self.r_r_bias * self.scale
- # Shape d_model x n_head x d_head
- w_r = self.r_kernel
-
- # Shape batch_size x sea_len x n_head x d_model
- q_r_attention = tf.einsum("binh,dnh->bind", q_head + u, w_r)
- q_r_attention_1 = q_r_attention * phi[:, None]
- q_r_attention_2 = q_r_attention * pi[:, None]
-
- # Shape batch_size x n_head x seq_len x context_len
- positional_attn = tf.einsum("bind,jd->bnij", q_r_attention_1, psi) + tf.einsum(
- "bind,jd->bnij", q_r_attention_2, omega
- )
- else:
- # Notations from the paper, appending A.2.1, final formula (https://arxiv.org/abs/2006.03236)
- # Grab the proper positional encoding, shape max_rel_len x d_model
- if shape_list(q_head)[1] != context_len:
- shift = 2
- r = position_embeds[self.block_index][1]
- else:
- shift = 1
- r = position_embeds[self.block_index][0]
- # Shape n_head x d_head
- v = self.r_r_bias * self.scale
- # Shape d_model x n_head x d_head
- w_r = self.r_kernel
-
- # Shape max_rel_len x n_head x d_model
- r_head = tf.einsum("td,dnh->tnh", r, w_r)
- # Shape batch_size x n_head x seq_len x max_rel_len
- positional_attn = tf.einsum("binh,tnh->bnit", q_head + v, r_head)
- # Shape batch_size x n_head x seq_len x context_len
- positional_attn = _relative_shift_gather(positional_attn, context_len, shift)
-
- if cls_mask is not None:
- positional_attn *= cls_mask
- return positional_attn
-
- def relative_token_type_attention(self, token_type_mat, q_head, cls_mask=None):
- """Relative attention score for the token_type_ids"""
- if token_type_mat is None:
- return 0
- batch_size, seq_len, context_len = shape_list(token_type_mat)
- # q_head has shape batch_size x seq_len x n_head x d_head
- # Shape n_head x d_head
- r_s_bias = self.r_s_bias * self.scale
-
- # Shape batch_size x n_head x seq_len x 2
- token_type_bias = tf.einsum("bind,snd->bnis", q_head + r_s_bias, self.seg_embed)
- # Shape batch_size x n_head x seq_len x context_len
- token_type_mat = tf.tile(token_type_mat[:, None], [1, shape_list(q_head)[2], 1, 1])
- # token_type_mat = tf.broadcast_to(token_type_mat[:, None], new_shape)
- # Shapes batch_size x n_head x seq_len
- diff_token_type, same_token_type = tf.split(token_type_bias, 2, axis=-1)
- # Shape batch_size x n_head x seq_len x context_len
- token_type_attn = tf.where(
- token_type_mat,
- tf.tile(same_token_type, [1, 1, 1, context_len]),
- tf.tile(diff_token_type, [1, 1, 1, context_len]),
- )
-
- if cls_mask is not None:
- token_type_attn *= cls_mask
- return token_type_attn
-
- def call(self, query, key, value, attention_inputs, output_attentions=False, training=False):
- # query has shape batch_size x seq_len x d_model
- # key and value have shapes batch_size x context_len x d_model
- position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
-
- batch_size, seq_len, _ = shape_list(query)
- context_len = shape_list(key)[1]
- n_head, d_head = self.n_head, self.d_head
-
- # Shape batch_size x seq_len x n_head x d_head
- q_head = tf.reshape(self.q_head(query), [batch_size, seq_len, n_head, d_head])
- # Shapes batch_size x context_len x n_head x d_head
- k_head = tf.reshape(self.k_head(key), [batch_size, context_len, n_head, d_head])
- v_head = tf.reshape(self.v_head(value), [batch_size, context_len, n_head, d_head])
-
- q_head = q_head * self.scale
- # Shape n_head x d_head
- r_w_bias = self.r_w_bias * self.scale
- # Shapes batch_size x n_head x seq_len x context_len
- content_score = tf.einsum("bind,bjnd->bnij", q_head + r_w_bias, k_head)
- positional_attn = self.relative_positional_attention(position_embeds, q_head, context_len, cls_mask)
- token_type_attn = self.relative_token_type_attention(token_type_mat, q_head, cls_mask)
-
- # merge attention scores
- attn_score = content_score + positional_attn + token_type_attn
-
- # perform masking
- if attention_mask is not None:
- attention_mask = tf.cast(attention_mask, dtype=attn_score.dtype)
- attn_score = attn_score - (INF * (1 - attention_mask[:, None, None]))
-
- # attention probability
- attn_prob = stable_softmax(attn_score, axis=-1)
- attn_prob = self.attention_dropout(attn_prob, training=training)
-
- # attention output, shape batch_size x seq_len x n_head x d_head
- attn_vec = tf.einsum("bnij,bjnd->bind", attn_prob, v_head)
-
- # Shape shape batch_size x seq_len x d_model
- attn_out = self.post_proj(tf.reshape(attn_vec, [batch_size, seq_len, n_head * d_head]))
- attn_out = self.hidden_dropout(attn_out, training=training)
-
- output = self.layer_norm(query + attn_out)
- return (output, attn_prob) if output_attentions else (output,)
-
-
-class TFFunnelPositionwiseFFN(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- initializer = get_initializer(config.initializer_range)
- self.linear_1 = keras.layers.Dense(config.d_inner, kernel_initializer=initializer, name="linear_1")
- self.activation_function = get_tf_activation(config.hidden_act)
- self.activation_dropout = keras.layers.Dropout(config.activation_dropout)
- self.linear_2 = keras.layers.Dense(config.d_model, kernel_initializer=initializer, name="linear_2")
- self.dropout = keras.layers.Dropout(config.hidden_dropout)
- self.layer_norm = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
- self.config = config
-
- def call(self, hidden, training=False):
- h = self.linear_1(hidden)
- h = self.activation_function(h)
- h = self.activation_dropout(h, training=training)
- h = self.linear_2(h)
- h = self.dropout(h, training=training)
- return self.layer_norm(hidden + h)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "linear_1", None) is not None:
- with tf.name_scope(self.linear_1.name):
- self.linear_1.build([None, None, self.config.d_model])
- if getattr(self, "linear_2", None) is not None:
- with tf.name_scope(self.linear_2.name):
- self.linear_2.build([None, None, self.config.d_inner])
- if getattr(self, "layer_norm", None) is not None:
- with tf.name_scope(self.layer_norm.name):
- self.layer_norm.build([None, None, self.config.d_model])
-
-
-class TFFunnelLayer(keras.layers.Layer):
- def __init__(self, config, block_index, **kwargs):
- super().__init__(**kwargs)
- self.attention = TFFunnelRelMultiheadAttention(config, block_index, name="attention")
- self.ffn = TFFunnelPositionwiseFFN(config, name="ffn")
-
- def call(self, query, key, value, attention_inputs, output_attentions=False, training=False):
- attn = self.attention(
- query, key, value, attention_inputs, output_attentions=output_attentions, training=training
- )
- output = self.ffn(attn[0], training=training)
- return (output, attn[1]) if output_attentions else (output,)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "attention", None) is not None:
- with tf.name_scope(self.attention.name):
- self.attention.build(None)
- if getattr(self, "ffn", None) is not None:
- with tf.name_scope(self.ffn.name):
- self.ffn.build(None)
-
-
-class TFFunnelEncoder(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- self.separate_cls = config.separate_cls
- self.pool_q_only = config.pool_q_only
- self.block_repeats = config.block_repeats
- self.attention_structure = TFFunnelAttentionStructure(config)
- self.blocks = [
- [TFFunnelLayer(config, block_index, name=f"blocks_._{block_index}_._{i}") for i in range(block_size)]
- for block_index, block_size in enumerate(config.block_sizes)
- ]
-
- def call(
- self,
- inputs_embeds,
- attention_mask=None,
- token_type_ids=None,
- output_attentions=False,
- output_hidden_states=False,
- return_dict=True,
- training=False,
- ):
- # The pooling is not implemented on long tensors, so we convert this mask.
- # attention_mask = tf.cast(attention_mask, inputs_embeds.dtype)
- attention_inputs = self.attention_structure.init_attention_inputs(
- inputs_embeds,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- training=training,
- )
- hidden = inputs_embeds
-
- all_hidden_states = (inputs_embeds,) if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- for block_index, block in enumerate(self.blocks):
- pooling_flag = shape_list(hidden)[1] > (2 if self.separate_cls else 1)
- pooling_flag = pooling_flag and block_index > 0
- pooled_hidden = tf.zeros(shape_list(hidden))
-
- if pooling_flag:
- pooled_hidden, attention_inputs = self.attention_structure.pre_attention_pooling(
- hidden, attention_inputs
- )
-
- for layer_index, layer in enumerate(block):
- for repeat_index in range(self.block_repeats[block_index]):
- do_pooling = (repeat_index == 0) and (layer_index == 0) and pooling_flag
- if do_pooling:
- query = pooled_hidden
- key = value = hidden if self.pool_q_only else pooled_hidden
- else:
- query = key = value = hidden
- layer_output = layer(
- query, key, value, attention_inputs, output_attentions=output_attentions, training=training
- )
- hidden = layer_output[0]
- if do_pooling:
- attention_inputs = self.attention_structure.post_attention_pooling(attention_inputs)
-
- if output_attentions:
- all_attentions = all_attentions + layer_output[1:]
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden,)
-
- if not return_dict:
- return tuple(v for v in [hidden, all_hidden_states, all_attentions] if v is not None)
- return TFBaseModelOutput(last_hidden_state=hidden, hidden_states=all_hidden_states, attentions=all_attentions)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- for block in self.blocks:
- for layer in block:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-def upsample(x, stride, target_len, separate_cls=True, truncate_seq=False):
- """
- Upsample tensor `x` to match `target_len` by repeating the tokens `stride` time on the sequence length dimension.
- """
- if stride == 1:
- return x
- if separate_cls:
- cls = x[:, :1]
- x = x[:, 1:]
- output = tf.repeat(x, repeats=stride, axis=1)
- if separate_cls:
- if truncate_seq:
- output = tf.pad(output, [[0, 0], [0, stride - 1], [0, 0]])
- output = output[:, : target_len - 1]
- output = tf.concat([cls, output], axis=1)
- else:
- output = output[:, :target_len]
- return output
-
-
-class TFFunnelDecoder(keras.layers.Layer):
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- self.separate_cls = config.separate_cls
- self.truncate_seq = config.truncate_seq
- self.stride = 2 ** (len(config.block_sizes) - 1)
- self.attention_structure = TFFunnelAttentionStructure(config)
- self.layers = [TFFunnelLayer(config, 0, name=f"layers_._{i}") for i in range(config.num_decoder_layers)]
-
- def call(
- self,
- final_hidden,
- first_block_hidden,
- attention_mask=None,
- token_type_ids=None,
- output_attentions=False,
- output_hidden_states=False,
- return_dict=True,
- training=False,
- ):
- upsampled_hidden = upsample(
- final_hidden,
- stride=self.stride,
- target_len=shape_list(first_block_hidden)[1],
- separate_cls=self.separate_cls,
- truncate_seq=self.truncate_seq,
- )
-
- hidden = upsampled_hidden + first_block_hidden
- all_hidden_states = (hidden,) if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- attention_inputs = self.attention_structure.init_attention_inputs(
- hidden,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- training=training,
- )
-
- for layer in self.layers:
- layer_output = layer(
- hidden, hidden, hidden, attention_inputs, output_attentions=output_attentions, training=training
- )
- hidden = layer_output[0]
-
- if output_attentions:
- all_attentions = all_attentions + layer_output[1:]
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden,)
-
- if not return_dict:
- return tuple(v for v in [hidden, all_hidden_states, all_attentions] if v is not None)
- return TFBaseModelOutput(last_hidden_state=hidden, hidden_states=all_hidden_states, attentions=all_attentions)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "layers", None) is not None:
- for layer in self.layers:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-@keras_serializable
-class TFFunnelBaseLayer(keras.layers.Layer):
- """Base model without decoder"""
-
- config_class = FunnelConfig
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.output_attentions = config.output_attentions
- self.output_hidden_states = config.output_hidden_states
- self.return_dict = config.use_return_dict
-
- self.embeddings = TFFunnelEmbeddings(config, name="embeddings")
- self.encoder = TFFunnelEncoder(config, name="encoder")
-
- def get_input_embeddings(self):
- return self.embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.weight = value
- self.embeddings.vocab_size = shape_list(value)[0]
-
- def _prune_heads(self, heads_to_prune):
- raise NotImplementedError # Not implemented yet in the library fr TF 2.0 models
-
- @unpack_inputs
- def call(
- self,
- input_ids=None,
- attention_mask=None,
- token_type_ids=None,
- inputs_embeds=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- training=False,
- ):
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = tf.fill(input_shape, 1)
-
- if token_type_ids is None:
- token_type_ids = tf.fill(input_shape, 0)
-
- if inputs_embeds is None:
- inputs_embeds = self.embeddings(input_ids, training=training)
-
- encoder_outputs = self.encoder(
- inputs_embeds,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return encoder_outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
-
-
-@keras_serializable
-class TFFunnelMainLayer(keras.layers.Layer):
- """Base model with decoder"""
-
- config_class = FunnelConfig
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
-
- self.config = config
- self.block_sizes = config.block_sizes
- self.output_attentions = config.output_attentions
- self.output_hidden_states = config.output_hidden_states
- self.return_dict = config.use_return_dict
-
- self.embeddings = TFFunnelEmbeddings(config, name="embeddings")
- self.encoder = TFFunnelEncoder(config, name="encoder")
- self.decoder = TFFunnelDecoder(config, name="decoder")
-
- def get_input_embeddings(self):
- return self.embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.weight = value
- self.embeddings.vocab_size = shape_list(value)[0]
-
- def _prune_heads(self, heads_to_prune):
- raise NotImplementedError # Not implemented yet in the library fr TF 2.0 models
-
- @unpack_inputs
- def call(
- self,
- input_ids=None,
- attention_mask=None,
- token_type_ids=None,
- inputs_embeds=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- training=False,
- ):
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if attention_mask is None:
- attention_mask = tf.fill(input_shape, 1)
-
- if token_type_ids is None:
- token_type_ids = tf.fill(input_shape, 0)
-
- if inputs_embeds is None:
- inputs_embeds = self.embeddings(input_ids, training=training)
-
- encoder_outputs = self.encoder(
- inputs_embeds,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- output_attentions=output_attentions,
- output_hidden_states=True,
- return_dict=return_dict,
- training=training,
- )
-
- decoder_outputs = self.decoder(
- final_hidden=encoder_outputs[0],
- first_block_hidden=encoder_outputs[1][self.block_sizes[0]],
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- if not return_dict:
- idx = 0
- outputs = (decoder_outputs[0],)
- if output_hidden_states:
- idx += 1
- outputs = outputs + (encoder_outputs[1] + decoder_outputs[idx],)
- if output_attentions:
- idx += 1
- outputs = outputs + (encoder_outputs[2] + decoder_outputs[idx],)
- return outputs
-
- return TFBaseModelOutput(
- last_hidden_state=decoder_outputs[0],
- hidden_states=(encoder_outputs.hidden_states + decoder_outputs.hidden_states)
- if output_hidden_states
- else None,
- attentions=(encoder_outputs.attentions + decoder_outputs.attentions) if output_attentions else None,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "embeddings", None) is not None:
- with tf.name_scope(self.embeddings.name):
- self.embeddings.build(None)
- if getattr(self, "encoder", None) is not None:
- with tf.name_scope(self.encoder.name):
- self.encoder.build(None)
- if getattr(self, "decoder", None) is not None:
- with tf.name_scope(self.decoder.name):
- self.decoder.build(None)
-
-
-class TFFunnelDiscriminatorPredictions(keras.layers.Layer):
- """Prediction module for the discriminator, made up of two dense layers."""
-
- def __init__(self, config, **kwargs):
- super().__init__(**kwargs)
- initializer = get_initializer(config.initializer_range)
- self.dense = keras.layers.Dense(config.d_model, kernel_initializer=initializer, name="dense")
- self.activation_function = get_tf_activation(config.hidden_act)
- self.dense_prediction = keras.layers.Dense(1, kernel_initializer=initializer, name="dense_prediction")
- self.config = config
-
- def call(self, discriminator_hidden_states):
- hidden_states = self.dense(discriminator_hidden_states)
- hidden_states = self.activation_function(hidden_states)
- logits = tf.squeeze(self.dense_prediction(hidden_states))
- return logits
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "dense", None) is not None:
- with tf.name_scope(self.dense.name):
- self.dense.build([None, None, self.config.d_model])
- if getattr(self, "dense_prediction", None) is not None:
- with tf.name_scope(self.dense_prediction.name):
- self.dense_prediction.build([None, None, self.config.d_model])
-
-
-class TFFunnelMaskedLMHead(keras.layers.Layer):
- def __init__(self, config, input_embeddings, **kwargs):
- super().__init__(**kwargs)
- self.config = config
- self.hidden_size = config.hidden_size
- self.input_embeddings = input_embeddings
-
- def build(self, input_shape):
- self.bias = self.add_weight(shape=(self.config.vocab_size,), initializer="zeros", trainable=True, name="bias")
-
- super().build(input_shape)
-
- def get_output_embeddings(self):
- return self.input_embeddings
-
- def set_output_embeddings(self, value):
- self.input_embeddings.weight = value
- self.input_embeddings.vocab_size = shape_list(value)[0]
-
- def get_bias(self):
- return {"bias": self.bias}
-
- def set_bias(self, value):
- self.bias = value["bias"]
- self.config.vocab_size = shape_list(value["bias"])[0]
-
- def call(self, hidden_states, training=False):
- seq_length = shape_list(tensor=hidden_states)[1]
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, self.hidden_size])
- hidden_states = tf.matmul(a=hidden_states, b=self.input_embeddings.weight, transpose_b=True)
- hidden_states = tf.reshape(tensor=hidden_states, shape=[-1, seq_length, self.config.vocab_size])
- hidden_states = tf.nn.bias_add(value=hidden_states, bias=self.bias)
-
- return hidden_states
-
-
-class TFFunnelClassificationHead(keras.layers.Layer):
- def __init__(self, config, n_labels, **kwargs):
- super().__init__(**kwargs)
- initializer = get_initializer(config.initializer_range)
- self.linear_hidden = keras.layers.Dense(config.d_model, kernel_initializer=initializer, name="linear_hidden")
- self.dropout = keras.layers.Dropout(config.hidden_dropout)
- self.linear_out = keras.layers.Dense(n_labels, kernel_initializer=initializer, name="linear_out")
- self.config = config
-
- def call(self, hidden, training=False):
- hidden = self.linear_hidden(hidden)
- hidden = keras.activations.tanh(hidden)
- hidden = self.dropout(hidden, training=training)
- return self.linear_out(hidden)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "linear_hidden", None) is not None:
- with tf.name_scope(self.linear_hidden.name):
- self.linear_hidden.build([None, None, self.config.d_model])
- if getattr(self, "linear_out", None) is not None:
- with tf.name_scope(self.linear_out.name):
- self.linear_out.build([None, None, self.config.d_model])
-
-
-class TFFunnelPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = FunnelConfig
- base_model_prefix = "funnel"
-
- @property
- def dummy_inputs(self):
- # Funnel misbehaves with very small inputs, so we override and make them a bit bigger
- return {"input_ids": tf.ones((1, 3), dtype=tf.int32)}
-
-
-@dataclass
-class TFFunnelForPreTrainingOutput(ModelOutput):
- """
- Output type of [`FunnelForPreTraining`].
-
- Args:
- logits (`tf.Tensor` of shape `(batch_size, sequence_length)`):
- Prediction scores of the head (scores for each token before SoftMax).
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- logits: tf.Tensor = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-FUNNEL_START_DOCSTRING = r"""
-
- The Funnel Transformer model was proposed in [Funnel-Transformer: Filtering out Sequential Redundancy for Efficient
- Language Processing](https://arxiv.org/abs/2006.03236) by Zihang Dai, Guokun Lai, Yiming Yang, Quoc V. Le.
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`XxxConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-FUNNEL_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- inputs_embeds (`tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- """
- The base Funnel Transformer Model transformer outputting raw hidden-states without upsampling head (also called
- decoder) or any task-specific head on top.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class TFFunnelBaseModel(TFFunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig, *inputs, **kwargs) -> None:
- super().__init__(config, *inputs, **kwargs)
- self.funnel = TFFunnelBaseLayer(config, name="funnel")
-
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small-base",
- output_type=TFBaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- @unpack_inputs
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[Tuple[tf.Tensor], TFBaseModelOutput]:
- return self.funnel(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- def serving_output(self, output):
- # hidden_states and attentions not converted to Tensor with tf.convert_to_tensor as they are all of
- # different dimensions
- return TFBaseModelOutput(
- last_hidden_state=output.last_hidden_state,
- hidden_states=output.hidden_states,
- attentions=output.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "funnel", None) is not None:
- with tf.name_scope(self.funnel.name):
- self.funnel.build(None)
-
-
-@add_start_docstrings(
- "The bare Funnel Transformer Model transformer outputting raw hidden-states without any specific head on top.",
- FUNNEL_START_DOCSTRING,
-)
-class TFFunnelModel(TFFunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig, *inputs, **kwargs) -> None:
- super().__init__(config, *inputs, **kwargs)
- self.funnel = TFFunnelMainLayer(config, name="funnel")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small",
- output_type=TFBaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- ) -> Union[Tuple[tf.Tensor], TFBaseModelOutput]:
- return self.funnel(
- input_ids=input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- def serving_output(self, output):
- # hidden_states and attentions not converted to Tensor with tf.convert_to_tensor as they are all of
- # different dimensions
- return TFBaseModelOutput(
- last_hidden_state=output.last_hidden_state,
- hidden_states=output.hidden_states,
- attentions=output.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "funnel", None) is not None:
- with tf.name_scope(self.funnel.name):
- self.funnel.build(None)
-
-
-@add_start_docstrings(
- """
- Funnel model with a binary classification head on top as used during pretraining for identifying generated tokens.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class TFFunnelForPreTraining(TFFunnelPreTrainedModel):
- def __init__(self, config: FunnelConfig, **kwargs) -> None:
- super().__init__(config, **kwargs)
-
- self.funnel = TFFunnelMainLayer(config, name="funnel")
- self.discriminator_predictions = TFFunnelDiscriminatorPredictions(config, name="discriminator_predictions")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=TFFunnelForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: bool = False,
- **kwargs,
- ) -> Union[Tuple[tf.Tensor], TFFunnelForPreTrainingOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoTokenizer, TFFunnelForPreTraining
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("funnel-transformer/small")
- >>> model = TFFunnelForPreTraining.from_pretrained("funnel-transformer/small")
-
- >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="tf")
- >>> logits = model(inputs).logits
- ```"""
- discriminator_hidden_states = self.funnel(
- input_ids,
- attention_mask,
- token_type_ids,
- inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- discriminator_sequence_output = discriminator_hidden_states[0]
- logits = self.discriminator_predictions(discriminator_sequence_output)
-
- if not return_dict:
- return (logits,) + discriminator_hidden_states[1:]
-
- return TFFunnelForPreTrainingOutput(
- logits=logits,
- hidden_states=discriminator_hidden_states.hidden_states,
- attentions=discriminator_hidden_states.attentions,
- )
-
- def serving_output(self, output):
- # hidden_states and attentions not converted to Tensor with tf.convert_to_tensor as they are all of
- # different dimensions
- return TFFunnelForPreTrainingOutput(
- logits=output.logits, hidden_states=output.hidden_states, attentions=output.attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "funnel", None) is not None:
- with tf.name_scope(self.funnel.name):
- self.funnel.build(None)
- if getattr(self, "discriminator_predictions", None) is not None:
- with tf.name_scope(self.discriminator_predictions.name):
- self.discriminator_predictions.build(None)
-
-
-@add_start_docstrings("""Funnel Model with a `language modeling` head on top.""", FUNNEL_START_DOCSTRING)
-class TFFunnelForMaskedLM(TFFunnelPreTrainedModel, TFMaskedLanguageModelingLoss):
- def __init__(self, config: FunnelConfig, *inputs, **kwargs) -> None:
- super().__init__(config, *inputs, **kwargs)
-
- self.funnel = TFFunnelMainLayer(config, name="funnel")
- self.lm_head = TFFunnelMaskedLMHead(config, self.funnel.embeddings, name="lm_head")
-
- def get_lm_head(self) -> TFFunnelMaskedLMHead:
- return self.lm_head
-
- def get_prefix_bias_name(self) -> str:
- warnings.warn("The method get_prefix_bias_name is deprecated. Please use `get_bias` instead.", FutureWarning)
- return self.name + "/" + self.lm_head.name
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small",
- output_type=TFMaskedLMOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[Tuple[tf.Tensor], TFMaskedLMOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
- config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
- loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
- """
- outputs = self.funnel(
- input_ids,
- attention_mask,
- token_type_ids,
- inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
- prediction_scores = self.lm_head(sequence_output, training=training)
-
- loss = None if labels is None else self.hf_compute_loss(labels, prediction_scores)
-
- if not return_dict:
- output = (prediction_scores,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFMaskedLMOutput(
- loss=loss,
- logits=prediction_scores,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def serving_output(self, output: TFMaskedLMOutput) -> TFMaskedLMOutput:
- # hidden_states and attentions not converted to Tensor with tf.convert_to_tensor as they are all of
- # different dimensions
- return TFMaskedLMOutput(logits=output.logits, hidden_states=output.hidden_states, attentions=output.attentions)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "funnel", None) is not None:
- with tf.name_scope(self.funnel.name):
- self.funnel.build(None)
- if getattr(self, "lm_head", None) is not None:
- with tf.name_scope(self.lm_head.name):
- self.lm_head.build(None)
-
-
-@add_start_docstrings(
- """
- Funnel Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
- output) e.g. for GLUE tasks.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class TFFunnelForSequenceClassification(TFFunnelPreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config: FunnelConfig, *inputs, **kwargs) -> None:
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
-
- self.funnel = TFFunnelBaseLayer(config, name="funnel")
- self.classifier = TFFunnelClassificationHead(config, config.num_labels, name="classifier")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small-base",
- output_type=TFSequenceClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[Tuple[tf.Tensor], TFSequenceClassifierOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- outputs = self.funnel(
- input_ids,
- attention_mask,
- token_type_ids,
- inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- last_hidden_state = outputs[0]
- pooled_output = last_hidden_state[:, 0]
- logits = self.classifier(pooled_output, training=training)
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def serving_output(self, output: TFSequenceClassifierOutput) -> TFSequenceClassifierOutput:
- # hidden_states and attentions not converted to Tensor with tf.convert_to_tensor as they are all of
- # different dimensions
- return TFSequenceClassifierOutput(
- logits=output.logits, hidden_states=output.hidden_states, attentions=output.attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "funnel", None) is not None:
- with tf.name_scope(self.funnel.name):
- self.funnel.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build(None)
-
-
-@add_start_docstrings(
- """
- Funnel Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
- softmax) e.g. for RocStories/SWAG tasks.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class TFFunnelForMultipleChoice(TFFunnelPreTrainedModel, TFMultipleChoiceLoss):
- def __init__(self, config: FunnelConfig, *inputs, **kwargs) -> None:
- super().__init__(config, *inputs, **kwargs)
-
- self.funnel = TFFunnelBaseLayer(config, name="funnel")
- self.classifier = TFFunnelClassificationHead(config, 1, name="classifier")
-
- @property
- def dummy_inputs(self):
- return {"input_ids": tf.ones((3, 3, 4), dtype=tf.int32)}
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small-base",
- output_type=TFMultipleChoiceModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[Tuple[tf.Tensor], TFMultipleChoiceModelOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]`
- where `num_choices` is the size of the second dimension of the input tensors. (See `input_ids` above)
- """
- if input_ids is not None:
- num_choices = shape_list(input_ids)[1]
- seq_length = shape_list(input_ids)[2]
- else:
- num_choices = shape_list(inputs_embeds)[1]
- seq_length = shape_list(inputs_embeds)[2]
-
- flat_input_ids = tf.reshape(input_ids, (-1, seq_length)) if input_ids is not None else None
- flat_attention_mask = tf.reshape(attention_mask, (-1, seq_length)) if attention_mask is not None else None
- flat_token_type_ids = tf.reshape(token_type_ids, (-1, seq_length)) if token_type_ids is not None else None
- flat_inputs_embeds = (
- tf.reshape(inputs_embeds, (-1, seq_length, shape_list(inputs_embeds)[3]))
- if inputs_embeds is not None
- else None
- )
-
- outputs = self.funnel(
- flat_input_ids,
- attention_mask=flat_attention_mask,
- token_type_ids=flat_token_type_ids,
- inputs_embeds=flat_inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- last_hidden_state = outputs[0]
- pooled_output = last_hidden_state[:, 0]
- logits = self.classifier(pooled_output, training=training)
- reshaped_logits = tf.reshape(logits, (-1, num_choices))
-
- loss = None if labels is None else self.hf_compute_loss(labels, reshaped_logits)
-
- if not return_dict:
- output = (reshaped_logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFMultipleChoiceModelOutput(
- loss=loss,
- logits=reshaped_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def serving_output(self, output: TFMultipleChoiceModelOutput) -> TFMultipleChoiceModelOutput:
- # hidden_states and attentions not converted to Tensor with tf.convert_to_tensor as they are all of
- # different dimensions
- return TFMultipleChoiceModelOutput(
- logits=output.logits, hidden_states=output.hidden_states, attentions=output.attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "funnel", None) is not None:
- with tf.name_scope(self.funnel.name):
- self.funnel.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build(None)
-
-
-@add_start_docstrings(
- """
- Funnel Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- FUNNEL_START_DOCSTRING,
-)
-class TFFunnelForTokenClassification(TFFunnelPreTrainedModel, TFTokenClassificationLoss):
- def __init__(self, config: FunnelConfig, *inputs, **kwargs) -> None:
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
-
- self.funnel = TFFunnelMainLayer(config, name="funnel")
- self.dropout = keras.layers.Dropout(config.hidden_dropout)
- self.classifier = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small",
- output_type=TFTokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[Tuple[tf.Tensor], TFTokenClassifierOutput]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
- """
- outputs = self.funnel(
- input_ids,
- attention_mask,
- token_type_ids,
- inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
-
- sequence_output = self.dropout(sequence_output, training=training)
- logits = self.classifier(sequence_output)
-
- loss = None if labels is None else self.hf_compute_loss(labels, logits)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFTokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def serving_output(self, output: TFTokenClassifierOutput) -> TFTokenClassifierOutput:
- # hidden_states and attentions not converted to Tensor with tf.convert_to_tensor as they are all of
- # different dimensions
- return TFTokenClassifierOutput(
- logits=output.logits, hidden_states=output.hidden_states, attentions=output.attentions
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "funnel", None) is not None:
- with tf.name_scope(self.funnel.name):
- self.funnel.build(None)
- if getattr(self, "classifier", None) is not None:
- with tf.name_scope(self.classifier.name):
- self.classifier.build([None, None, self.config.hidden_size])
-
-
-@add_start_docstrings(
- """
- Funnel Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
- layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- FUNNEL_START_DOCSTRING,
-)
-class TFFunnelForQuestionAnswering(TFFunnelPreTrainedModel, TFQuestionAnsweringLoss):
- def __init__(self, config: FunnelConfig, *inputs, **kwargs) -> None:
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
-
- self.funnel = TFFunnelMainLayer(config, name="funnel")
- self.qa_outputs = keras.layers.Dense(
- config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="funnel-transformer/small",
- output_type=TFQuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- start_positions: np.ndarray | tf.Tensor | None = None,
- end_positions: np.ndarray | tf.Tensor | None = None,
- training: bool = False,
- ) -> Union[Tuple[tf.Tensor], TFQuestionAnsweringModelOutput]:
- r"""
- start_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
-
- outputs = self.funnel(
- input_ids,
- attention_mask,
- token_type_ids,
- inputs_embeds,
- output_attentions,
- output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = tf.split(logits, 2, axis=-1)
- start_logits = tf.squeeze(start_logits, axis=-1)
- end_logits = tf.squeeze(end_logits, axis=-1)
-
- loss = None
- if start_positions is not None and end_positions is not None:
- labels = {"start_position": start_positions, "end_position": end_positions}
- loss = self.hf_compute_loss(labels, (start_logits, end_logits))
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFQuestionAnsweringModelOutput(
- loss=loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def serving_output(self, output: TFQuestionAnsweringModelOutput) -> TFQuestionAnsweringModelOutput:
- # hidden_states and attentions not converted to Tensor with tf.convert_to_tensor as they are all of
- # different dimensions
- return TFQuestionAnsweringModelOutput(
- start_logits=output.start_logits,
- end_logits=output.end_logits,
- hidden_states=output.hidden_states,
- attentions=output.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "funnel", None) is not None:
- with tf.name_scope(self.funnel.name):
- self.funnel.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.config.hidden_size])
diff --git a/transformers/models/funnel/tokenization_funnel.py b/transformers/models/funnel/tokenization_funnel.py
deleted file mode 100644
index a1580deccfb3f75b4dcc3aa50573f542ec3a4117..0000000000000000000000000000000000000000
--- a/transformers/models/funnel/tokenization_funnel.py
+++ /dev/null
@@ -1,534 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Tokenization class for Funnel Transformer."""
-
-import collections
-import os
-import unicodedata
-from typing import List, Optional, Tuple
-
-from ...tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
-
-_model_names = [
- "small",
- "small-base",
- "medium",
- "medium-base",
- "intermediate",
- "intermediate-base",
- "large",
- "large-base",
- "xlarge",
- "xlarge-base",
-]
-
-
-# Copied from transformers.models.bert.tokenization_bert.load_vocab
-def load_vocab(vocab_file):
- """Loads a vocabulary file into a dictionary."""
- vocab = collections.OrderedDict()
- with open(vocab_file, "r", encoding="utf-8") as reader:
- tokens = reader.readlines()
- for index, token in enumerate(tokens):
- token = token.rstrip("\n")
- vocab[token] = index
- return vocab
-
-
-# Copied from transformers.models.bert.tokenization_bert.whitespace_tokenize
-def whitespace_tokenize(text):
- """Runs basic whitespace cleaning and splitting on a piece of text."""
- text = text.strip()
- if not text:
- return []
- tokens = text.split()
- return tokens
-
-
-class FunnelTokenizer(PreTrainedTokenizer):
- r"""
- Construct a Funnel Transformer tokenizer. Based on WordPiece.
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- do_basic_tokenize (`bool`, *optional*, defaults to `True`):
- Whether or not to do basic tokenization before WordPiece.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `""`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `""`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `""`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- bos_token (`str`, *optional*, defaults to `""`):
- The beginning of sentence token.
- eos_token (`str`, *optional*, defaults to `""`):
- The end of sentence token.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- cls_token_type_id: int = 2
-
- def __init__(
- self,
- vocab_file,
- do_lower_case=True,
- do_basic_tokenize=True,
- never_split=None,
- unk_token="",
- sep_token="",
- pad_token="",
- cls_token="",
- mask_token="",
- bos_token="",
- eos_token="",
- tokenize_chinese_chars=True,
- strip_accents=None,
- **kwargs,
- ):
- if not os.path.isfile(vocab_file):
- raise ValueError(
- f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
- " model use `tokenizer = FunnelTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- self.vocab = load_vocab(vocab_file)
- self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
- self.do_basic_tokenize = do_basic_tokenize
- if do_basic_tokenize:
- self.basic_tokenizer = BasicTokenizer(
- do_lower_case=do_lower_case,
- never_split=never_split,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- )
- self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token))
-
- super().__init__(
- do_lower_case=do_lower_case,
- do_basic_tokenize=do_basic_tokenize,
- never_split=never_split,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- bos_token=bos_token,
- eos_token=eos_token,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- **kwargs,
- )
-
- @property
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.do_lower_case
- def do_lower_case(self):
- return self.basic_tokenizer.do_lower_case
-
- @property
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.vocab_size
- def vocab_size(self):
- return len(self.vocab)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_vocab
- def get_vocab(self):
- return dict(self.vocab, **self.added_tokens_encoder)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._tokenize
- def _tokenize(self, text, split_special_tokens=False):
- split_tokens = []
- if self.do_basic_tokenize:
- for token in self.basic_tokenizer.tokenize(
- text, never_split=self.all_special_tokens if not split_special_tokens else None
- ):
- # If the token is part of the never_split set
- if token in self.basic_tokenizer.never_split:
- split_tokens.append(token)
- else:
- split_tokens += self.wordpiece_tokenizer.tokenize(token)
- else:
- split_tokens = self.wordpiece_tokenizer.tokenize(text)
- return split_tokens
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.vocab.get(token, self.vocab.get(self.unk_token))
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.ids_to_tokens.get(index, self.unk_token)
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.convert_tokens_to_string
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- out_string = " ".join(tokens).replace(" ##", "").strip()
- return out_string
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A BERT sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- if token_ids_1 is None:
- return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
- cls = [self.cls_token_id]
- sep = [self.sep_token_id]
- return cls + token_ids_0 + sep + token_ids_1 + sep
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_special_tokens_mask
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
-
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if token_ids_1 is not None:
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
- return [1] + ([0] * len(token_ids_0)) + [1]
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A Funnel
- Transformer sequence pair mask has the following format:
-
- ```
- 2 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls) * [self.cls_token_type_id] + len(token_ids_0 + sep) * [0]
- return len(cls) * [self.cls_token_type_id] + len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- index = 0
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- else:
- vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(token + "\n")
- index += 1
- return (vocab_file,)
-
-
-# Copied from transformers.models.bert.tokenization_bert.BasicTokenizer
-class BasicTokenizer(object):
- """
- Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
-
- Args:
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- never_split (`Iterable`, *optional*):
- Collection of tokens which will never be split during tokenization. Only has an effect when
- `do_basic_tokenize=True`
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters.
-
- This should likely be deactivated for Japanese (see this
- [issue](https://github.com/huggingface/transformers/issues/328)).
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- do_split_on_punc (`bool`, *optional*, defaults to `True`):
- In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
- the full context of the words, such as contractions.
- """
-
- def __init__(
- self,
- do_lower_case=True,
- never_split=None,
- tokenize_chinese_chars=True,
- strip_accents=None,
- do_split_on_punc=True,
- ):
- if never_split is None:
- never_split = []
- self.do_lower_case = do_lower_case
- self.never_split = set(never_split)
- self.tokenize_chinese_chars = tokenize_chinese_chars
- self.strip_accents = strip_accents
- self.do_split_on_punc = do_split_on_punc
-
- def tokenize(self, text, never_split=None):
- """
- Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
-
- Args:
- never_split (`List[str]`, *optional*)
- Kept for backward compatibility purposes. Now implemented directly at the base class level (see
- [`PreTrainedTokenizer.tokenize`]) List of token not to split.
- """
- # union() returns a new set by concatenating the two sets.
- never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
- text = self._clean_text(text)
-
- # This was added on November 1st, 2018 for the multilingual and Chinese
- # models. This is also applied to the English models now, but it doesn't
- # matter since the English models were not trained on any Chinese data
- # and generally don't have any Chinese data in them (there are Chinese
- # characters in the vocabulary because Wikipedia does have some Chinese
- # words in the English Wikipedia.).
- if self.tokenize_chinese_chars:
- text = self._tokenize_chinese_chars(text)
- # prevents treating the same character with different unicode codepoints as different characters
- unicode_normalized_text = unicodedata.normalize("NFC", text)
- orig_tokens = whitespace_tokenize(unicode_normalized_text)
- split_tokens = []
- for token in orig_tokens:
- if token not in never_split:
- if self.do_lower_case:
- token = token.lower()
- if self.strip_accents is not False:
- token = self._run_strip_accents(token)
- elif self.strip_accents:
- token = self._run_strip_accents(token)
- split_tokens.extend(self._run_split_on_punc(token, never_split))
-
- output_tokens = whitespace_tokenize(" ".join(split_tokens))
- return output_tokens
-
- def _run_strip_accents(self, text):
- """Strips accents from a piece of text."""
- text = unicodedata.normalize("NFD", text)
- output = []
- for char in text:
- cat = unicodedata.category(char)
- if cat == "Mn":
- continue
- output.append(char)
- return "".join(output)
-
- def _run_split_on_punc(self, text, never_split=None):
- """Splits punctuation on a piece of text."""
- if not self.do_split_on_punc or (never_split is not None and text in never_split):
- return [text]
- chars = list(text)
- i = 0
- start_new_word = True
- output = []
- while i < len(chars):
- char = chars[i]
- if _is_punctuation(char):
- output.append([char])
- start_new_word = True
- else:
- if start_new_word:
- output.append([])
- start_new_word = False
- output[-1].append(char)
- i += 1
-
- return ["".join(x) for x in output]
-
- def _tokenize_chinese_chars(self, text):
- """Adds whitespace around any CJK character."""
- output = []
- for char in text:
- cp = ord(char)
- if self._is_chinese_char(cp):
- output.append(" ")
- output.append(char)
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
- def _is_chinese_char(self, cp):
- """Checks whether CP is the codepoint of a CJK character."""
- # This defines a "chinese character" as anything in the CJK Unicode block:
- # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
- #
- # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
- # despite its name. The modern Korean Hangul alphabet is a different block,
- # as is Japanese Hiragana and Katakana. Those alphabets are used to write
- # space-separated words, so they are not treated specially and handled
- # like the all of the other languages.
- if (
- (cp >= 0x4E00 and cp <= 0x9FFF)
- or (cp >= 0x3400 and cp <= 0x4DBF) #
- or (cp >= 0x20000 and cp <= 0x2A6DF) #
- or (cp >= 0x2A700 and cp <= 0x2B73F) #
- or (cp >= 0x2B740 and cp <= 0x2B81F) #
- or (cp >= 0x2B820 and cp <= 0x2CEAF) #
- or (cp >= 0xF900 and cp <= 0xFAFF)
- or (cp >= 0x2F800 and cp <= 0x2FA1F) #
- ): #
- return True
-
- return False
-
- def _clean_text(self, text):
- """Performs invalid character removal and whitespace cleanup on text."""
- output = []
- for char in text:
- cp = ord(char)
- if cp == 0 or cp == 0xFFFD or _is_control(char):
- continue
- if _is_whitespace(char):
- output.append(" ")
- else:
- output.append(char)
- return "".join(output)
-
-
-# Copied from transformers.models.bert.tokenization_bert.WordpieceTokenizer
-class WordpieceTokenizer(object):
- """Runs WordPiece tokenization."""
-
- def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
- self.vocab = vocab
- self.unk_token = unk_token
- self.max_input_chars_per_word = max_input_chars_per_word
-
- def tokenize(self, text):
- """
- Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
- tokenization using the given vocabulary.
-
- For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
-
- Args:
- text: A single token or whitespace separated tokens. This should have
- already been passed through *BasicTokenizer*.
-
- Returns:
- A list of wordpiece tokens.
- """
-
- output_tokens = []
- for token in whitespace_tokenize(text):
- chars = list(token)
- if len(chars) > self.max_input_chars_per_word:
- output_tokens.append(self.unk_token)
- continue
-
- is_bad = False
- start = 0
- sub_tokens = []
- while start < len(chars):
- end = len(chars)
- cur_substr = None
- while start < end:
- substr = "".join(chars[start:end])
- if start > 0:
- substr = "##" + substr
- if substr in self.vocab:
- cur_substr = substr
- break
- end -= 1
- if cur_substr is None:
- is_bad = True
- break
- sub_tokens.append(cur_substr)
- start = end
-
- if is_bad:
- output_tokens.append(self.unk_token)
- else:
- output_tokens.extend(sub_tokens)
- return output_tokens
diff --git a/transformers/models/funnel/tokenization_funnel_fast.py b/transformers/models/funnel/tokenization_funnel_fast.py
deleted file mode 100644
index 9ff2a3bfefc57eb7dd67eb0e0d810cb492e87521..0000000000000000000000000000000000000000
--- a/transformers/models/funnel/tokenization_funnel_fast.py
+++ /dev/null
@@ -1,200 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Tokenization class for Funnel Transformer."""
-
-import json
-from typing import List, Optional, Tuple
-
-from tokenizers import normalizers
-
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import logging
-from .tokenization_funnel import FunnelTokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
-
-_model_names = [
- "small",
- "small-base",
- "medium",
- "medium-base",
- "intermediate",
- "intermediate-base",
- "large",
- "large-base",
- "xlarge",
- "xlarge-base",
-]
-
-
-class FunnelTokenizerFast(PreTrainedTokenizerFast):
- r"""
- Construct a "fast" Funnel Transformer tokenizer (backed by HuggingFace's *tokenizers* library). Based on WordPiece.
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- do_lower_case (`bool`, *optional*, defaults to `True`):
- Whether or not to lowercase the input when tokenizing.
- unk_token (`str`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- sep_token (`str`, *optional*, defaults to `""`):
- The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
- sequence classification or for a text and a question for question answering. It is also used as the last
- token of a sequence built with special tokens.
- pad_token (`str`, *optional*, defaults to `""`):
- The token used for padding, for example when batching sequences of different lengths.
- cls_token (`str`, *optional*, defaults to `""`):
- The classifier token which is used when doing sequence classification (classification of the whole sequence
- instead of per-token classification). It is the first token of the sequence when built with special tokens.
- mask_token (`str`, *optional*, defaults to `""`):
- The token used for masking values. This is the token used when training this model with masked language
- modeling. This is the token which the model will try to predict.
- clean_text (`bool`, *optional*, defaults to `True`):
- Whether or not to clean the text before tokenization by removing any control characters and replacing all
- whitespaces by the classic one.
- tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
- Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see [this
- issue](https://github.com/huggingface/transformers/issues/328)).
- bos_token (`str`, `optional`, defaults to `""`):
- The beginning of sentence token.
- eos_token (`str`, `optional`, defaults to `""`):
- The end of sentence token.
- strip_accents (`bool`, *optional*):
- Whether or not to strip all accents. If this option is not specified, then it will be determined by the
- value for `lowercase` (as in the original BERT).
- wordpieces_prefix (`str`, *optional*, defaults to `"##"`):
- The prefix for subwords.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = FunnelTokenizer
- cls_token_type_id: int = 2
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- do_lower_case=True,
- unk_token="",
- sep_token="",
- pad_token="",
- cls_token="",
- mask_token="",
- bos_token="",
- eos_token="",
- clean_text=True,
- tokenize_chinese_chars=True,
- strip_accents=None,
- wordpieces_prefix="##",
- **kwargs,
- ):
- super().__init__(
- vocab_file,
- tokenizer_file=tokenizer_file,
- do_lower_case=do_lower_case,
- unk_token=unk_token,
- sep_token=sep_token,
- pad_token=pad_token,
- cls_token=cls_token,
- mask_token=mask_token,
- bos_token=bos_token,
- eos_token=eos_token,
- clean_text=clean_text,
- tokenize_chinese_chars=tokenize_chinese_chars,
- strip_accents=strip_accents,
- wordpieces_prefix=wordpieces_prefix,
- **kwargs,
- )
-
- normalizer_state = json.loads(self.backend_tokenizer.normalizer.__getstate__())
- if (
- normalizer_state.get("lowercase", do_lower_case) != do_lower_case
- or normalizer_state.get("strip_accents", strip_accents) != strip_accents
- or normalizer_state.get("handle_chinese_chars", tokenize_chinese_chars) != tokenize_chinese_chars
- ):
- normalizer_class = getattr(normalizers, normalizer_state.pop("type"))
- normalizer_state["lowercase"] = do_lower_case
- normalizer_state["strip_accents"] = strip_accents
- normalizer_state["handle_chinese_chars"] = tokenize_chinese_chars
- self.backend_tokenizer.normalizer = normalizer_class(**normalizer_state)
-
- self.do_lower_case = do_lower_case
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.build_inputs_with_special_tokens with BERT->Funnel
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- """
- Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
- adding special tokens. A Funnel sequence has the following format:
-
- - single sequence: `[CLS] X [SEP]`
- - pair of sequences: `[CLS] A [SEP] B [SEP]`
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs to which the special tokens will be added.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
- """
- output = [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
-
- if token_ids_1 is not None:
- output += token_ids_1 + [self.sep_token_id]
-
- return output
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Create a mask from the two sequences passed to be used in a sequence-pair classification task. A Funnel
- Transformer sequence pair mask has the following format:
-
- ```
- 2 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- sep = [self.sep_token_id]
- cls = [self.cls_token_id]
- if token_ids_1 is None:
- return len(cls) * [self.cls_token_type_id] + len(token_ids_0 + sep) * [0]
- return len(cls) * [self.cls_token_type_id] + len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
-
- # Copied from transformers.models.bert.tokenization_bert_fast.BertTokenizerFast.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
diff --git a/transformers/models/fuyu/__init__.py b/transformers/models/fuyu/__init__.py
deleted file mode 100644
index 51a72a5366140362b65c77c7cd9dbb337ee992ea..0000000000000000000000000000000000000000
--- a/transformers/models/fuyu/__init__.py
+++ /dev/null
@@ -1,73 +0,0 @@
-# Copyright 2023 AdeptAI and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {
- "configuration_fuyu": ["FUYU_PRETRAINED_CONFIG_ARCHIVE_MAP", "FuyuConfig"],
-}
-
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["image_processing_fuyu"] = ["FuyuImageProcessor"]
- _import_structure["processing_fuyu"] = ["FuyuProcessor"]
-
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_fuyu"] = [
- "FuyuForCausalLM",
- "FuyuPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_fuyu import FUYU_PRETRAINED_CONFIG_ARCHIVE_MAP, FuyuConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .image_processing_fuyu import FuyuImageProcessor
- from .processing_fuyu import FuyuProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_fuyu import (
- FuyuForCausalLM,
- FuyuPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/fuyu/__pycache__/__init__.cpython-310.pyc b/transformers/models/fuyu/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index b25a1d831e796f410ded61d79d043d0bafe7b399..0000000000000000000000000000000000000000
Binary files a/transformers/models/fuyu/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fuyu/__pycache__/configuration_fuyu.cpython-310.pyc b/transformers/models/fuyu/__pycache__/configuration_fuyu.cpython-310.pyc
deleted file mode 100644
index a160dac6ab7b869337d14ebce8f3ecfb3b1f5fe8..0000000000000000000000000000000000000000
Binary files a/transformers/models/fuyu/__pycache__/configuration_fuyu.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fuyu/__pycache__/convert_fuyu_model_weights_to_hf.cpython-310.pyc b/transformers/models/fuyu/__pycache__/convert_fuyu_model_weights_to_hf.cpython-310.pyc
deleted file mode 100644
index 20e94ef41b7c0461857584de76dad6e25c059bfa..0000000000000000000000000000000000000000
Binary files a/transformers/models/fuyu/__pycache__/convert_fuyu_model_weights_to_hf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fuyu/__pycache__/image_processing_fuyu.cpython-310.pyc b/transformers/models/fuyu/__pycache__/image_processing_fuyu.cpython-310.pyc
deleted file mode 100644
index 06e67488f1e95fbe53e098cd04f02368e1d4253a..0000000000000000000000000000000000000000
Binary files a/transformers/models/fuyu/__pycache__/image_processing_fuyu.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fuyu/__pycache__/modeling_fuyu.cpython-310.pyc b/transformers/models/fuyu/__pycache__/modeling_fuyu.cpython-310.pyc
deleted file mode 100644
index 6e5be5fc66a742099cc2c5b04565bfa22c4d0a86..0000000000000000000000000000000000000000
Binary files a/transformers/models/fuyu/__pycache__/modeling_fuyu.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fuyu/__pycache__/processing_fuyu.cpython-310.pyc b/transformers/models/fuyu/__pycache__/processing_fuyu.cpython-310.pyc
deleted file mode 100644
index f347181320a5419ffafe28c4957b96ca8247a772..0000000000000000000000000000000000000000
Binary files a/transformers/models/fuyu/__pycache__/processing_fuyu.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/fuyu/configuration_fuyu.py b/transformers/models/fuyu/configuration_fuyu.py
deleted file mode 100644
index 40b09492d8f16130f90465a6890fb467ccbca296..0000000000000000000000000000000000000000
--- a/transformers/models/fuyu/configuration_fuyu.py
+++ /dev/null
@@ -1,211 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Adept AI and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Fuyu model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-from ..auto import CONFIG_MAPPING
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import FUYU_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class FuyuConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`FuyuForCausalLM`]. It is used to instantiate an
- Fuyu model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the
- [adept/fuyu-8b](https://huggingface.co/adept/fuyu-8b).
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 262144):
- Vocabulary size of the Fuyu model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`FuyuForCausalLM`]
- hidden_size (`int`, *optional*, defaults to 4096):
- Dimension of the hidden representations.
- intermediate_size (`int`, *optional*, defaults to 16384):
- Dimension of the MLP representations.
- num_hidden_layers (`int`, *optional*, defaults to 36):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 64):
- Number of attention heads for each attention layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"relu2"`):
- The non-linear activation function (function or string) in the decoder.
- max_position_embeddings (`int`, *optional*, defaults to 16384):
- The maximum sequence length that this model might ever be used with.
- image_size (`int`, *optional*, defaults to 300):
- The input image size.
- patch_size (`int`, *optional*, defaults to 30):
- The input vision transformer encoding patch size.
- num_channels (`int`, *optional*, defaults to 3):
- The input image number of channels.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the rms normalization layers.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`. Whether to tie weight embeddings
- tie_word_embeddings (`bool`, *optional*, defaults to `False`):
- Whether to tie input and output embeddings.
- rope_theta (`float`, *optional*, defaults to 25000.0):
- The base period of the RoPE embeddings.
- rope_scaling (`Dict`, *optional*):
- Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
- strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
- `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
- `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
- these scaling strategies behave:
- https://www.reddit.com/r/LocalFuyu/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
- experimental feature, subject to breaking API changes in future versions.
- qk_layernorm (`bool`, *optional*, defaults to `True`):
- Whether or not to normalize the Queries and Keys after projecting the hidden states
- hidden_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio after applying the MLP to the hidden states.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio after computing the attention scores.
- partial_rotary_factor (`float`, *optional*, defaults to 0.5):
- Percentage of the query and keys which will have rotary embedding.
-
- pad_token_id (`int`, *optional*):
- The id of the *padding* token.
- bos_token_id (`int`, *optional*, defaults to 1):
- The id of the *beginning-of-sequence* token.
- eos_token_id (`Union[int, List[int]]`, *optional*, defaults to 2):
- The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens.
- text_config (`dict`, *optional*):
- Dictionary of configuration options used to initialize the `language``[`Aut`].
-
- ```python
- >>> from transformers import FuyuConfig
-
- >>> # Initializing a Fuyu fuyu-7b style configuration
- >>> configuration = FuyuConfig()
- ```"""
-
- model_type = "fuyu"
- keys_to_ignore_at_inference = ["past_key_values"]
-
- def __init__(
- self,
- vocab_size=262144,
- hidden_size=4096,
- intermediate_size=16384,
- num_hidden_layers=36,
- num_attention_heads=64,
- hidden_act="relu2",
- max_position_embeddings=16384,
- image_size=300,
- patch_size=30,
- num_channels=3,
- initializer_range=0.02,
- layer_norm_eps=1e-5,
- use_cache=True,
- tie_word_embeddings=False,
- rope_theta=25000.0,
- rope_scaling=None,
- qk_layernorm=True,
- hidden_dropout=0.0,
- attention_dropout=0.0,
- partial_rotary_factor=0.5,
- pad_token_id=None,
- bos_token_id=1,
- eos_token_id=2,
- text_config=None,
- **kwargs,
- ):
- if text_config is None:
- text_config = {
- "vocab_size": vocab_size,
- "max_position_embeddings": max_position_embeddings,
- "hidden_size": hidden_size,
- "intermediate_size": intermediate_size,
- "num_hidden_layers": num_hidden_layers,
- "num_attention_heads": num_attention_heads,
- "hidden_act": hidden_act,
- "initializer_range": initializer_range,
- "layer_norm_eps": layer_norm_eps,
- "use_cache": use_cache,
- "rope_theta": rope_theta,
- "rope_scaling": rope_scaling,
- "qk_layernorm": qk_layernorm,
- "hidden_dropout": hidden_dropout,
- "attention_dropout": attention_dropout,
- "partial_rotary_factor": partial_rotary_factor,
- "pad_token_id": pad_token_id,
- "bos_token_id": bos_token_id,
- "eos_token_id": eos_token_id,
- "tie_word_embeddings": tie_word_embeddings,
- }
- logger.info("text_config is None. initializing the text model with default values.")
- text_model_type = text_config["model_type"] if "model_type" in text_config else "persimmon"
- self.text_config = CONFIG_MAPPING[text_model_type](**text_config)
-
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.image_size = image_size
- self.patch_size = patch_size
- self.num_channels = num_channels
- self.hidden_size = hidden_size
- self.intermediate_size = intermediate_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.hidden_act = hidden_act
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.rope_scaling = rope_scaling
- self.qk_layernorm = qk_layernorm
- self.hidden_dropout = hidden_dropout
- self.attention_dropout = attention_dropout
- self.partial_rotary_factor = partial_rotary_factor
- self._rope_scaling_validation()
-
- super().__init__(
- pad_token_id=pad_token_id,
- bos_token_id=bos_token_id,
- eos_token_id=eos_token_id,
- tie_word_embeddings=tie_word_embeddings,
- **kwargs,
- )
-
- # Copied from transformers.models.llama.configuration_llama.LlamaConfig._rope_scaling_validation
- def _rope_scaling_validation(self):
- """
- Validate the `rope_scaling` configuration.
- """
- if self.rope_scaling is None:
- return
-
- if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
- raise ValueError(
- "`rope_scaling` must be a dictionary with two fields, `type` and `factor`, " f"got {self.rope_scaling}"
- )
- rope_scaling_type = self.rope_scaling.get("type", None)
- rope_scaling_factor = self.rope_scaling.get("factor", None)
- if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
- raise ValueError(
- f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
- )
- if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
- raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
diff --git a/transformers/models/fuyu/convert_fuyu_model_weights_to_hf.py b/transformers/models/fuyu/convert_fuyu_model_weights_to_hf.py
deleted file mode 100644
index 6d029c0d13ab850e80f3a36f0a48cb51360a8ce1..0000000000000000000000000000000000000000
--- a/transformers/models/fuyu/convert_fuyu_model_weights_to_hf.py
+++ /dev/null
@@ -1,134 +0,0 @@
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import argparse
-import os
-import sys
-import warnings
-
-import flatdict
-import torch
-
-from transformers import FuyuConfig, FuyuForCausalLM, LlamaTokenizer
-
-
-try:
- from transformers import LlamaTokenizerFast
-
- tokenizer_class = LlamaTokenizerFast
-except ImportError as e:
- warnings.warn(e)
- warnings.warn(
- "The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion"
- )
- tokenizer_class = LlamaTokenizer
-
-"""
-Sample usage: # TODO fix clone links from persimmon to fuyu
-```
-git clone https://github.com/adept-ai-labs/adept-inference
-wget https://axtkn4xl5cip.objectstorage.us-phoenix-1.oci.customer-oci.com/n/axtkn4xl5cip/b/adept-public-data/o/8b_base_model_release.tar
-wget https://axtkn4xl5cip.objectstorage.us-phoenix-1.oci.customer-oci.com/n/axtkn4xl5cip/b/adept-public-data/o/8b_chat_model_release.tar
-python src/transformers/models/fuyu/convert_fuyu_weights_to_hf.py --input_dir /path/to/downloaded/fuyu/weights/ --output_dir /output/path
-```
-
-Thereafter, models can be loaded via:
-
-```py
-from transformers import FuyuForCausalLM, FuyuTokenizer
-
-model = FuyuForCausalLM.from_pretrained("/output/path")
-tokenizer = FuyuTokenizer.from_pretrained("/output/path")
-```
-
-Important note: you need to be able to host the whole model in RAM to execute this script (even if the biggest versions
-come in several checkpoints they each contain a part of each weight of the model, so we need to load them all in RAM).
-"""
-
-
-KEYS_TO_MODIFY_MAPPING = {
- "self_attention": "self_attn",
- "language_model.encoder": "language_model.model",
- "word_embeddings_for_head": "language_model.lm_head",
- "language_model.embedding.word_embeddings": "language_model.model.embed_tokens",
- "vit_encoder.linear_encoder": "vision_embed_tokens",
-}
-
-KEYS_TO_REMOVE = {
- "rotary_emb.inv_freq",
- "image_patch_projection",
- "image_patch_projection.weight",
- "image_patch_projection.bias",
-}
-
-
-def rename_state_dict(state_dict):
- model_state_dict = {}
- for key, value in state_dict.items():
- for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
- if key_to_modify in key:
- key = key.replace(key_to_modify, new_key)
- # if KEYS_TO_REMOVE in key:
- if key in KEYS_TO_REMOVE:
- continue
- model_state_dict[key] = value
- return model_state_dict
-
-
-def convert_fuyu_checkpoint(pytorch_dump_folder_path, ada_lib_path, pt_model_path, safe_serialization=False):
- sys.path.insert(0, ada_lib_path)
- model_state_dict_base = torch.load(pt_model_path, map_location="cpu")
- state_dict = flatdict.FlatDict(model_state_dict_base["model"], ".")
- state_dict = rename_state_dict(state_dict)
-
- transformers_config = FuyuConfig()
- model = FuyuForCausalLM(transformers_config).to(torch.bfloat16)
- model.load_state_dict(state_dict)
- model.save_pretrained(pytorch_dump_folder_path, safe_serialization=safe_serialization)
- transformers_config.save_pretrained(pytorch_dump_folder_path)
-
-
-def main():
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--input_dir",
- help="Location of Fuyu weights, which contains tokenizer.model and model folders",
- )
- parser.add_argument(
- "--pt_model_path",
- help="Location of Fuyu `model_optim_rng.pt`",
- )
- parser.add_argument(
- "--output_dir",
- help="Location to write HF model and tokenizer",
- )
- parser.add_argument(
- "--ada_lib_path",
- help="Location of original source code from adept to deserialize .pt checkpoint",
- )
- parser.add_argument("--safe_serialization", type=bool, help="Whether or not to save using `safetensors`.")
- args = parser.parse_args()
- spm_path = os.path.join(args.input_dir, "adept_vocab.model")
-
- convert_fuyu_checkpoint(
- pytorch_dump_folder_path=args.output_dir,
- pt_model_path=args.pt_model_path,
- safe_serialization=args.safe_serialization,
- ada_lib_path=args.ada_lib_path,
- )
- tokenizer = tokenizer_class(spm_path, bos_token="|ENDOFTEXT|", eos_token="|ENDOFTEXT|")
- tokenizer.save_pretrained(args.output_dir)
-
-
-if __name__ == "__main__":
- main()
diff --git a/transformers/models/fuyu/image_processing_fuyu.py b/transformers/models/fuyu/image_processing_fuyu.py
deleted file mode 100644
index ec5e1a36abb75ceb6cd9b817d3166451d559f611..0000000000000000000000000000000000000000
--- a/transformers/models/fuyu/image_processing_fuyu.py
+++ /dev/null
@@ -1,736 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for Fuyu."""
-
-import math
-from typing import Dict, List, Optional, Union
-
-import numpy as np
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature
-from ...image_transforms import (
- pad,
- resize,
- to_channel_dimension_format,
-)
-from ...image_utils import (
- ChannelDimension,
- ImageInput,
- PILImageResampling,
- get_image_size,
- infer_channel_dimension_format,
- is_scaled_image,
- is_valid_image,
- make_list_of_images,
- to_numpy_array,
- validate_preprocess_arguments,
-)
-from ...utils import (
- TensorType,
- is_torch_available,
- is_torch_device,
- is_torch_dtype,
- logging,
- requires_backends,
-)
-
-
-if is_torch_available():
- import torch
-
-
-logger = logging.get_logger(__name__)
-
-
-def make_list_of_list_of_images(
- images: Union[List[List[ImageInput]], List[ImageInput], ImageInput],
-) -> List[List[ImageInput]]:
- if is_valid_image(images):
- return [[images]]
-
- if isinstance(images, list) and all(isinstance(image, list) for image in images):
- return images
-
- if isinstance(images, list):
- return [make_list_of_images(image) for image in images]
-
- raise ValueError("images must be a list of list of images or a list of images or an image.")
-
-
-class FuyuBatchFeature(BatchFeature):
- """
- BatchFeature class for Fuyu image processor and processor.
-
- The outputs dictionary from the processors contains a mix of tensors and lists of tensors.
- """
-
- def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None):
- """
- Convert the inner content to tensors.
-
- Args:
- tensor_type (`str` or [`~utils.TensorType`], *optional*):
- The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If
- `None`, no modification is done.
- """
- if tensor_type is None:
- return self
-
- is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type=tensor_type)
-
- def _convert_tensor(elem):
- if is_tensor(elem):
- return elem
- return as_tensor(elem)
-
- def _safe_convert_tensor(elem):
- try:
- return _convert_tensor(elem)
- except: # noqa E722
- if key == "overflowing_values":
- raise ValueError("Unable to create tensor returning overflowing values of different lengths. ")
- raise ValueError(
- "Unable to create tensor, you should probably activate padding "
- "with 'padding=True' to have batched tensors with the same length."
- )
-
- # Do the tensor conversion in batch
- for key, value in self.items():
- if isinstance(value, list) and isinstance(value[0], list):
- # List[List[Any]] -> List[List[Tensor]]
- self[key] = [[_safe_convert_tensor(elem) for elem in elems] for elems in value]
- elif isinstance(value, list):
- # List[Any] -> List[Tensor]
- self[key] = [_safe_convert_tensor(elem) for elem in value]
- else:
- # Any -> Tensor
- self[key] = _safe_convert_tensor(value)
- return self
-
- def to(self, *args, **kwargs) -> "BatchFeature":
- """
- Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in
- different `dtypes` and sending the `BatchFeature` to a different `device`.
-
- Args:
- args (`Tuple`):
- Will be passed to the `to(...)` function of the tensors.
- kwargs (`Dict`, *optional*):
- Will be passed to the `to(...)` function of the tensors.
-
- Returns:
- [`BatchFeature`]: The same instance after modification.
- """
- requires_backends(self, ["torch"])
- import torch # noqa
-
- new_data = {}
- device = kwargs.get("device")
- # Check if the args are a device or a dtype
- if device is None and len(args) > 0:
- # device should be always the first argument
- arg = args[0]
- if is_torch_dtype(arg):
- # The first argument is a dtype
- pass
- elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
- device = arg
- else:
- # it's something else
- raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
-
- def _to(elem):
- # check if v is a floating point
- if torch.is_floating_point(elem):
- # cast and send to device
- return elem.to(*args, **kwargs)
- if device is not None:
- return elem.to(device=device)
-
- return elem
-
- # We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
- for k, v in self.items():
- if isinstance(v, list) and isinstance(v[0], list):
- # Data structure is a list of lists
- new_v = []
- for elems in v:
- new_v.append([_to(elem) for elem in elems])
- new_data[k] = new_v
- elif isinstance(v, list):
- # Data structure is a list
- new_data[k] = [_to(elem) for elem in v]
- else:
- new_data[k] = _to(v)
- self.data = new_data
- return self
-
-
-class FuyuImageProcessor(BaseImageProcessor):
- """
- This class should handle the image processing part before the main FuyuForCausalLM. In particular, it should
- handle:
-
- - Processing Images:
- Taking a batch of images as input. If the images are variable-sized, it resizes them based on the desired patch
- dimensions. The image output is always img_h, img_w of (1080, 1920)
-
- Then, it patches up these images using the patchify_image function.
-
- - Creating Image Input IDs:
- For each patch, a placeholder ID is given to identify where these patches belong in a token sequence. For
- variable-sized images, each line of patches is terminated with a newline ID.
-
- - Image Patch Indices:
- For each image patch, the code maintains an index where these patches should be inserted in a token stream.
-
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the image to `size`.
- size (`Dict[str, int]`, *optional*, defaults to `{"height": 1080, "width": 1920}`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
- resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
- `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BILINEAR`.
- do_pad (`bool`, *optional*, defaults to `True`):
- Whether to pad the image to `size`.
- padding_value (`float`, *optional*, defaults to 1.0):
- The value to pad the image with.
- padding_mode (`str`, *optional*, defaults to `"constant"`):
- The padding mode to use when padding the image.
- do_normalize (`bool`, *optional*, defaults to `True`):
- Whether to normalize the image.
- image_mean (`float`, *optional*, defaults to 0.5):
- The mean to use when normalizing the image.
- image_std (`float`, *optional*, defaults to 0.5):
- The standard deviation to use when normalizing the image.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether to rescale the image.
- rescale_factor (`float`, *optional*, defaults to `1 / 255`):
- The factor to use when rescaling the image.
- patch_size (`Dict[str, int]`, *optional*, defaults to `{"height": 30, "width": 30}`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the patches.
- """
-
- model_input_names = [
- "images",
- "image_input_ids",
- "image_patches",
- "image_patch_indices_per_batch",
- "image_patch_indices_per_subsequence",
- ]
-
- def __init__(
- self,
- do_resize: bool = True,
- size: Optional[Dict[str, int]] = None,
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- do_pad: bool = True,
- padding_value: float = 1.0,
- padding_mode: str = "constant",
- do_normalize: bool = True,
- image_mean: Union[float, List[float]] = 0.5,
- image_std: Union[float, List[float]] = 0.5,
- do_rescale: bool = True,
- rescale_factor: float = 1 / 255,
- patch_size: Optional[Dict[str, int]] = None,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.do_resize = do_resize
- self.size = size if size is not None else {"height": 1080, "width": 1920}
- self.resample = resample
- self.do_pad = do_pad
- self.padding_value = padding_value
- self.padding_mode = padding_mode
- self.do_normalize = do_normalize
- self.image_mean = image_mean
- self.image_std = image_std
- self.do_rescale = do_rescale
- self.rescale_factor = rescale_factor
- self.patch_size = patch_size if patch_size is not None else {"height": 30, "width": 30}
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size",
- "resample",
- "do_pad",
- "padding_value",
- "padding_mode",
- "do_normalize",
- "image_mean",
- "image_std",
- "do_rescale",
- "rescale_factor",
- "patch_size",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- def resize(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize an image to `(size["height"], size["width"])`.
-
- Args:
- image (`np.ndarray`):
- Image to resize.
- size (`Dict[str, int]`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
- resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
- `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BILINEAR`.
- data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the output image. If unset, the channel dimension format of the input
- image is used. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
-
- Returns:
- `np.ndarray`: The resized image.
- """
- image_height, image_width = get_image_size(image, input_data_format)
- target_height, target_width = size["height"], size["width"]
-
- if image_width <= target_width and image_height <= target_height:
- return image
-
- height_scale_factor = target_height / image_height
- width_scale_factor = target_width / image_width
- optimal_scale_factor = min(height_scale_factor, width_scale_factor)
-
- new_height = int(image_height * optimal_scale_factor)
- new_width = int(image_width * optimal_scale_factor)
-
- scaled_image = resize(
- image=image,
- size=(new_height, new_width),
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
- return scaled_image
-
- def pad_image(
- self,
- image: np.ndarray,
- size: Dict[str, int],
- mode: str = "constant",
- constant_values: float = 1.0,
- data_format: Optional[Union[str, ChannelDimension]] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- ) -> np.ndarray:
- """
- Pad an image to `(size["height"], size["width"])`.
-
- Args:
- image (`np.ndarray`):
- Image to pad.
- size (`Dict[str, int]`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
- data_format (`ChannelDimension` or `str`, *optional*):
- The data format of the output image. If unset, the same format as the input image is used.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not provided, it will be inferred.
- """
- image_height, image_width = get_image_size(image, input_data_format)
- target_height, target_width = size["height"], size["width"]
- padding_top = 0
- padding_left = 0
- padding_bottom = target_height - image_height
- padding_right = target_width - image_width
- padded_image = pad(
- image,
- padding=((padding_top, padding_bottom), (padding_left, padding_right)),
- mode=mode,
- constant_values=constant_values,
- data_format=data_format,
- input_data_format=input_data_format,
- )
- return padded_image
-
- def preprocess(
- self,
- images,
- do_resize: Optional[bool] = None,
- size: Optional[Dict[str, int]] = None,
- resample: Optional[PILImageResampling] = None,
- do_pad: Optional[bool] = None,
- padding_value: Optional[float] = None,
- padding_mode: Optional[str] = None,
- do_normalize: Optional[bool] = None,
- image_mean: Optional[float] = None,
- image_std: Optional[float] = None,
- do_rescale: Optional[bool] = None,
- rescale_factor: Optional[float] = None,
- patch_size: Optional[Dict[str, int]] = None,
- data_format: Optional[Union[str, ChannelDimension]] = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- return_tensors: Optional[TensorType] = None,
- ):
- """
-
- Utility function to preprocess the images and extract necessary information about original formats.
-
- Args:
- images (`ImageInput`):
- Images to preprocess. Expects a single image, a list or images or a list of lists of images. Pixel
- values range from 0 to 255, or between 0 and 1 if `do_rescale` is `False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the image to `size`.
- size (`Dict[str, int]`, *optional*, defaults to `self.size`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
- resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
- `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BILINEAR`.
- do_pad (`bool`, *optional*, defaults to `self.do_pad`):
- Whether to pad the image to `size`.
- padding_value (`float`, *optional*, defaults to `self.padding_value`):
- The value to pad the image with.
- padding_mode (`str`, *optional*, defaults to `self.padding_mode`):
- The padding mode to use when padding the image.
- do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
- Whether to normalize the image.
- image_mean (`float`, *optional*, defaults to `self.image_mean`):
- The mean to use when normalizing the image.
- image_std (`float`, *optional*, defaults to `self.image_std`):
- The standard deviation to use when normalizing the image.
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether to rescale the image.
- rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
- The factor to use when rescaling the image.
- patch_size (`Dict[str, int]`, *optional*, defaults to `self.patch_size`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the patches.
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - Unset: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format of the output image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- """
-
- do_resize = do_resize if do_resize is not None else self.do_resize
- size = size if size is not None else self.size
- resample = resample if resample is not None else self.resample
- do_pad = do_pad if do_pad is not None else self.do_pad
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- do_normalize = do_normalize if do_normalize is not None else self.do_normalize
- image_mean = image_mean if image_mean is not None else self.image_mean
- image_std = image_std if image_std is not None else self.image_std
- padding_value = padding_value if padding_value is not None else self.padding_value
- padding_mode = padding_mode if padding_mode is not None else self.padding_mode
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
- patch_size = patch_size if patch_size is not None else self.patch_size
-
- if isinstance(images, list) and any(isinstance(elem, list) and len(elem) >= 2 for elem in images):
- raise ValueError("Multiple images for a single sample are not yet supported.")
-
- batch_images = make_list_of_list_of_images(images)
-
- validate_preprocess_arguments(
- do_rescale=do_rescale,
- rescale_factor=rescale_factor,
- do_normalize=do_normalize,
- image_mean=image_mean,
- image_std=image_std,
- do_pad=do_pad,
- size_divisibility=size, # There is no pad divisibility in this processor, but pad requires the size arg.
- do_resize=do_resize,
- size=size,
- resample=resample,
- )
- # All transformations expect numpy arrays.
- batch_images = [[to_numpy_array(image) for image in images] for images in batch_images]
-
- if is_scaled_image(batch_images[0][0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(batch_images[0][0])
-
- original_image_sizes = [get_image_size(images[0], channel_dim=input_data_format) for images in batch_images]
-
- if do_resize:
- batch_images = [
- [self.resize(image, size=size, input_data_format=input_data_format) for image in images]
- for images in batch_images
- ]
-
- image_sizes = [get_image_size(images[0], channel_dim=input_data_format) for images in batch_images]
- image_unpadded_heights = [[image_size[0]] for image_size in image_sizes]
- image_unpadded_widths = [[image_size[1]] for image_size in image_sizes]
-
- # scale_h is the same as scale_w
- image_scale_factors = [
- [resized_size[0] / original_size[0]]
- for original_size, resized_size in zip(original_image_sizes, image_sizes)
- ]
-
- if do_pad:
- batch_images = [
- [
- self.pad_image(
- image,
- size=size,
- mode=padding_mode,
- constant_values=padding_value,
- input_data_format=input_data_format,
- )
- for image in images
- ]
- for images in batch_images
- ]
-
- if do_rescale:
- batch_images = [
- [self.rescale(image, scale=rescale_factor, input_data_format=input_data_format) for image in images]
- for images in batch_images
- ]
-
- if do_normalize:
- batch_images = [
- [
- self.normalize(image, mean=image_mean, std=image_std, input_data_format=input_data_format)
- for image in images
- ]
- for images in batch_images
- ]
-
- if data_format is not None:
- batch_images = [
- [to_channel_dimension_format(image, data_format, input_data_format) for image in images]
- for images in batch_images
- ]
-
- data = {
- "images": batch_images,
- "image_unpadded_heights": image_unpadded_heights,
- "image_unpadded_widths": image_unpadded_widths,
- "image_scale_factors": image_scale_factors,
- }
- return FuyuBatchFeature(data=data, tensor_type=return_tensors)
-
- def get_num_patches(self, image_height: int, image_width: int, patch_size: Dict[str, int] = None) -> int:
- """
- Calculate number of patches required to encode an image.
-
- Args:
- image_height (`int`):
- Height of the image.
- image_width (`int`):
- Width of the image.
- patch_size (`Dict[str, int]`, *optional*, defaults to `self.patch_size`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the patches.
- """
- patch_size = patch_size if patch_size is not None else self.patch_size
- patch_height, patch_width = self.patch_size["height"], self.patch_size["width"]
-
- if image_height % patch_height != 0:
- raise ValueError(f"{image_height=} must be divisible by {patch_height}")
- if image_width % patch_width != 0:
- raise ValueError(f"{image_width=} must be divisible by {patch_width}")
-
- num_patches_per_dim_h = image_height // patch_height
- num_patches_per_dim_w = image_width // patch_width
- num_patches = num_patches_per_dim_h * num_patches_per_dim_w
- return num_patches
-
- def patchify_image(self, image: "torch.Tensor", patch_size: Optional[Dict[str, int]] = None) -> "torch.Tensor":
- """
- Convert an image into a tensor of patches.
-
- Args:
- image (`torch.Tensor`):
- Image to convert. Shape: [batch, channels, height, width]
- patch_size (`Dict[str, int]`, *optional*, defaults to `self.patch_size`):
- Dictionary in the format `{"height": int, "width": int}` specifying the size of the patches.
- """
- requires_backends(self, ["torch"])
- patch_size = patch_size if patch_size is not None else self.patch_size
- patch_height, patch_width = patch_size["height"], patch_size["width"]
-
- # TODO refer to https://github.com/ArthurZucker/transformers/blob/0f0a3fe5ca5697ee58faeb5b53f049af720b5e98/src/transformers/models/vit_mae/modeling_vit_mae.py#L871
- # torch implementation is faster but does not handle non-squares
-
- batch_size, channels, _, _ = image.shape
- unfolded_along_height = image.unfold(2, patch_height, patch_height)
- patches = unfolded_along_height.unfold(3, patch_width, patch_width)
- patches = patches.contiguous()
- patches = patches.view(batch_size, channels, -1, patch_height, patch_width)
- patches = patches.permute(0, 2, 3, 4, 1)
- patches = patches.reshape(batch_size, -1, channels * patch_height * patch_width)
- return patches
-
- def preprocess_with_tokenizer_info(
- self,
- image_input: "torch.Tensor",
- image_present: "torch.Tensor",
- image_unpadded_h: "torch.Tensor",
- image_unpadded_w: "torch.Tensor",
- image_placeholder_id: int,
- image_newline_id: int,
- variable_sized: bool,
- patch_size: Optional[Dict[str, int]] = None,
- ) -> FuyuBatchFeature:
- """Process images for model input. In particular, variable-sized images are handled here.
-
- Args:
- image_input (`torch.Tensor` of shape [batch_size, subsequence_size, num_channels, height, width]):
- Tensor of images padded to model input size.
- image_present (`torch.Tensor` of shape [batch_size, subsequence_size, num_images]):
- Tensor of 1s and 0s indicating whether an image is present.
- image_unpadded_h (`torch.Tensor` of shape [batch_size, subsequence_size]):
- Tensor of unpadded image heights.
- image_unpadded_w (`torch.Tensor` of shape [batch_size, subsequence_size]):
- Tensor of unpadded image widths.
- image_placeholder_id (int):
- The id of the image placeholder token. Comes from an associated tokenizer.
- image_newline_id (int):
- The id of the image newline token. Comes from an associated tokenizer.
- variable_sized (bool):
- Whether to process images as variable-sized.
- patch_size (`Dict[str, int]`, *optional*, defaults to `self.patch_size`):
- Size of the patches.
- """
- requires_backends(self, ["torch"])
-
- patch_size = patch_size if patch_size is not None else self.patch_size
- patch_height, patch_width = patch_size["height"], patch_size["width"]
-
- # Only images that are present.
- images: List[List[torch.Tensor]] = []
- batch_image_patches: List[List[torch.Tensor]] = []
- # Image input ids for every subsequence, including ones with no image present.
- batch_image_input_ids: List[List[torch.Tensor]] = []
- for batch_index in range(image_input.shape[0]):
- image_input_ids = []
- image_patches = []
- for subseq_index in range(image_input.shape[1]):
- if image_present[batch_index, subseq_index]:
- image = image_input[batch_index, subseq_index]
- image_height, image_width = image.shape[1], image.shape[2]
- if variable_sized:
- # The min() is required here due to floating point issues:
- # math.ceil(torch.tensor(300).cuda() / 30) == 11
- new_h = min(
- image_height,
- math.ceil(image_unpadded_h[batch_index, subseq_index] / patch_height) * patch_height,
- )
- new_w = min(
- image_width,
- math.ceil(image_unpadded_w[batch_index, subseq_index] / patch_width) * patch_width,
- )
- image = image[:, :new_h, :new_w]
- image_height, image_width = new_h, new_w
-
- num_patches = self.get_num_patches(image_height=image_height, image_width=image_width)
- tensor_of_image_ids = torch.full(
- [num_patches], image_placeholder_id, dtype=torch.int32, device=image_input.device
- )
- patches = self.patchify_image(image=image.unsqueeze(0)).squeeze(0)
- assert num_patches == patches.shape[0]
-
- if variable_sized:
- # Now terminate each line with |NEWLINE|.
- tensor_of_image_ids = tensor_of_image_ids.reshape(-1, image_width // patch_width)
- newline_ids = torch.full(
- [tensor_of_image_ids.shape[0], 1],
- image_newline_id,
- dtype=torch.int32,
- device=image_input.device,
- )
- tensor_of_image_ids = torch.cat([tensor_of_image_ids, newline_ids], dim=1)
- tensor_of_image_ids = tensor_of_image_ids.reshape(-1)
-
- images.append([image])
- image_input_ids.append(tensor_of_image_ids)
- image_patches.append(patches)
- else:
- image_input_ids.append(torch.tensor([], dtype=torch.int32, device=image_input.device))
-
- batch_image_input_ids.append(image_input_ids)
- batch_image_patches.append(image_patches)
-
- # Create image_patch_input_indices, where non-negative values correspond to image patches to be inserted in
- # the stream.
- image_patch_indices_per_batch: List[List[torch.Tensor]] = []
- image_patch_indices_per_subsequence: List[List[torch.Tensor]] = []
-
- for sample_image_input_ids in batch_image_input_ids:
- index_offset = 0
- per_batch_indices = []
- per_subsequence_indices = []
- for subseq_image_input_ids in sample_image_input_ids:
- # Indices of image patches.
- patches_mask = subseq_image_input_ids == image_placeholder_id
- num_patches = torch.count_nonzero(patches_mask)
- indices = torch.arange(num_patches, dtype=torch.int64, device=subseq_image_input_ids.device).type_as(
- subseq_image_input_ids
- )
-
- # Place those indices in the image input ids token stream, with -1 representing non-index tokens.
- indices_in_stream_per_batch = torch.full_like(subseq_image_input_ids, -1)
- indices_in_stream_per_subsequence = torch.full_like(subseq_image_input_ids, -1)
- patches_inds = torch.nonzero(patches_mask, as_tuple=True)[0]
-
- indices_in_stream_per_batch[patches_inds] = indices + index_offset
- indices_in_stream_per_subsequence[patches_inds] = indices
-
- per_batch_indices.append(indices_in_stream_per_batch)
- per_subsequence_indices.append(indices_in_stream_per_subsequence)
- index_offset += num_patches
-
- image_patch_indices_per_batch.append(per_batch_indices)
- image_patch_indices_per_subsequence.append(per_subsequence_indices)
-
- return FuyuBatchFeature(
- data={
- "images": images,
- "image_input_ids": batch_image_input_ids,
- "image_patches": batch_image_patches,
- "image_patch_indices_per_batch": image_patch_indices_per_batch,
- "image_patch_indices_per_subsequence": image_patch_indices_per_subsequence,
- }
- )
diff --git a/transformers/models/fuyu/modeling_fuyu.py b/transformers/models/fuyu/modeling_fuyu.py
deleted file mode 100644
index 8e9a41954aee9c96945a787f08d698091baa0d4b..0000000000000000000000000000000000000000
--- a/transformers/models/fuyu/modeling_fuyu.py
+++ /dev/null
@@ -1,358 +0,0 @@
-# coding=utf-8
-# Copyright 2023 HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Fuyu model."""
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from ...modeling_outputs import CausalLMOutputWithPast
-from ...modeling_utils import PreTrainedModel
-from ...models.auto.modeling_auto import AutoModelForCausalLM
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from .configuration_fuyu import FuyuConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "FuyuConfig"
-
-
-FUYU_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`FuyuConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-@add_start_docstrings(
- "The bare Fuyu Model outputting raw hidden-states without any specific head on top.",
- FUYU_START_DOCSTRING,
-)
-class FuyuPreTrainedModel(PreTrainedModel):
- config_class = FuyuConfig
- base_model_prefix = "fuyu"
- supports_gradient_checkpointing = True
- _no_split_modules = []
- _skip_keys_device_placement = "past_key_values"
-
- def _init_weights(self, module):
- std = self.config.initializer_range
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
-
-FUYU_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
- it.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
- `past_key_values`).
-
- If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
- and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
- information on the default strategy.
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
- image_patches (`torch.FloatTensor` of shape `(batch_size, num_total_patches, patch_size_ x patch_size x num_channels)`, *optional*):
- Image patches to be used as continuous embeddings. The patches are flattened and then projected to the
- hidden size of the model.
- image_patches_indices (`torch.LongTensor` of shape `(batch_size, num_total_patches + number_of_newline_tokens + number_of_text_tokens, patch_size_ x patch_size x num_channels )`, *optional*):
- Indices indicating at which position the image_patches have to be inserted in input_embeds.
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
- `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
- `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
-
- Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
- blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "Fuyu Model with a language modeling head on top for causal language model conditioned on image patches and text.",
- FUYU_START_DOCSTRING,
-)
-class FuyuForCausalLM(FuyuPreTrainedModel):
- def __init__(self, config: FuyuConfig):
- super().__init__(config)
- self.padding_idx = config.pad_token_id
- self.vocab_size = config.vocab_size
- self.language_model = AutoModelForCausalLM.from_config(config.text_config)
-
- self.vision_embed_tokens = nn.Linear(
- config.patch_size * config.patch_size * config.num_channels, config.hidden_size
- )
-
- self.gradient_checkpointing = False
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.language_model.get_input_embeddings()
-
- def set_input_embeddings(self, value):
- self.language_model.set_input_embeddings(value)
-
- def gather_continuous_embeddings(
- self,
- word_embeddings: torch.Tensor,
- continuous_embeddings: List[torch.Tensor],
- image_patch_input_indices: torch.Tensor,
- ) -> torch.Tensor:
- """This function places the continuous_embeddings into the word_embeddings at the locations
- indicated by image_patch_input_indices. Different batch elements can have different numbers of continuous
- embeddings.
-
- Args:
- word_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Tensor of word embeddings.
- continuous_embeddings (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`):
- Tensor of continuous embeddings. The length of the list is the batch size. Each entry is shape
- [num_image_embeddings, hidden], and num_image_embeddings needs to match the number of non-negative
- indices in image_patch_input_indices for that batch element.
- image_patch_input_indices (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Tensor of indices of the image patches in the input_ids tensor.
- """
- if not (word_embeddings.shape[0] == len(continuous_embeddings)):
- raise ValueError(
- f"Batch sizes must match! Got {len(continuous_embeddings)=} and {word_embeddings.shape[0]=}"
- )
-
- output_embeddings = word_embeddings.clone()
- for batch_idx in range(word_embeddings.shape[0]):
- # First, find the positions of all the non-negative values in image_patch_input_indices, those are the
- # positions in word_embeddings that we want to replace with content from continuous_embeddings.
- dst_indices = torch.nonzero(image_patch_input_indices[batch_idx] >= 0, as_tuple=True)[0]
- # Next look up those indices in image_patch_input_indices to find the indices in continuous_embeddings that we
- # want to use to replace the values in word_embeddings.
- src_indices = image_patch_input_indices[batch_idx][dst_indices]
- # Check if we have more indices than embeddings. Note that we could have fewer indices if images got truncated.
- if src_indices.shape[0] > continuous_embeddings[batch_idx].shape[0]:
- raise ValueError(
- f"Number of continuous embeddings {continuous_embeddings[batch_idx].shape=} does not match "
- f"number of continuous token ids {src_indices.shape=} in batch element {batch_idx}."
- )
- output_embeddings[batch_idx, dst_indices] = continuous_embeddings[batch_idx][src_indices]
- return output_embeddings
-
- @add_start_docstrings_to_model_forward(FUYU_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- image_patches: torch.Tensor = None, # [batch_size, num_total_patches, patch_size_ x patch_size x num_channels ]
- image_patches_indices: torch.Tensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- labels: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import FuyuProcessor, FuyuForCausalLM
- >>> from PIL import Image
- >>> import requests
-
- >>> processor = FuyuProcessor.from_pretrained("adept/fuyu-8b")
- >>> model = FuyuForCausalLM.from_pretrained("adept/fuyu-8b")
-
- >>> url = "https://huggingface.co/datasets/hf-internal-testing/fixtures-captioning/resolve/main/bus.png"
- >>> image = Image.open(requests.get(url, stream=True).raw)
- >>> prompt = "Generate a coco-style caption.\n"
-
- >>> inputs = processor(text=prompt, images=image, return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> generated_ids = model.generate(**inputs, max_new_tokens=7)
- >>> generation_text = processor.batch_decode(generated_ids[:, -7:], skip_special_tokens=True)
- >>> print(generation_text[0])
- A blue bus parked on the side of a road.
- ```"""
-
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
-
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- batch_size, seq_length = input_ids.shape
- elif inputs_embeds is not None:
- batch_size, seq_length, _ = inputs_embeds.shape
- else:
- raise ValueError("You have to specify either input_is or inputs_embeds")
-
- seq_length_with_past = seq_length
- past_key_values_length = 0
-
- if past_key_values is not None:
- past_key_values_length = past_key_values[0][0].shape[2]
- seq_length_with_past = seq_length_with_past + past_key_values_length
-
- if position_ids is None:
- device = input_ids.device if input_ids is not None else inputs_embeds.device
- position_ids = torch.arange(
- past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
- )
- position_ids = position_ids.unsqueeze(0)
-
- if inputs_embeds is None:
- inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
- if image_patches is not None and past_key_values is None:
- patch_embeddings = [
- self.vision_embed_tokens(patch.to(self.vision_embed_tokens.weight.dtype))
- .squeeze(0)
- .to(inputs_embeds.device)
- for patch in image_patches
- ]
- inputs_embeds = self.gather_continuous_embeddings(
- word_embeddings=inputs_embeds,
- continuous_embeddings=patch_embeddings,
- image_patch_input_indices=image_patches_indices,
- )
-
- outputs = self.language_model(
- inputs_embeds=inputs_embeds,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- labels=labels,
- use_cache=use_cache,
- return_dict=return_dict,
- )
-
- return outputs
-
- def prepare_inputs_for_generation(
- self,
- input_ids,
- past_key_values=None,
- attention_mask=None,
- inputs_embeds=None,
- image_patches=None,
- image_patches_indices=None,
- **kwargs,
- ):
- if past_key_values:
- input_ids = input_ids[:, -1:]
-
- position_ids = kwargs.get("position_ids", None)
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -1].unsqueeze(-1)
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- model_inputs = {"input_ids": input_ids}
-
- if image_patches_indices is not None:
- model_inputs["image_patches_indices"] = image_patches_indices
-
- model_inputs.update(
- {
- "position_ids": position_ids,
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "attention_mask": attention_mask,
- "image_patches_indices": image_patches_indices if past_key_values is None else None,
- "image_patches": image_patches if past_key_values is None else None,
- }
- )
- return model_inputs
diff --git a/transformers/models/fuyu/processing_fuyu.py b/transformers/models/fuyu/processing_fuyu.py
deleted file mode 100644
index ffa215f1a0652ec5df567ddf28b8975cbb2e467c..0000000000000000000000000000000000000000
--- a/transformers/models/fuyu/processing_fuyu.py
+++ /dev/null
@@ -1,694 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Image/Text processor class for GIT
-"""
-import re
-from typing import Dict, List, Optional, Tuple, Union
-
-import numpy as np
-
-from ...processing_utils import ProcessorMixin
-from ...tokenization_utils_base import PaddingStrategy, TruncationStrategy
-from ...utils import TensorType, is_torch_available, logging, requires_backends
-
-
-if is_torch_available():
- from .image_processing_fuyu import FuyuBatchFeature
-
-
-logger = logging.get_logger(__name__)
-
-
-if is_torch_available():
- import torch
-
-
-TEXT_REPR_BBOX_OPEN = ""
-TEXT_REPR_BBOX_CLOSE = ""
-TEXT_REPR_POINT_OPEN = ""
-TEXT_REPR_POINT_CLOSE = ""
-
-TOKEN_BBOX_OPEN_STRING = "<0x00>" #
-TOKEN_BBOX_CLOSE_STRING = "<0x01>" #
-TOKEN_POINT_OPEN_STRING = "<0x02>" #
-TOKEN_POINT_CLOSE_STRING = "<0x03>" #
-BEGINNING_OF_ANSWER_STRING = "<0x04>" #
-
-
-def full_unpacked_stream_to_tensor(
- all_bi_tokens_to_place: List[int],
- full_unpacked_stream: List["torch.Tensor"],
- fill_value: int,
- batch_size: int,
- new_seq_len: int,
- offset: int,
-) -> "torch.Tensor":
- """Takes an unpacked stream of tokens (i.e. a list of tensors, one for each item in the batch) and does
- the required padding to create a single tensor for the batch of shape batch_size x new_seq_len.
- """
-
- assert len(all_bi_tokens_to_place) == batch_size
- assert len(full_unpacked_stream) == batch_size
-
- # Create padded tensors for the full batch.
- new_padded_tensor = torch.full(
- [batch_size, new_seq_len],
- fill_value=fill_value,
- dtype=full_unpacked_stream[0].dtype,
- device=full_unpacked_stream[0].device,
- )
-
- # Place each batch entry into the batch tensor.
- for bi in range(batch_size):
- tokens_to_place = all_bi_tokens_to_place[bi]
- new_padded_tensor[bi, :tokens_to_place] = full_unpacked_stream[bi][offset : tokens_to_place + offset]
-
- return new_padded_tensor
-
-
-def construct_full_unpacked_stream(
- num_real_text_tokens: Union[List[List[int]], "torch.Tensor"],
- input_stream: "torch.Tensor",
- image_tokens: List[List["torch.Tensor"]],
- batch_size: int,
- num_sub_sequences: int,
-) -> List["torch.Tensor"]:
- """Takes an input_stream tensor of shape B x S x ?. For each subsequence, adds any required
- padding to account for images and then unpacks the subsequences to create a single sequence per item in the batch.
- Returns a list of tensors, one for each item in the batch."""
-
- all_bi_stream = []
-
- for batch_index in range(batch_size):
- all_si_stream = []
-
- # First, construct full token stream (including image placeholder tokens) and loss mask for each subsequence
- # and append to lists. We use lists rather than tensors because each subsequence is variable-sized.
- # TODO Remove this logic in a subsequent release since subsequences are not supported.
- image_adjustment = image_tokens[batch_index][0]
- subsequence_stream = torch.cat([image_adjustment, input_stream[batch_index, 0]], dim=0)
- num_real_tokens = image_adjustment.shape[0] + num_real_text_tokens[batch_index][0]
- all_si_stream.append(subsequence_stream[:num_real_tokens])
- all_bi_stream.append(torch.cat(all_si_stream, dim=0))
-
- return all_bi_stream
-
-
-def _replace_string_repr_with_token_tags(prompt: str) -> str:
- prompt = prompt.replace(TEXT_REPR_POINT_OPEN, TOKEN_POINT_OPEN_STRING)
- prompt = prompt.replace(TEXT_REPR_POINT_CLOSE, TOKEN_POINT_CLOSE_STRING)
- prompt = prompt.replace(TEXT_REPR_BBOX_OPEN, TOKEN_BBOX_OPEN_STRING)
- prompt = prompt.replace(TEXT_REPR_BBOX_CLOSE, TOKEN_BBOX_CLOSE_STRING)
- return prompt
-
-
-def _segment_prompt_into_text_token_conversions(prompt: str) -> List:
- """
- Given a string prompt, converts the prompt into a list of TextTokenConversions.
- """
- # Wherever, we notice the [TOKEN_OPEN_STRING, TOKEN_CLOSE_STRING], we split the prompt
- prompt_text_list: List = []
- regex_pattern = re.compile(
- f"({TOKEN_BBOX_OPEN_STRING}|{TOKEN_BBOX_CLOSE_STRING}|{TOKEN_POINT_OPEN_STRING}|{TOKEN_POINT_CLOSE_STRING})"
- )
- # Split by the regex pattern
- prompt_split = regex_pattern.split(prompt)
- for i, elem in enumerate(prompt_split):
- if len(elem) == 0 or elem in [
- TOKEN_BBOX_OPEN_STRING,
- TOKEN_BBOX_CLOSE_STRING,
- TOKEN_POINT_OPEN_STRING,
- TOKEN_POINT_CLOSE_STRING,
- ]:
- continue
- prompt_text_list.append(
- (elem, i > 1 and prompt_split[i - 1] in [TOKEN_BBOX_OPEN_STRING, TOKEN_POINT_OPEN_STRING])
- )
- return prompt_text_list
-
-
-def _transform_coordinates_and_tokenize(prompt: str, scale_factor: float, tokenizer) -> List[int]:
- """
- This function transforms the prompt in the following fashion:
- - and to their respective token mappings
- - extract the coordinates from the tag
- - transform the coordinates into the transformed image space
- - return the prompt tokens with the transformed coordinates and new tags
-
- Bounding boxes and points MUST be in the following format: y1, x1, y2, x2 x, y The spaces
- and punctuation added above are NOT optional.
- """
- # Make a namedtuple that stores "text" and "is_bbox"
-
- # We want to do the following: Tokenize the code normally -> when we see a point or box, tokenize using the tokenize_within_tag function
- # When point or box close tag, continue tokenizing normally
- # First, we replace the point and box tags with their respective tokens
- prompt = _replace_string_repr_with_token_tags(prompt)
- # Tokenize the prompt
- # Convert prompt into a list split
- prompt_text_list = _segment_prompt_into_text_token_conversions(prompt)
- transformed_prompt_tokens: List[int] = []
- for elem in prompt_text_list:
- if elem[1]:
- # This is a location, we need to tokenize it
- within_tag_tokenized = _transform_within_tags(elem[0], scale_factor, tokenizer)
- # Surround the text with the open and close tags
- transformed_prompt_tokens.extend(within_tag_tokenized)
- else:
- transformed_prompt_tokens.extend(tokenizer(elem[0], add_special_tokens=False).input_ids)
- return transformed_prompt_tokens
-
-
-def _transform_within_tags(text: str, scale_factor: float, tokenizer) -> List[int]:
- """
- Given a bounding box of the fashion 1, 2, 3, 4 | 1, 2 This function is responsible for
- converting 1, 2, 3, 4 into tokens of 1 2 3 4 without any commas.
- """
- # Convert the text into a list of strings.
- num_int_strs = text.split(",")
- if len(num_int_strs) == 2:
- # If there are any open or close tags, remove them.
- token_space_open_string = tokenizer.vocab[TOKEN_POINT_OPEN_STRING]
- token_space_close_string = tokenizer.vocab[TOKEN_POINT_CLOSE_STRING]
- else:
- token_space_open_string = tokenizer.vocab[TOKEN_BBOX_OPEN_STRING]
- token_space_close_string = tokenizer.vocab[TOKEN_BBOX_CLOSE_STRING]
-
- # Remove all spaces from num_ints
- num_ints = [float(num.strip()) for num in num_int_strs]
- # scale to transformed image siz
- if len(num_ints) == 2:
- num_ints_translated = scale_point_to_transformed_image(x=num_ints[0], y=num_ints[1], scale_factor=scale_factor)
- elif len(num_ints) == 4:
- num_ints_translated = scale_bbox_to_transformed_image(
- top=num_ints[0],
- left=num_ints[1],
- bottom=num_ints[2],
- right=num_ints[3],
- scale_factor=scale_factor,
- )
- else:
- raise ValueError(f"Invalid number of ints: {len(num_ints)}")
- # Tokenize the text, skipping the
- tokens = [tokenizer.vocab[str(num)] for num in num_ints_translated]
- return [token_space_open_string] + tokens + [token_space_close_string]
-
-
-def _tokenize_prompts_with_image_and_batch(
- tokenizer,
- prompts: List[List[str]],
- scale_factors: Optional[List[List["torch.Tensor"]]],
- max_tokens_to_generate: int,
- max_position_embeddings: int,
- add_BOS: bool, # Same issue with types as above
- add_beginning_of_answer_token: bool,
-) -> Tuple["torch.Tensor", "torch.Tensor"]:
- """
- Given a set of prompts and number of tokens to generate:
- - tokenize prompts
- - set the sequence length to be the max of length of prompts plus the number of tokens we would like to generate
- - pad all the sequences to this length so we can convert them into a 3D tensor.
- """
-
- # If not tool use, tranform the coordinates while tokenizing
- if scale_factors is not None:
- transformed_prompt_tokens = []
- for prompt_seq, scale_factor_seq in zip(prompts, scale_factors):
- transformed_prompt_tokens.append(
- [
- _transform_coordinates_and_tokenize(prompt, scale_factor.item(), tokenizer)
- for prompt, scale_factor in zip(prompt_seq, scale_factor_seq)
- ]
- )
- else:
- transformed_prompt_tokens = [[tokenizer.tokenize(prompt) for prompt in prompt_seq] for prompt_seq in prompts]
-
- prompts_tokens = transformed_prompt_tokens
-
- if add_BOS:
- bos_token = tokenizer.vocab[""]
- else:
- bos_token = tokenizer.vocab["|ENDOFTEXT|"]
- prompts_tokens = [[[bos_token] + x for x in prompt_seq] for prompt_seq in prompts_tokens]
- if add_beginning_of_answer_token:
- boa = tokenizer.vocab[BEGINNING_OF_ANSWER_STRING]
- # Only add bbox open token to the last subsequence since that is what will be completed
- for token_seq in prompts_tokens:
- token_seq[-1].append(boa)
-
- # Now we have a list of list of tokens which each list has a different
- # size. We want to extend this list to:
- # - incorporate the tokens that need to be generated
- # - make all the sequences equal length.
- # Get the prompts length.
-
- prompts_length = [[len(x) for x in prompts_tokens_seq] for prompts_tokens_seq in prompts_tokens]
- # Get the max prompts length.
- max_prompt_len: int = np.max(prompts_length)
- # Number of tokens in the each sample of the batch.
- samples_length = min(max_prompt_len + max_tokens_to_generate, max_position_embeddings)
- if max_prompt_len + max_tokens_to_generate > max_position_embeddings:
- logger.warning(
- f"Max subsequence prompt length of {max_prompt_len} + max tokens to generate {max_tokens_to_generate}",
- f"exceeds context length of {max_position_embeddings}. Will generate as many tokens as possible.",
- )
- # Now update the list of list to be of the same size: samples_length.
- for prompt_tokens_seq, prompts_length_seq in zip(prompts_tokens, prompts_length):
- for prompt_tokens, prompt_length in zip(prompt_tokens_seq, prompts_length_seq):
- if len(prompt_tokens) > samples_length:
- raise ValueError("Length of subsequence prompt exceeds sequence length.")
- padding_size = samples_length - prompt_length
- prompt_tokens.extend([tokenizer.vocab["|ENDOFTEXT|"]] * padding_size)
-
- # Now we are in a structured format, we can convert to tensors.
- prompts_tokens_tensor = torch.tensor(prompts_tokens, dtype=torch.int64)
- prompts_length_tensor = torch.tensor(prompts_length, dtype=torch.int64)
-
- return prompts_tokens_tensor, prompts_length_tensor
-
-
-# Simplified assuming self.crop_top = self.padding_top = 0
-def original_to_transformed_h_coords(original_coords, scale_h):
- return np.round(original_coords * scale_h).astype(np.int32)
-
-
-# Simplified assuming self.crop_left = self.padding_left = 0
-def original_to_transformed_w_coords(original_coords, scale_w):
- return np.round(original_coords * scale_w).astype(np.int32)
-
-
-def scale_point_to_transformed_image(x: float, y: float, scale_factor: float) -> List[int]:
- x_scaled = original_to_transformed_w_coords(np.array([x / 2]), scale_factor)[0]
- y_scaled = original_to_transformed_h_coords(np.array([y / 2]), scale_factor)[0]
- return [x_scaled, y_scaled]
-
-
-def scale_bbox_to_transformed_image(
- top: float, left: float, bottom: float, right: float, scale_factor: float
-) -> List[int]:
- top_scaled = original_to_transformed_w_coords(np.array([top / 2]), scale_factor)[0]
- left_scaled = original_to_transformed_h_coords(np.array([left / 2]), scale_factor)[0]
- bottom_scaled = original_to_transformed_w_coords(np.array([bottom / 2]), scale_factor)[0]
- right_scaled = original_to_transformed_h_coords(np.array([right / 2]), scale_factor)[0]
- return [top_scaled, left_scaled, bottom_scaled, right_scaled]
-
-
-class FuyuProcessor(ProcessorMixin):
- r"""
- Constructs a Fuyu processor which wraps a Fuyu image processor and a Llama tokenizer into a single processor.
-
- [`FuyuProcessor`] offers all the functionalities of [`FuyuImageProcessor`] and [`LlamaTokenizerFast`]. See the
- [`~FuyuProcessor.__call__`] and [`~FuyuProcessor.decode`] for more information.
-
- Args:
- image_processor ([`FuyuImageProcessor`]):
- The image processor is a required input.
- tokenizer ([`LlamaTokenizerFast`]):
- The tokenizer is a required input.
- """
-
- attributes = ["image_processor", "tokenizer"]
- image_processor_class = "FuyuImageProcessor"
- tokenizer_class = "AutoTokenizer"
-
- def __init__(self, image_processor, tokenizer):
- super().__init__(image_processor=image_processor, tokenizer=tokenizer)
- self.image_processor = image_processor
- self.tokenizer = tokenizer
- self.max_tokens_to_generate = 10
- self.max_position_embeddings = 16384 # TODO Can't derive this from model files: where to set it?
- self.pad_token_id = 0
- self.dummy_image_index = -1
-
- def _left_pad_inputs_with_attention_mask(self, model_inputs: List[Dict], return_attention_mask: bool):
- max_length_input_ids = max(entry["input_ids"].shape[1] for entry in model_inputs)
- max_length_image_patch_indices = max(entry["image_patches_indices"].shape[1] for entry in model_inputs)
-
- batched_inputs = {"input_ids": [], "image_patches": [], "image_patches_indices": [], "attention_mask": []}
-
- for entry in model_inputs:
- for key, tensor in entry.items():
- if key == "input_ids":
- num_padding_tokens = max_length_input_ids - tensor.shape[1]
- padded_input_ids = torch.cat(
- [
- torch.full((tensor.shape[0], num_padding_tokens), self.pad_token_id, dtype=torch.long),
- tensor,
- ],
- dim=1,
- )
- batched_inputs[key].append(padded_input_ids)
-
- attention_mask = torch.cat(
- [torch.zeros(tensor.shape[0], num_padding_tokens, dtype=torch.long), torch.ones_like(tensor)],
- dim=1,
- )
- batched_inputs["attention_mask"].append(attention_mask)
-
- elif key == "image_patches":
- # For image_patches, we don't pad but just append them to the list.
- batched_inputs[key].append(tensor)
-
- else: # for image_patches_indices
- num_padding_indices = max_length_image_patch_indices - tensor.shape[1]
- padded_indices = torch.cat(
- [
- torch.full(
- (tensor.shape[0], num_padding_indices), self.dummy_image_index, dtype=torch.long
- ),
- tensor,
- ],
- dim=1,
- )
- batched_inputs[key].append(padded_indices)
- batched_keys = ["input_ids", "image_patches_indices"]
- if return_attention_mask:
- batched_keys.append("attention_mask")
- for key in batched_keys:
- batched_inputs[key] = torch.cat(batched_inputs[key], dim=0)
-
- return batched_inputs
-
- def get_sample_encoding(
- self,
- prompts,
- scale_factors,
- image_unpadded_heights,
- image_unpadded_widths,
- image_placeholder_id,
- image_newline_id,
- tensor_batch_images,
- ):
- image_present = torch.ones(1, 1, 1)
- model_image_input = self.image_processor.preprocess_with_tokenizer_info(
- image_input=tensor_batch_images,
- image_present=image_present,
- image_unpadded_h=image_unpadded_heights,
- image_unpadded_w=image_unpadded_widths,
- image_placeholder_id=image_placeholder_id,
- image_newline_id=image_newline_id,
- variable_sized=True,
- )
- # FIXME max_tokens_to_generate is embedded into this processor's call.
- prompt_tokens, prompts_length = _tokenize_prompts_with_image_and_batch(
- tokenizer=self.tokenizer,
- prompts=prompts,
- scale_factors=scale_factors,
- max_tokens_to_generate=self.max_tokens_to_generate,
- max_position_embeddings=self.max_position_embeddings,
- add_BOS=True,
- add_beginning_of_answer_token=True,
- )
- image_padded_unpacked_tokens = construct_full_unpacked_stream(
- num_real_text_tokens=prompts_length,
- input_stream=prompt_tokens,
- image_tokens=model_image_input["image_input_ids"],
- batch_size=1,
- num_sub_sequences=self.subsequence_length,
- )
- # Construct inputs for image patch indices.
- unpacked_image_patch_indices_per_batch = construct_full_unpacked_stream(
- num_real_text_tokens=prompts_length,
- input_stream=torch.full_like(prompt_tokens, -1),
- image_tokens=model_image_input["image_patch_indices_per_batch"],
- batch_size=1,
- num_sub_sequences=self.subsequence_length,
- )
- max_prompt_length = max(x.shape[-1] for x in image_padded_unpacked_tokens)
- max_seq_len_batch = min(max_prompt_length + self.max_tokens_to_generate, self.max_position_embeddings)
- tokens_to_place = min(max_seq_len_batch, max(0, image_padded_unpacked_tokens[0].shape[0]))
-
- # Use same packing logic for the image patch indices.
- image_patch_input_indices = full_unpacked_stream_to_tensor(
- all_bi_tokens_to_place=[tokens_to_place],
- full_unpacked_stream=unpacked_image_patch_indices_per_batch,
- fill_value=-1,
- batch_size=1,
- new_seq_len=max_seq_len_batch,
- offset=0,
- )
- image_patches_tensor = torch.stack([img[0] for img in model_image_input["image_patches"]])
- batch_encoding = {
- "input_ids": image_padded_unpacked_tokens[0].unsqueeze(0),
- "image_patches": image_patches_tensor,
- "image_patches_indices": image_patch_input_indices,
- }
- return batch_encoding
-
- def __call__(
- self,
- text=None,
- images=None,
- add_special_tokens: bool = True,
- return_attention_mask: bool = True,
- padding: Union[bool, str, PaddingStrategy] = False,
- truncation: Union[bool, str, TruncationStrategy] = None,
- max_length: Optional[int] = None,
- stride: int = 0,
- pad_to_multiple_of: Optional[int] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_token_type_ids: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- return_tensors: Optional[Union[str, TensorType]] = None,
- **kwargs,
- ) -> "FuyuBatchFeature":
- """
- Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
- and `kwargs` arguments to LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`] if `text` is not `None` to
- encode the text. To prepare the image(s), this method forwards the `images` and `kwargs` arguments to
- FuyuImageProcessor's [`~FuyuImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
- of the above two methods for more information.
-
- Args:
- text (`str`, `List[str]`):
- The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
- (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
- `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
- images (`PIL.Image.Image`, `List[PIL.Image.Image]`):
- The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
- tensor. Both channels-first and channels-last formats are supported.
-
- Returns:
- [`FuyuBatchEncoding`]: A [`FuyuBatchEncoding`] with the following fields:
-
- - **input_ids** -- Tensor of token ids to be fed to a model. Returned when `text` is not `None`.
- - **image_patches** -- List of Tensor of image patches. Returned when `images` is not `None`.
- - **image_patches_indices** -- Tensor of indices where patch embeddings have to be inserted by the model.
- - **attention_mask** -- List of indices specifying which tokens should be attended to by the model when
- `return_attention_mask=True`.
- """
- requires_backends(self, ["torch"])
-
- # --- Check input validity ---
- if not return_attention_mask:
- raise ValueError("`return_attention_mask=False` is not supported for this model.")
- if text is None and images is None:
- raise ValueError("You have to specify either text or images. Both cannot be None.")
- if text is not None and images is None:
- logger.warning("You are processing a text with no associated image. Make sure it is intended.")
- self.current_processor = self.tokenizer
- text_encoding = self.tokenizer(
- text=text,
- add_special_tokens=add_special_tokens,
- padding=padding,
- truncation=truncation,
- max_length=max_length,
- stride=stride,
- pad_to_multiple_of=pad_to_multiple_of,
- return_attention_mask=return_attention_mask,
- return_overflowing_tokens=return_overflowing_tokens,
- return_special_tokens_mask=return_special_tokens_mask,
- return_offsets_mapping=return_offsets_mapping,
- return_token_type_ids=return_token_type_ids,
- return_length=return_length,
- verbose=verbose,
- return_tensors=return_tensors,
- **kwargs,
- )
- return text_encoding
-
- if text is None and images is not None:
- logger.warning("You are processing an image with no associated text. Make sure it is intended.")
- prompts = [[""]]
- if text is not None and images is not None:
- if isinstance(text, str):
- prompts = [[text]]
- elif isinstance(text, list):
- prompts = [[text_seq] for text_seq in text]
-
- # --- Preprocess images using self.image_processor ---
-
- # FIXME - We hard code "pt" here because the rest of the processing assumes torch tensors
- image_encoding = self.image_processor.preprocess(images, return_tensors="pt")
- batch_images = image_encoding["images"]
- image_unpadded_heights = image_encoding["image_unpadded_heights"]
- image_unpadded_widths = image_encoding["image_unpadded_widths"]
- scale_factors = image_encoding["image_scale_factors"]
- self.subsequence_length = 1 # Each batch contains only one sequence.
- self.batch_size = len(batch_images)
-
- # --- Use self.tokenizer to get the ids of special tokens to insert into image ids ---
-
- image_placeholder_id = self.tokenizer("|SPEAKER|", add_special_tokens=False)["input_ids"][1]
- image_newline_id = self.tokenizer("|NEWLINE|", add_special_tokens=False)["input_ids"][1]
- tensor_batch_images = torch.stack([img[0] for img in batch_images]).unsqueeze(1)
-
- # --- Use self.image_processor again to obtain the full token ids and batch inputs ---
- all_encodings = []
-
- for prompt, scale_factor, image_unpadded_height, image_unpadded_width, tensor_batch_image in zip(
- prompts, scale_factors, image_unpadded_heights, image_unpadded_widths, tensor_batch_images
- ):
- sample_encoding = self.get_sample_encoding(
- prompts=[prompt],
- scale_factors=[scale_factor],
- image_unpadded_heights=torch.tensor([image_unpadded_height]),
- image_unpadded_widths=torch.tensor([image_unpadded_width]),
- image_placeholder_id=image_placeholder_id,
- image_newline_id=image_newline_id,
- tensor_batch_images=tensor_batch_image.unsqueeze(0),
- )
- all_encodings.append(sample_encoding)
- batch_encoding = self._left_pad_inputs_with_attention_mask(
- model_inputs=all_encodings, return_attention_mask=return_attention_mask
- )
- return FuyuBatchFeature(data=batch_encoding)
-
- def post_process_box_coordinates(self, outputs, target_sizes=None):
- """
- Transforms raw coordinates detected by [`FuyuForCausalLM`] to the original images' coordinate space.
- Coordinates will be returned in "box" format, with the following pattern:
- `top, left, bottom, right`
-
- Point coordinates are not supported yet.
-
- Args:
- outputs ([`GenerateOutput`]):
- Raw outputs from `generate`.
- target_sizes (`torch.Tensor`, *optional*):
- Tensor of shape (batch_size, 2) where each entry is the (height, width) of the corresponding image in
- the batch. If set, found coordinates in the output sequence are rescaled to the target sizes. If left
- to None, coordinates will not be rescaled.
-
- Returns:
- `GenerateOutput`: Same output type returned by `generate`, with output token ids replaced with
- boxed and possible rescaled coordinates.
- """
-
- def scale_factor_to_fit(original_size, target_size=None):
- height, width = original_size
- if target_size is None:
- max_height = self.image_processor.size["height"]
- max_width = self.image_processor.size["width"]
- else:
- max_height, max_width = target_size
- if width <= max_width and height <= max_height:
- return 1.0
- return min(max_height / height, max_width / width)
-
- def find_delimiters_pair(tokens, start_token, end_token):
- start_id = self.tokenizer.convert_tokens_to_ids(start_token)
- end_id = self.tokenizer.convert_tokens_to_ids(end_token)
-
- starting_positions = (tokens == start_id).nonzero(as_tuple=True)[0]
- ending_positions = (tokens == end_id).nonzero(as_tuple=True)[0]
-
- if torch.any(starting_positions) and torch.any(ending_positions):
- return (starting_positions[0], ending_positions[0])
- return (None, None)
-
- def tokens_to_boxes(tokens, original_size):
- while (pair := find_delimiters_pair(tokens, TOKEN_BBOX_OPEN_STRING, TOKEN_BBOX_CLOSE_STRING)) != (
- None,
- None,
- ):
- start, end = pair
- if end != start + 5:
- continue
-
- # Retrieve transformed coordinates from tokens
- coords = self.tokenizer.convert_ids_to_tokens(tokens[start + 1 : end])
-
- # Scale back to original image size and multiply by 2
- scale = scale_factor_to_fit(original_size)
- top, left, bottom, right = [2 * int(float(c) / scale) for c in coords]
-
- # Replace the IDs so they get detokenized right
- replacement = f" {TEXT_REPR_BBOX_OPEN}{top}, {left}, {bottom}, {right}{TEXT_REPR_BBOX_CLOSE}"
- replacement = self.tokenizer.tokenize(replacement)[1:]
- replacement = self.tokenizer.convert_tokens_to_ids(replacement)
- replacement = torch.tensor(replacement).to(tokens)
-
- tokens = torch.cat([tokens[:start], replacement, tokens[end + 1 :]], 0)
- return tokens
-
- def tokens_to_points(tokens, original_size):
- while (pair := find_delimiters_pair(tokens, TOKEN_POINT_OPEN_STRING, TOKEN_POINT_CLOSE_STRING)) != (
- None,
- None,
- ):
- start, end = pair
- if end != start + 3:
- continue
-
- # Retrieve transformed coordinates from tokens
- coords = self.tokenizer.convert_ids_to_tokens(tokens[start + 1 : end])
-
- # Scale back to original image size and multiply by 2
- scale = scale_factor_to_fit(original_size)
- x, y = [2 * int(float(c) / scale) for c in coords]
-
- # Replace the IDs so they get detokenized right
- replacement = f" {TEXT_REPR_POINT_OPEN}{x}, {y}{TEXT_REPR_POINT_CLOSE}"
- replacement = self.tokenizer.tokenize(replacement)[1:]
- replacement = self.tokenizer.convert_tokens_to_ids(replacement)
- replacement = torch.tensor(replacement).to(tokens)
-
- tokens = torch.cat([tokens[:start], replacement, tokens[end + 1 :]], 0)
- return tokens
-
- if target_sizes is None:
- target_sizes = ((self.image_processor.size["height"], self.image_processor.size["width"]),) * len(outputs)
- elif target_sizes.shape[1] != 2:
- raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch")
-
- if len(outputs) != len(target_sizes):
- raise ValueError("Make sure that you pass in as many target sizes as output sequences")
-
- results = []
- for seq, size in zip(outputs, target_sizes):
- seq = tokens_to_boxes(seq, size)
- seq = tokens_to_points(seq, size)
- results.append(seq)
-
- return results
-
- def batch_decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
- refer to the docstring of this method for more information.
- """
- return self.tokenizer.batch_decode(*args, **kwargs)
-
- def decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
- the docstring of this method for more information.
- """
- return self.tokenizer.decode(*args, **kwargs)
diff --git a/transformers/models/gemma/__init__.py b/transformers/models/gemma/__init__.py
deleted file mode 100644
index 64ff34453828164159c6e26220b7ea1f23d3c35c..0000000000000000000000000000000000000000
--- a/transformers/models/gemma/__init__.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_flax_available,
- is_sentencepiece_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_gemma": ["GEMMA_PRETRAINED_CONFIG_ARCHIVE_MAP", "GemmaConfig"],
-}
-
-try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_gemma"] = ["GemmaTokenizer"]
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_gemma_fast"] = ["GemmaTokenizerFast"]
-
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_gemma"] = [
- "GemmaForCausalLM",
- "GemmaModel",
- "GemmaPreTrainedModel",
- "GemmaForSequenceClassification",
- ]
-
-try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flax_gemma"] = [
- "FlaxGemmaForCausalLM",
- "FlaxGemmaModel",
- "FlaxGemmaPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_gemma import GEMMA_PRETRAINED_CONFIG_ARCHIVE_MAP, GemmaConfig
-
- try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_gemma import GemmaTokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_gemma_fast import GemmaTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_gemma import (
- GemmaForCausalLM,
- GemmaForSequenceClassification,
- GemmaModel,
- GemmaPreTrainedModel,
- )
-
- try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flax_gemma import (
- FlaxGemmaForCausalLM,
- FlaxGemmaModel,
- FlaxGemmaPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gemma/__pycache__/__init__.cpython-310.pyc b/transformers/models/gemma/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index dcaf9e4e1b5a239a805b9d3f47256516dfc1f905..0000000000000000000000000000000000000000
Binary files a/transformers/models/gemma/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gemma/__pycache__/configuration_gemma.cpython-310.pyc b/transformers/models/gemma/__pycache__/configuration_gemma.cpython-310.pyc
deleted file mode 100644
index 830b540a5a36c4db29a1856a4d4fd080b3c7e306..0000000000000000000000000000000000000000
Binary files a/transformers/models/gemma/__pycache__/configuration_gemma.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gemma/__pycache__/convert_gemma_weights_to_hf.cpython-310.pyc b/transformers/models/gemma/__pycache__/convert_gemma_weights_to_hf.cpython-310.pyc
deleted file mode 100644
index b0f64497609e7f9274d008c22dc4197e25ca8f8c..0000000000000000000000000000000000000000
Binary files a/transformers/models/gemma/__pycache__/convert_gemma_weights_to_hf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gemma/__pycache__/modeling_flax_gemma.cpython-310.pyc b/transformers/models/gemma/__pycache__/modeling_flax_gemma.cpython-310.pyc
deleted file mode 100644
index 0fc3d43b2982d196a323bca58b407a4fda256455..0000000000000000000000000000000000000000
Binary files a/transformers/models/gemma/__pycache__/modeling_flax_gemma.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gemma/__pycache__/modeling_gemma.cpython-310.pyc b/transformers/models/gemma/__pycache__/modeling_gemma.cpython-310.pyc
deleted file mode 100644
index a460a6ae91615138523a65cc9ec218dd108ce103..0000000000000000000000000000000000000000
Binary files a/transformers/models/gemma/__pycache__/modeling_gemma.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gemma/__pycache__/tokenization_gemma.cpython-310.pyc b/transformers/models/gemma/__pycache__/tokenization_gemma.cpython-310.pyc
deleted file mode 100644
index e97538df6af35fe9ccf4c554722546780987f287..0000000000000000000000000000000000000000
Binary files a/transformers/models/gemma/__pycache__/tokenization_gemma.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gemma/__pycache__/tokenization_gemma_fast.cpython-310.pyc b/transformers/models/gemma/__pycache__/tokenization_gemma_fast.cpython-310.pyc
deleted file mode 100644
index 4bcf1fb60230cfc8af29049112e4426472e24587..0000000000000000000000000000000000000000
Binary files a/transformers/models/gemma/__pycache__/tokenization_gemma_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gemma/configuration_gemma.py b/transformers/models/gemma/configuration_gemma.py
deleted file mode 100644
index 87e5a2c6693f0d44f8c32bc02a9c6a69541c95e6..0000000000000000000000000000000000000000
--- a/transformers/models/gemma/configuration_gemma.py
+++ /dev/null
@@ -1,153 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Gemma model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GEMMA_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GemmaConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GemmaModel`]. It is used to instantiate an Gemma
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the Gemma-7B.
-
- e.g. [google/gemma-7b](https://huggingface.co/google/gemma-7b)
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 256000):
- Vocabulary size of the Gemma model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`GemmaModel`]
- hidden_size (`int`, *optional*, defaults to 3072):
- Dimension of the hidden representations.
- intermediate_size (`int`, *optional*, defaults to 24576):
- Dimension of the MLP representations.
- num_hidden_layers (`int`, *optional*, defaults to 28):
- Number of hidden layers in the Transformer decoder.
- num_attention_heads (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer decoder.
- num_key_value_heads (`int`, *optional*, defaults to 16):
- This is the number of key_value heads that should be used to implement Grouped Query Attention. If
- `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
- `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
- converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
- by meanpooling all the original heads within that group. For more details checkout [this
- paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
- `num_attention_heads`.
- head_dim (`int`, *optional*, defaults to 256):
- The attention head dimension.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
- The legacy activation function. It is overwritten by the `hidden_activation`.
- hidden_activation (`str` or `function`, *optional*):
- The non-linear activation function (function or string) in the decoder. Will default to `"gelu_pytorch_tanh"`
- if not specified. `"gelu_pytorch_tanh"` uses an approximation of the `"gelu"` activation function.
- max_position_embeddings (`int`, *optional*, defaults to 8192):
- The maximum sequence length that this model might ever be used with.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- rms_norm_eps (`float`, *optional*, defaults to 1e-06):
- The epsilon used by the rms normalization layers.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- pad_token_id (`int`, *optional*, defaults to 0):
- Padding token id.
- eos_token_id (`int`, *optional*, defaults to 1):
- End of stream token id.
- bos_token_id (`int`, *optional*, defaults to 2):
- Beginning of stream token id.
- tie_word_embeddings (`bool`, *optional*, defaults to `True`):
- Whether to tie weight embeddings
- rope_theta (`float`, *optional*, defaults to 10000.0):
- The base period of the RoPE embeddings.
- attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
- Whether to use a bias in the query, key, value and output projection layers during self-attention.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
-
- ```python
- >>> from transformers import GemmaModel, GemmaConfig
-
- >>> # Initializing a Gemma gemma-7b style configuration
- >>> configuration = GemmaConfig()
-
- >>> # Initializing a model from the gemma-7b style configuration
- >>> model = GemmaModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "gemma"
- keys_to_ignore_at_inference = ["past_key_values"]
-
- def __init__(
- self,
- vocab_size=256000,
- hidden_size=3072,
- intermediate_size=24576,
- num_hidden_layers=28,
- num_attention_heads=16,
- num_key_value_heads=16,
- head_dim=256,
- hidden_act="gelu_pytorch_tanh",
- hidden_activation=None,
- max_position_embeddings=8192,
- initializer_range=0.02,
- rms_norm_eps=1e-6,
- use_cache=True,
- pad_token_id=0,
- eos_token_id=1,
- bos_token_id=2,
- tie_word_embeddings=True,
- rope_theta=10000.0,
- attention_bias=False,
- attention_dropout=0.0,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.hidden_size = hidden_size
- self.intermediate_size = intermediate_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.head_dim = head_dim
- self.num_key_value_heads = num_key_value_heads
- self.hidden_act = hidden_act
- self.hidden_activation = hidden_activation
- self.initializer_range = initializer_range
- self.rms_norm_eps = rms_norm_eps
- self.use_cache = use_cache
- self.rope_theta = rope_theta
- self.attention_bias = attention_bias
- self.attention_dropout = attention_dropout
-
- super().__init__(
- pad_token_id=pad_token_id,
- bos_token_id=bos_token_id,
- eos_token_id=eos_token_id,
- tie_word_embeddings=tie_word_embeddings,
- **kwargs,
- )
diff --git a/transformers/models/gemma/convert_gemma_weights_to_hf.py b/transformers/models/gemma/convert_gemma_weights_to_hf.py
deleted file mode 100644
index 9b71be35bfa167f4c51eb2e30a345929ea9f54ee..0000000000000000000000000000000000000000
--- a/transformers/models/gemma/convert_gemma_weights_to_hf.py
+++ /dev/null
@@ -1,206 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import argparse
-import os
-import warnings
-
-import torch
-from accelerate import init_empty_weights
-
-from transformers import GemmaConfig, GemmaForCausalLM, GemmaTokenizer
-
-
-try:
- from transformers import GemmaTokenizerFast
-except ImportError as e:
- warnings.warn(e)
- warnings.warn(
- "The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion"
- )
- GemmaTokenizerFast = None
-
-"""
-Sample usage:
-
-```
-python src/transformers/models/gemma/convert_gemma_weights_to_hf.py \
- --input_dir /path/to/downloaded/gemma/weights --model_size 7B --output_dir /output/path
-```
-
-Thereafter, models can be loaded via:
-
-```py
-from transformers import GemmaForCausalLM, GemmaTokenizerFast
-
-model = GemmaForCausalLM.from_pretrained("/output/path")
-tokenizer = GemmaTokenizerFast.from_pretrained("/output/path")
-```
-
-Important note: you need to be able to host the whole model in RAM to execute this script (even if the biggest versions
-come in several checkpoints they each contain a part of each weight of the model, so we need to load them all in RAM).
-"""
-
-gemma_2b_config = GemmaConfig(
- num_hidden_layers=18,
- num_attention_heads=8,
- num_key_value_heads=1,
- hidden_size=2048,
- intermediate_size=16384,
-)
-
-gemma_7b_config = GemmaConfig()
-
-CONFIG_MAPPING = {"2B": gemma_2b_config, "7B": gemma_7b_config}
-LAYER_NAME_MAPPING = {"embedder.weight": "model.embed_tokens.weight"}
-
-
-def write_model(save_path, input_base_path, config, safe_serialization=True, push_to_hub=False, dtype=torch.float32):
- num_attn_heads = config.num_attention_heads
- hidden_size = config.hidden_size
- num_kv_heads = config.num_key_value_heads
- head_dim = config.head_dim
-
- print(f"Fetching all parameters from the checkpoint at '{input_base_path}'")
- model_state_dict = torch.load(input_base_path, map_location="cpu")["model_state_dict"]
- model_state_dict.pop("freqs_cis")
-
- state_dict = {}
- for k, v in model_state_dict.items():
- if "qkv_proj" in k:
- if num_kv_heads == 1:
- v = v.reshape(num_attn_heads + num_kv_heads * 2, head_dim, hidden_size)
- q_proj = v[:num_attn_heads, ...]
- k_proj = v[num_attn_heads : num_attn_heads + num_kv_heads, ...].repeat(num_kv_heads, 1, 1)
- v_proj = v[-num_kv_heads:, ...].repeat(num_kv_heads, 1, 1)
-
- state_dict[k.replace("qkv_proj", "q_proj")] = q_proj.reshape(
- num_attn_heads * head_dim, hidden_size
- ).clone()
- state_dict[k.replace("qkv_proj", "k_proj")] = k_proj.reshape(
- num_kv_heads * head_dim, hidden_size
- ).clone()
- state_dict[k.replace("qkv_proj", "v_proj")] = v_proj[0].clone()
- else:
- q_proj, k_proj, v_proj = torch.split(v, v.shape[0] // 3, 0)
- state_dict[k.replace("qkv_proj", "q_proj")] = q_proj.reshape(
- num_attn_heads * head_dim, hidden_size
- ).clone()
- state_dict[k.replace("qkv_proj", "k_proj")] = k_proj.reshape(
- num_kv_heads * head_dim, hidden_size
- ).clone()
- state_dict[k.replace("qkv_proj", "v_proj")] = v_proj.clone()
-
- elif k == "embedder.weight":
- state_dict[LAYER_NAME_MAPPING[k]] = v
- state_dict["lm_head.weight"] = v
- else:
- state_dict[k] = v
-
- torch.set_default_dtype(dtype)
-
- print("Loading the checkpoint in a Gemma model.")
- with init_empty_weights():
- model = GemmaForCausalLM(config)
- model.load_state_dict(state_dict, assign=True, strict=False)
-
- model.config.torch_dtype = torch.float32
- del model.config._name_or_path
- print("Saving in the Transformers format.")
-
- if push_to_hub:
- print(f"pushing the model to {save_path}")
- model.push_to_hub(save_path, safe_serialization=safe_serialization, private=True)
- else:
- model.save_pretrained(save_path, safe_serialization=safe_serialization)
-
-
-def write_tokenizer(input_tokenizer_path, save_path, push_to_hub=False):
- # Initialize the tokenizer based on the `spm` model
- tokenizer_class = GemmaTokenizer if GemmaTokenizerFast is None else GemmaTokenizerFast
- print(f"Saving a {tokenizer_class.__name__} to {save_path}.")
- tokenizer = tokenizer_class(input_tokenizer_path)
- if push_to_hub:
- tokenizer.push_to_hub(save_path)
- else:
- tokenizer.save_pretrained(save_path)
-
-
-def main():
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--input_checkpoint",
- help="Absolute path to the target Gemma weights.",
- required=True,
- )
- parser.add_argument(
- "--tokenizer_checkpoint",
- help="Location of Gemma tokenizer model",
- )
- parser.add_argument(
- "--model_size",
- default="7B",
- choices=["2B", "7B", "tokenizer_only"],
- help="'f' models correspond to the finetuned versions, and are specific to the Gemma2 official release. For more details on Gemma2, checkout the original repo: https://huggingface.co/google/gemma-7b",
- )
- parser.add_argument(
- "--output_dir",
- default="google/gemma-7b",
- help="Location to write HF model and tokenizer",
- )
- parser.add_argument(
- "--pickle_serialization",
- help="Whether or not to save using `safetensors`.",
- action="store_true",
- default=False,
- )
- parser.add_argument(
- "--convert_tokenizer",
- help="Whether or not to convert the tokenizer as well.",
- action="store_true",
- default=False,
- )
- parser.add_argument(
- "--push_to_hub",
- help="Whether or not to push the model to the hub at `output_dir` instead of saving it locally.",
- action="store_true",
- default=False,
- )
- parser.add_argument(
- "--dtype",
- default="float32",
- help="Target dtype of the converted model",
- )
- args = parser.parse_args()
-
- if args.convert_tokenizer:
- if args.tokenizer_checkpoint is None:
- raise ValueError("Path to the tokenizer is required when passing --convert_tokenizer")
-
- spm_path = os.path.join(args.tokenizer_checkpoint)
- write_tokenizer(spm_path, args.output_dir, args.push_to_hub)
-
- config = CONFIG_MAPPING[args.model_size]
- dtype = getattr(torch, args.dtype)
- write_model(
- config=config,
- input_base_path=args.input_checkpoint,
- save_path=args.output_dir,
- safe_serialization=not args.pickle_serialization,
- push_to_hub=args.push_to_hub,
- dtype=dtype,
- )
-
-
-if __name__ == "__main__":
- main()
diff --git a/transformers/models/gemma/modeling_flax_gemma.py b/transformers/models/gemma/modeling_flax_gemma.py
deleted file mode 100644
index 235f65680fad3e9a15bcd38efb0a1b15e1e7533b..0000000000000000000000000000000000000000
--- a/transformers/models/gemma/modeling_flax_gemma.py
+++ /dev/null
@@ -1,773 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Google Inc., and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Flax Gemma model."""
-from typing import Optional, Tuple
-
-import flax.linen as nn
-import jax
-import jax.numpy as jnp
-import numpy as np
-from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
-from flax.linen import combine_masks, make_causal_mask
-from flax.linen.attention import dot_product_attention_weights
-from flax.traverse_util import flatten_dict, unflatten_dict
-from jax import lax
-
-from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput
-from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_gemma import GemmaConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "GemmaConfig"
-_CHECKPOINT_FOR_DOC = "google/gemma-2b"
-_REAL_CHECKPOINT_FOR_DOC = "openlm-research/open_llama_3b_v2"
-
-GEMMA_START_DOCSTRING = r"""
-
- This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a Flax Linen
- [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a
- regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.
-
- Finally, this model supports inherent JAX features such as:
-
- - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
-
- Parameters:
- config ([`GemmaConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.
- dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
- The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16`, or
- `jax.numpy.bfloat16`.
-
- This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
- specified all the computation will be performed with the given `dtype`.
-
- **Note that this only specifies the dtype of the computation and does not influence the dtype of model
- parameters.**
-
- If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and
- [`~FlaxPreTrainedModel.to_bf16`].
-"""
-
-GEMMA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`numpy.ndarray` of shape `(batch_size, input_ids_length)`):
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
- it.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
- `past_key_values`).
-
- If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
- and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
- information on the default strategy.
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
- position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- past_key_values (`Dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):
- Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
- auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-def create_sinusoidal_positions(num_pos, dim):
- inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2)[: (dim // 2)] / dim))
- freqs = np.einsum("i , j -> i j", np.arange(num_pos), inv_freq).astype("float32")
-
- emb = np.concatenate((freqs, freqs), axis=-1)
- out = np.concatenate((np.sin(emb)[:, None, :], np.cos(emb)[:, None, :]), axis=-1)
- return jnp.array(out[:, :, :num_pos])
-
-
-# Copied from transformers.models.llama.modeling_flax_llama.rotate_half
-def rotate_half(tensor):
- """Rotates half the hidden dims of the input."""
- rotate_half_tensor = jnp.concatenate(
- (-tensor[..., tensor.shape[-1] // 2 :], tensor[..., : tensor.shape[-1] // 2]), axis=-1
- )
- return rotate_half_tensor
-
-
-# Copied from transformers.models.llama.modeling_flax_llama.apply_rotary_pos_emb
-def apply_rotary_pos_emb(tensor, sin_pos, cos_pos):
- return (tensor * cos_pos) + (rotate_half(tensor) * sin_pos)
-
-
-class FlaxGemmaRMSNorm(nn.Module):
- config: GemmaConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.epsilon = self.config.rms_norm_eps
- self.weight = self.param("weight", lambda _, shape: jnp.ones(shape), self.config.hidden_size)
-
- def __call__(self, hidden_states):
- variance = jnp.asarray(hidden_states, dtype=jnp.float32)
- variance = jnp.power(variance, 2)
- variance = variance.mean(-1, keepdims=True)
- # use `jax.numpy.sqrt` as `jax.lax.rsqrt` does not match `torch.rsqrt`
- hidden_states = hidden_states / jnp.sqrt(variance + self.epsilon)
-
- return (1 + self.weight) * jnp.asarray(hidden_states, dtype=self.dtype)
-
-
-# Copied from transformers.models.llama.modeling_flax_llama.FlaxLlamaRotaryEmbedding with Llama->Gemma
-class FlaxGemmaRotaryEmbedding(nn.Module):
- config: GemmaConfig
- dtype: jnp.dtype = jnp.float32
-
- # Ignore copy
- def setup(self):
- head_dim = self.config.head_dim
- self.sincos = create_sinusoidal_positions(self.config.max_position_embeddings, head_dim)
-
- def __call__(self, key, query, position_ids):
- sincos = self.sincos[position_ids]
- sin_pos, cos_pos = jnp.split(sincos, 2, axis=-1)
-
- key = apply_rotary_pos_emb(key, sin_pos, cos_pos)
- query = apply_rotary_pos_emb(query, sin_pos, cos_pos)
-
- key = jnp.asarray(key, dtype=self.dtype)
- query = jnp.asarray(query, dtype=self.dtype)
-
- return key, query
-
-
-class FlaxGemmaAttention(nn.Module):
- config: GemmaConfig
- dtype: jnp.dtype = jnp.float32
- causal: bool = True
- is_cross_attention: bool = False
-
- def setup(self):
- config = self.config
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = config.head_dim
- self.attention_softmax_in_fp32 = self.dtype is not jnp.float32
-
- self.num_key_value_heads = config.num_key_value_heads
- self.num_key_value_groups = self.num_heads // self.num_key_value_heads
-
- kernel = jax.nn.initializers.normal(self.config.initializer_range)
- self.q_proj = nn.Dense(
- self.num_heads * self.head_dim, use_bias=config.attention_bias, dtype=self.dtype, kernel_init=kernel
- )
- self.k_proj = nn.Dense(
- self.num_key_value_heads * self.head_dim,
- use_bias=config.attention_bias,
- dtype=self.dtype,
- kernel_init=kernel,
- )
- self.v_proj = nn.Dense(
- self.num_key_value_heads * self.head_dim,
- use_bias=config.attention_bias,
- dtype=self.dtype,
- kernel_init=kernel,
- )
- self.o_proj = nn.Dense(self.embed_dim, use_bias=config.attention_bias, dtype=self.dtype, kernel_init=kernel)
-
- self.causal_mask = make_causal_mask(jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool")
- self.rotary_emb = FlaxGemmaRotaryEmbedding(config, dtype=self.dtype)
-
- def _split_heads(self, hidden_states, num_heads):
- return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim))
-
- def _merge_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads * self.head_dim,))
-
- @nn.compact
- # Copied from transformers.models.gpt_neo.modeling_flax_gpt_neo.FlaxGPTNeoSelfAttention._concatenate_to_cache
- def _concatenate_to_cache(self, key, value, query, attention_mask):
- """
- This function takes projected key, value states from a single input token and concatenates the states to cached
- states from previous steps. This function is slighly adapted from the official Flax repository:
- https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252
- """
- # detect if we're initializing by absence of existing cache data.
- is_initialized = self.has_variable("cache", "cached_key")
- cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype)
- cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype)
- cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
-
- if is_initialized:
- *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape
- # update key, value caches with our new 1d spatial slices
- cur_index = cache_index.value
- indices = (0,) * len(batch_dims) + (cur_index, 0, 0)
- key = lax.dynamic_update_slice(cached_key.value, key, indices)
- value = lax.dynamic_update_slice(cached_value.value, value, indices)
- cached_key.value = key
- cached_value.value = value
- num_updated_cache_vectors = query.shape[1]
- cache_index.value = cache_index.value + num_updated_cache_vectors
- # causal mask for cached decoder self-attention: our single query position should only attend to those key positions that have already been generated and cached, not the remaining zero elements.
- pad_mask = jnp.broadcast_to(
- jnp.arange(max_length) < cur_index + num_updated_cache_vectors,
- tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),
- )
- attention_mask = combine_masks(pad_mask, attention_mask)
- return key, value, attention_mask
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- position_ids,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- query = self.q_proj(hidden_states)
- key = self.k_proj(hidden_states)
- value = self.v_proj(hidden_states)
-
- query = self._split_heads(query, self.num_heads)
- key = self._split_heads(key, self.num_key_value_heads)
- value = self._split_heads(value, self.num_key_value_heads)
-
- key, query = self.rotary_emb(key, query, position_ids)
-
- query_length, key_length = query.shape[1], key.shape[1]
-
- if self.has_variable("cache", "cached_key"):
- mask_shift = self.variables["cache"]["cache_index"]
- max_decoder_length = self.variables["cache"]["cached_key"].shape[1]
- causal_mask = lax.dynamic_slice(
- self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)
- )
- else:
- causal_mask = self.causal_mask[:, :, :query_length, :key_length]
-
- batch_size = hidden_states.shape[0]
- causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])
-
- attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
- attention_mask = combine_masks(attention_mask, causal_mask)
-
- dropout_rng = None
- if not deterministic and self.config.attention_dropout > 0.0:
- dropout_rng = self.make_rng("dropout")
-
- # During fast autoregressive decoding, we feed one position at a time,
- # and cache the keys and values step by step.
- if self.has_variable("cache", "cached_key") or init_cache:
- key, value, attention_mask = self._concatenate_to_cache(key, value, query, attention_mask)
-
- # transform boolean mask into float mask
- attention_bias = lax.select(
- attention_mask > 0,
- jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
- jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
- )
-
- key = jnp.repeat(key, repeats=self.num_key_value_groups, axis=2)
- value = jnp.repeat(value, repeats=self.num_key_value_groups, axis=2)
-
- # usual dot product attention
- attention_dtype = jnp.float32 if self.attention_softmax_in_fp32 else self.dtype
- attn_weights = dot_product_attention_weights(
- query,
- key,
- bias=attention_bias,
- dropout_rng=dropout_rng,
- dropout_rate=self.config.attention_dropout,
- deterministic=deterministic,
- dtype=attention_dtype,
- )
-
- if self.attention_softmax_in_fp32:
- attn_weights = attn_weights.astype(self.dtype)
-
- attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value)
- attn_output = self._merge_heads(attn_output)
- attn_output = self.o_proj(attn_output)
-
- outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)
- return outputs
-
-
-class FlaxGemmaMLP(nn.Module):
- config: GemmaConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- embed_dim = self.config.hidden_size
- inner_dim = self.config.intermediate_size if self.config.intermediate_size is not None else 4 * embed_dim
-
- kernel_init = jax.nn.initializers.normal(self.config.initializer_range)
- if self.config.hidden_activation is None:
- logger.warning_once(
- "Gemma's activation function should be approximate GeLU and not exact GeLU. "
- "Changing the activation function to `gelu_pytorch_tanh`."
- f"if you want to use the legacy `{self.config.hidden_act}`, "
- f"edit the `model.config` to set `hidden_activation={self.config.hidden_act}` "
- " instead of `hidden_act`. See https://github.com/huggingface/transformers/pull/29402 for more details."
- )
- hidden_activation = "gelu_pytorch_tanh"
- else:
- hidden_activation = self.config.hidden_activation
- self.act = ACT2FN[hidden_activation]
-
- self.gate_proj = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, kernel_init=kernel_init)
- self.down_proj = nn.Dense(embed_dim, use_bias=False, dtype=self.dtype, kernel_init=kernel_init)
- self.up_proj = nn.Dense(inner_dim, use_bias=False, dtype=self.dtype, kernel_init=kernel_init)
-
- def __call__(self, hidden_states):
- up_proj_states = self.up_proj(hidden_states)
- gate_states = self.act(self.gate_proj(hidden_states))
-
- hidden_states = self.down_proj(up_proj_states * gate_states)
- return hidden_states
-
-
-# Copied from transformers.models.llama.modeling_flax_llama.FlaxLlamaDecoderLayer with Llama->Gemma
-class FlaxGemmaDecoderLayer(nn.Module):
- config: GemmaConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.input_layernorm = FlaxGemmaRMSNorm(self.config, dtype=self.dtype)
- self.self_attn = FlaxGemmaAttention(self.config, dtype=self.dtype)
- self.post_attention_layernorm = FlaxGemmaRMSNorm(self.config, dtype=self.dtype)
- self.mlp = FlaxGemmaMLP(self.config, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- position_ids=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- residual = hidden_states
- hidden_states = self.input_layernorm(hidden_states)
- outputs = self.self_attn(
- hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
- # residual connection
- attn_output = outputs[0]
- hidden_states = residual + attn_output
-
- residual = hidden_states
- hidden_states = self.post_attention_layernorm(hidden_states)
- hidden_states = self.mlp(hidden_states)
- # residual connection
- hidden_states = residual + hidden_states
-
- return (hidden_states,) + outputs[1:]
-
-
-# Copied from transformers.models.gpt_neo.modeling_flax_gpt_neo.FlaxGPTNeoPreTrainedModel with GPTNeo->Gemma, GPT_NEO->GEMMA, transformer->model
-class FlaxGemmaPreTrainedModel(FlaxPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GemmaConfig
- base_model_prefix = "model"
- module_class: nn.Module = None
-
- def __init__(
- self,
- config: GemmaConfig,
- input_shape: Tuple = (1, 1),
- seed: int = 0,
- dtype: jnp.dtype = jnp.float32,
- _do_init: bool = True,
- **kwargs,
- ):
- module = self.module_class(config=config, dtype=dtype, **kwargs)
- super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
-
- def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
- # init input tensors
- input_ids = jnp.zeros(input_shape, dtype="i4")
- attention_mask = jnp.ones_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape)
- params_rng, dropout_rng = jax.random.split(rng)
- rngs = {"params": params_rng, "dropout": dropout_rng}
-
- random_params = self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False)["params"]
-
- if params is not None:
- random_params = flatten_dict(unfreeze(random_params))
- params = flatten_dict(unfreeze(params))
- for missing_key in self._missing_keys:
- params[missing_key] = random_params[missing_key]
- self._missing_keys = set()
- return freeze(unflatten_dict(params))
- else:
- return random_params
-
- def init_cache(self, batch_size, max_length):
- r"""
- Args:
- batch_size (`int`):
- batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
- max_length (`int`):
- maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
- cache.
- """
- # init input variables to retrieve cache
- input_ids = jnp.ones((batch_size, max_length))
- attention_mask = jnp.ones_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
-
- init_variables = self.module.init(
- jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True
- )
- return unfreeze(init_variables["cache"])
-
- @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- position_ids=None,
- params: dict = None,
- past_key_values: dict = None,
- dropout_rng: jax.random.PRNGKey = None,
- train: bool = False,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ):
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- batch_size, sequence_length = input_ids.shape
-
- if position_ids is None:
- if past_key_values is not None:
- raise ValueError("Make sure to provide `position_ids` when passing `past_key_values`.")
-
- position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
-
- if attention_mask is None:
- attention_mask = jnp.ones((batch_size, sequence_length))
-
- # Handle any PRNG if needed
- rngs = {}
- if dropout_rng is not None:
- rngs["dropout"] = dropout_rng
-
- inputs = {"params": params or self.params}
-
- # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be changed by FlaxGemmaAttention module
- if past_key_values:
- inputs["cache"] = past_key_values
- mutable = ["cache"]
- else:
- mutable = False
-
- outputs = self.module.apply(
- inputs,
- jnp.array(input_ids, dtype="i4"),
- jnp.array(attention_mask, dtype="i4"),
- jnp.array(position_ids, dtype="i4"),
- not train,
- False,
- output_attentions,
- output_hidden_states,
- return_dict,
- rngs=rngs,
- mutable=mutable,
- )
-
- # add updated cache to model output
- if past_key_values is not None and return_dict:
- outputs, past_key_values = outputs
- outputs["past_key_values"] = unfreeze(past_key_values["cache"])
- return outputs
- elif past_key_values is not None and not return_dict:
- outputs, past_key_values = outputs
- outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:]
-
- return outputs
-
-
-# Copied from transformers.models.llama.modeling_flax_llama.FlaxLlamaLayerCollection with Llama->Gemma
-class FlaxGemmaLayerCollection(nn.Module):
- config: GemmaConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.blocks = [
- FlaxGemmaDecoderLayer(self.config, dtype=self.dtype, name=str(i))
- for i in range(self.config.num_hidden_layers)
- ]
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- position_ids=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = False,
- ):
- all_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
-
- for block in self.blocks:
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
- layer_outputs = block(
- hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions += (layer_outputs[1],)
-
- # this contains possible `None` values - `FlaxGemmaModule` will filter them out
- outputs = (hidden_states, all_hidden_states, all_attentions)
-
- return outputs
-
-
-# Copied from transformers.models.llama.modeling_flax_llama.FlaxLlamaModule with Llama->Gemma
-class FlaxGemmaModule(nn.Module):
- config: GemmaConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.hidden_size = self.config.hidden_size
- embedding_init = jax.nn.initializers.normal(stddev=self.config.initializer_range)
- self.embed_tokens = nn.Embed(
- self.config.vocab_size,
- self.hidden_size,
- embedding_init=embedding_init,
- dtype=self.dtype,
- )
- self.layers = FlaxGemmaLayerCollection(self.config, dtype=self.dtype)
- self.norm = FlaxGemmaRMSNorm(self.config, dtype=self.dtype)
-
- # Ignore copy
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- position_ids=None,
- deterministic=True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- input_embeds = self.embed_tokens(input_ids.astype("i4"))
-
- input_embeds = input_embeds * (self.config.hidden_size**0.5)
-
- outputs = self.layers(
- input_embeds,
- position_ids=position_ids,
- attention_mask=attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- hidden_states = self.norm(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = outputs[1] + (hidden_states,)
- outputs = (hidden_states, all_hidden_states) + outputs[2:]
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- if not return_dict:
- return tuple(v for v in outputs if v is not None)
-
- return FlaxBaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=outputs[1],
- attentions=outputs[-1],
- )
-
-
-@add_start_docstrings(
- "The bare Gemma Model transformer outputting raw hidden-states without any specific head on top.",
- GEMMA_START_DOCSTRING,
-)
-# Copied from transformers.models.llama.modeling_flax_llama.FlaxLlamaModel with Llama->Gemma
-class FlaxGemmaModel(FlaxGemmaPreTrainedModel):
- module_class = FlaxGemmaModule
-
-
-append_call_sample_docstring(
- FlaxGemmaModel,
- _CHECKPOINT_FOR_DOC,
- FlaxBaseModelOutput,
- _CONFIG_FOR_DOC,
- real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
-)
-
-
-# Copied from transformers.models.llama.modeling_flax_llama.FlaxLlamaForCausalLMModule with Llama->Gemma
-class FlaxGemmaForCausalLMModule(nn.Module):
- config: GemmaConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.model = FlaxGemmaModule(self.config, dtype=self.dtype)
- self.lm_head = nn.Dense(
- self.config.vocab_size,
- use_bias=False,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
-
- # Ignore copy
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- position_ids=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- outputs = self.model(
- input_ids,
- position_ids=position_ids,
- attention_mask=attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- if self.config.tie_word_embeddings:
- shared_kernel = self.model.variables["params"]["embed_tokens"]["embedding"].T
- lm_logits = self.lm_head.apply({"params": {"kernel": shared_kernel}}, hidden_states)
- else:
- lm_logits = self.lm_head(hidden_states)
-
- if not return_dict:
- return (lm_logits,) + outputs[1:]
-
- return FlaxCausalLMOutput(logits=lm_logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
-
-
-@add_start_docstrings(
- """
- The Gemma Model transformer with a language modeling head (linear layer) on top.
- """,
- GEMMA_START_DOCSTRING,
-)
-# Copied from transformers.models.gptj.modeling_flax_gptj.FlaxGPTJForCausalLM with GPTJ->Gemma
-class FlaxGemmaForCausalLM(FlaxGemmaPreTrainedModel):
- module_class = FlaxGemmaForCausalLMModule
-
- def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jax.Array] = None):
- # initializing the cache
- batch_size, seq_length = input_ids.shape
-
- past_key_values = self.init_cache(batch_size, max_length)
- # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
- # But since Gemma uses a causal mask, those positions are masked anyways.
- # Thus we can create a single static attention_mask here, which is more efficient for compilation
- extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
- if attention_mask is not None:
- position_ids = attention_mask.cumsum(axis=-1) - 1
- extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, attention_mask, (0, 0))
- else:
- position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))
-
- return {
- "past_key_values": past_key_values,
- "attention_mask": extended_attention_mask,
- "position_ids": position_ids,
- }
-
- def update_inputs_for_generation(self, model_outputs, model_kwargs):
- model_kwargs["past_key_values"] = model_outputs.past_key_values
- model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1
- return model_kwargs
-
-
-append_call_sample_docstring(
- FlaxGemmaForCausalLM,
- _CHECKPOINT_FOR_DOC,
- FlaxCausalLMOutput,
- _CONFIG_FOR_DOC,
- real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
-)
diff --git a/transformers/models/gemma/modeling_gemma.py b/transformers/models/gemma/modeling_gemma.py
deleted file mode 100644
index 6077259d0b0fac7b4bf09c69664c3878dc676279..0000000000000000000000000000000000000000
--- a/transformers/models/gemma/modeling_gemma.py
+++ /dev/null
@@ -1,1372 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved.
-#
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch Gemma model."""
-
-import math
-import warnings
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...cache_utils import Cache, DynamicCache, StaticCache
-from ...modeling_attn_mask_utils import (
- AttentionMaskConverter,
- _prepare_4d_causal_attention_mask,
-)
-from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
-from ...utils import (
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- logging,
- replace_return_docstrings,
-)
-from ...utils.import_utils import is_torch_fx_available
-from .configuration_gemma import GemmaConfig
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-
-# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
-# It means that the function will not be traced through and simply appear as a node in the graph.
-if is_torch_fx_available():
- if not is_torch_greater_or_equal_than_1_13:
- import torch.fx
-
- _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "GemmaConfig"
-
-
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-class GemmaRMSNorm(nn.Module):
- def __init__(self, dim: int, eps: float = 1e-6):
- super().__init__()
- self.eps = eps
- self.weight = nn.Parameter(torch.zeros(dim))
-
- def _norm(self, x):
- return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
- def forward(self, x):
- output = self._norm(x.float())
- # Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16)
- # See https://github.com/huggingface/transformers/pull/29402
- output = output * (1.0 + self.weight.float())
- return output.type_as(x)
-
-
-ALL_LAYERNORM_LAYERS.append(GemmaRMSNorm)
-
-
-class GemmaRotaryEmbedding(nn.Module):
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
- super().__init__()
-
- self.dim = dim
- self.max_position_embeddings = max_position_embeddings
- self.base = base
- self.register_buffer("inv_freq", None, persistent=False)
-
- @torch.no_grad()
- def forward(self, x, position_ids, seq_len=None):
- # x: [bs, num_attention_heads, seq_len, head_size]
- if self.inv_freq is None:
- self.inv_freq = 1.0 / (
- self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim)
- )
- inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
- position_ids_expanded = position_ids[:, None, :].float()
- # Force float32 since bfloat16 loses precision on long contexts
- # See https://github.com/huggingface/transformers/pull/29285
- device_type = x.device.type
- device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
- with torch.autocast(device_type=device_type, enabled=False):
- freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
- emb = torch.cat((freqs, freqs), dim=-1)
- cos = emb.cos()
- sin = emb.sin()
- return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-# Copied from transformers.models.llama.modeling_llama.rotate_half
-def rotate_half(x):
- """Rotates half the hidden dims of the input."""
- x1 = x[..., : x.shape[-1] // 2]
- x2 = x[..., x.shape[-1] // 2 :]
- return torch.cat((-x2, x1), dim=-1)
-
-
-# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
- """Applies Rotary Position Embedding to the query and key tensors.
-
- Args:
- q (`torch.Tensor`): The query tensor.
- k (`torch.Tensor`): The key tensor.
- cos (`torch.Tensor`): The cosine part of the rotary embedding.
- sin (`torch.Tensor`): The sine part of the rotary embedding.
- position_ids (`torch.Tensor`, *optional*):
- Deprecated and unused.
- unsqueeze_dim (`int`, *optional*, defaults to 1):
- The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
- sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
- that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
- k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
- cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
- the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
- Returns:
- `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
- """
- cos = cos.unsqueeze(unsqueeze_dim)
- sin = sin.unsqueeze(unsqueeze_dim)
- q_embed = (q * cos) + (rotate_half(q) * sin)
- k_embed = (k * cos) + (rotate_half(k) * sin)
- return q_embed, k_embed
-
-
-class GemmaMLP(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.hidden_size = config.hidden_size
- self.intermediate_size = config.intermediate_size
- self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
- self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
- self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
- if config.hidden_activation is None:
- logger.warning_once(
- "Gemma's activation function should be approximate GeLU and not exact GeLU.\n"
- "Changing the activation function to `gelu_pytorch_tanh`."
- f"if you want to use the legacy `{config.hidden_act}`, "
- f"edit the `model.config` to set `hidden_activation={config.hidden_act}` "
- " instead of `hidden_act`. See https://github.com/huggingface/transformers/pull/29402 for more details."
- )
- hidden_activation = "gelu_pytorch_tanh"
- else:
- hidden_activation = config.hidden_activation
- self.act_fn = ACT2FN[hidden_activation]
-
- def forward(self, x):
- return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-
-
-# Copied from transformers.models.llama.modeling_llama.repeat_kv
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
- """
- This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
- num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
- """
- batch, num_key_value_heads, slen, head_dim = hidden_states.shape
- if n_rep == 1:
- return hidden_states
- hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
- return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-class GemmaAttention(nn.Module):
- """Multi-headed attention from 'Attention Is All You Need' paper"""
-
- # Ignore copy
- def __init__(self, config: GemmaConfig, layer_idx: Optional[int] = None):
- super().__init__()
- self.config = config
- self.layer_idx = layer_idx
- if layer_idx is None:
- logger.warning_once(
- f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
- "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
- "when creating this class."
- )
-
- self.attention_dropout = config.attention_dropout
- self.hidden_size = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = config.head_dim
- self.num_key_value_heads = config.num_key_value_heads
- self.num_key_value_groups = self.num_heads // self.num_key_value_heads
- self.max_position_embeddings = config.max_position_embeddings
- self.rope_theta = config.rope_theta
- self.is_causal = True
-
- if self.hidden_size % self.num_heads != 0:
- raise ValueError(
- f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
- f" and `num_heads`: {self.num_heads})."
- )
-
- self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
- self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
- self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
- self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
- self.rotary_emb = GemmaRotaryEmbedding(
- self.head_dim,
- max_position_embeddings=self.max_position_embeddings,
- base=self.rope_theta,
- )
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- bsz, q_len, _ = hidden_states.size()
-
- query_states = self.q_proj(hidden_states)
- key_states = self.k_proj(hidden_states)
- value_states = self.v_proj(hidden_states)
-
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- past_key_value = getattr(self, "past_key_value", past_key_value)
- cos, sin = self.rotary_emb(value_states, position_ids, seq_len=None)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, None)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; cache_position needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
- key_states = repeat_kv(key_states, self.num_key_value_groups)
- value_states = repeat_kv(value_states, self.num_key_value_groups)
-
- attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-
- if attention_mask is not None: # no matter the length, we just slice it
- causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
- attn_weights = attn_weights + causal_mask
-
- # upcast attention to fp32
- attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
- attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
- attn_output = torch.matmul(attn_weights, value_states)
-
- if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.transpose(1, 2).contiguous()
-
- attn_output = attn_output.view(bsz, q_len, -1)
- attn_output = self.o_proj(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights, past_key_value
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->Gemma
-class GemmaFlashAttention2(GemmaAttention):
- """
- Gemma flash attention module. This module inherits from `GemmaAttention` as the weights of the module stays
- untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
- flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- # Ignore copy
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- output_attentions = False
-
- bsz, q_len, _ = hidden_states.size()
-
- query_states = self.q_proj(hidden_states)
- key_states = self.k_proj(hidden_states)
- value_states = self.v_proj(hidden_states)
-
- # Flash attention requires the input to have the shape
- # batch_size x seq_length x head_dim x hidden_dim
- # therefore we just need to keep the original shape
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- cos, sin = self.rotary_emb(value_states, position_ids, seq_len=None)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, None)
-
- past_key_value = getattr(self, "past_key_value", past_key_value)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; cache_position needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
- # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
- # to be able to avoid many of these transpose/reshape/view.
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.transpose(1, 2)
-
- dropout_rate = self.attention_dropout if self.training else 0.0
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in the correct dtype just to be sure everything works as expected.
- # This might slowdown training & inference so it is recommended to not cast the LayerNorms
- # in fp32. (GemmaRMSNorm handles it correctly)
-
- input_dtype = query_states.dtype
- if input_dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.q_proj.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
-
- query_states = query_states.to(target_dtype)
- key_states = key_states.to(target_dtype)
- value_states = value_states.to(target_dtype)
-
- attn_output = self._flash_attention_forward(
- query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
- )
-
- attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
- attn_output = self.o_proj(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights, past_key_value
-
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in GemmaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Gemma
-class GemmaSdpaAttention(GemmaAttention):
- """
- Gemma attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
- `GemmaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
- SDPA API.
- """
-
- # Ignore copy
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Cache] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- if output_attentions:
- # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
- logger.warning_once(
- "GemmaModel is using GemmaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
- 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
- )
- return super().forward(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- )
-
- bsz, q_len, _ = hidden_states.size()
-
- query_states = self.q_proj(hidden_states)
- key_states = self.k_proj(hidden_states)
- value_states = self.v_proj(hidden_states)
-
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
- cos, sin = self.rotary_emb(value_states, position_ids, seq_len=None)
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, None)
-
- past_key_value = getattr(self, "past_key_value", past_key_value)
-
- if past_key_value is not None:
- # sin and cos are specific to RoPE models; cache_position needed for the static cache
- cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
- key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
- key_states = repeat_kv(key_states, self.num_key_value_groups)
- value_states = repeat_kv(value_states, self.num_key_value_groups)
-
- causal_mask = attention_mask
- if attention_mask is not None:
- causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
-
- # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
- # Reference: https://github.com/pytorch/pytorch/issues/112577.
- if query_states.device.type == "cuda" and causal_mask is not None:
- query_states = query_states.contiguous()
- key_states = key_states.contiguous()
- value_states = value_states.contiguous()
-
- # In case we are not compiling, we may set `causal_mask` to None, which is required to dispatch to SDPA's Flash Attention 2 backend, rather
- # relying on the `is_causal` argument.
- attn_output = torch.nn.functional.scaled_dot_product_attention(
- query_states,
- key_states,
- value_states,
- attn_mask=causal_mask,
- dropout_p=self.attention_dropout if self.training else 0.0,
- is_causal=causal_mask is None and q_len > 1,
- )
-
- attn_output = attn_output.transpose(1, 2).contiguous()
- attn_output = attn_output.view(bsz, q_len, -1)
-
- attn_output = self.o_proj(attn_output)
-
- return attn_output, None, past_key_value
-
-
-GEMMA_ATTENTION_CLASSES = {
- "eager": GemmaAttention,
- "flash_attention_2": GemmaFlashAttention2,
- "sdpa": GemmaSdpaAttention,
-}
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaDecoderLayer with LLAMA->GEMMA,Llama->Gemma
-class GemmaDecoderLayer(nn.Module):
- def __init__(self, config: GemmaConfig, layer_idx: int):
- super().__init__()
- self.hidden_size = config.hidden_size
-
- self.self_attn = GEMMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
-
- self.mlp = GemmaMLP(config)
- self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
- self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: Optional[bool] = False,
- use_cache: Optional[bool] = False,
- cache_position: Optional[torch.LongTensor] = None,
- **kwargs,
- ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`, *optional*):
- attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
- query_sequence_length, key_sequence_length)` if default attention is used.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
- """
- if "padding_mask" in kwargs:
- warnings.warn(
- "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
- )
-
- residual = hidden_states
-
- hidden_states = self.input_layernorm(hidden_states)
-
- # Self Attention
- hidden_states, self_attn_weights, present_key_value = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- **kwargs,
- )
- hidden_states = residual + hidden_states
-
- # Fully Connected
- residual = hidden_states
- hidden_states = self.post_attention_layernorm(hidden_states)
- hidden_states = self.mlp(hidden_states)
- hidden_states = residual + hidden_states
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights,)
-
- if use_cache:
- outputs += (present_key_value,)
-
- return outputs
-
-
-GEMMA_START_DOCSTRING = r"""
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`GemmaConfig`]):
- Model configuration class with all the parameters of the model. Initializing with a config file does not
- load the weights associated with the model, only the configuration. Check out the
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-
-@add_start_docstrings(
- "The bare Gemma Model outputting raw hidden-states without any specific head on top.",
- GEMMA_START_DOCSTRING,
-)
-class GemmaPreTrainedModel(PreTrainedModel):
- config_class = GemmaConfig
- base_model_prefix = "model"
- supports_gradient_checkpointing = True
- _keep_in_fp32_modules = ["inv_freq", "rotary_emb", "cos_cached", "sin_cached"]
- _no_split_modules = ["GemmaDecoderLayer"]
- _skip_keys_device_placement = ["past_key_values", "causal_mask"]
- _supports_flash_attn_2 = True
- _supports_sdpa = True
- _supports_cache_class = True
-
- def _init_weights(self, module):
- std = self.config.initializer_range
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=std)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
-
- def _setup_cache(self, cache_cls, max_batch_size, max_cache_len: Optional[int] = None):
- if self.config._attn_implementation == "flash_attention_2" and cache_cls == StaticCache:
- raise ValueError(
- "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "
- "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers"
- )
-
- for layer in self.model.layers:
- weights = layer.self_attn.o_proj.weight
- layer.self_attn.past_key_value = cache_cls(
- self.config, max_batch_size, max_cache_len, device=weights.device, dtype=weights.dtype
- )
-
- def _reset_cache(self):
- for layer in self.model.layers:
- layer.self_attn.past_key_value = None
-
-
-GEMMA_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
- it.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
- `past_key_values`).
-
- If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
- and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
- information on the default strategy.
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
- Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
- blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
- returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
-
- Two formats are allowed:
- - a [`~cache_utils.Cache`] instance;
- - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
- shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
- cache format.
-
- The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
- legacy cache format will be returned.
-
- If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
- have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
- of shape `(batch_size, sequence_length)`.
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
- Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
- this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
- the complete sequence length.
-"""
-
-
-@add_start_docstrings(
- "The bare Gemma Model outputting raw hidden-states without any specific head on top.",
- GEMMA_START_DOCSTRING,
-)
-# Copied from transformers.models.llama.modeling_llama.LlamaModel with LLAMA->GEMMA,Llama->Gemma
-class GemmaModel(GemmaPreTrainedModel):
- """
- Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`GemmaDecoderLayer`]
-
- Args:
- config: GemmaConfig
- """
-
- def __init__(self, config: GemmaConfig):
- super().__init__(config)
- self.padding_idx = config.pad_token_id
- self.vocab_size = config.vocab_size
-
- self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
- self.layers = nn.ModuleList(
- [GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
- )
- self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embed_tokens
-
- def set_input_embeddings(self, value):
- self.embed_tokens = value
-
- @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
- # Ignore copy
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Union[Tuple, BaseModelOutputWithPast]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if (input_ids is None) ^ (inputs_embeds is not None):
- raise ValueError(
- "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
- )
-
- if self.gradient_checkpointing and self.training and use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
- )
- use_cache = False
-
- if inputs_embeds is None:
- inputs_embeds = self.embed_tokens(input_ids)
-
- past_seen_tokens = 0
- if use_cache: # kept for BC (cache positions)
- if not isinstance(past_key_values, StaticCache):
- past_key_values = DynamicCache.from_legacy_cache(past_key_values)
- past_seen_tokens = past_key_values.get_seq_length()
-
- if cache_position is None:
- cache_position = torch.arange(
- past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
- )
-
- if position_ids is None:
- position_ids = cache_position.unsqueeze(0)
-
- causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_seen_tokens)
-
- # embed positions
- hidden_states = inputs_embeds
-
- # normalized
- # Gemma downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5
- # See https://github.com/huggingface/transformers/pull/29402
- normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)
- hidden_states = hidden_states * normalizer
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- next_decoder_cache = None
-
- for decoder_layer in self.layers:
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- decoder_layer.__call__,
- hidden_states,
- causal_mask,
- position_ids,
- past_key_values,
- output_attentions,
- use_cache,
- cache_position,
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- attention_mask=causal_mask,
- position_ids=position_ids,
- past_key_value=past_key_values,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cache_position=cache_position,
- )
-
- hidden_states = layer_outputs[0]
-
- if use_cache:
- next_decoder_cache = layer_outputs[2 if output_attentions else 1]
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- hidden_states = self.norm(hidden_states)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- next_cache = None
- if use_cache:
- next_cache = (
- next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache
- )
- if not return_dict:
- return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=next_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- )
-
- def _update_causal_mask(
- self,
- attention_mask: torch.Tensor,
- input_tensor: torch.Tensor,
- cache_position: torch.Tensor,
- past_seen_tokens: int,
- ):
- # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static
- # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.
- # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using
- # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114
-
- if self.config._attn_implementation == "flash_attention_2":
- if attention_mask is not None and 0.0 in attention_mask:
- return attention_mask
- return None
-
- if self.config._attn_implementation == "sdpa":
- # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument,
- # in order to dispatch on Flash Attention 2.
- if AttentionMaskConverter._ignore_causal_mask_sdpa(
- attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens
- ):
- return None
-
- dtype, device = input_tensor.dtype, input_tensor.device
- min_dtype = torch.finfo(dtype).min
- sequence_length = input_tensor.shape[1]
- if hasattr(getattr(self.layers[0], "self_attn", {}), "past_key_value"): # static cache
- target_length = self.config.max_position_embeddings
- else: # dynamic cache
- target_length = (
- attention_mask.shape[-1]
- if isinstance(attention_mask, torch.Tensor)
- else past_seen_tokens + sequence_length + 1
- )
-
- causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
- if sequence_length != 1:
- causal_mask = torch.triu(causal_mask, diagonal=1)
- causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
- causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
- if attention_mask is not None:
- causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
- if attention_mask.dim() == 2:
- mask_length = attention_mask.shape[-1]
- padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
- causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype)
- elif attention_mask.dim() == 4:
- # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with
- # cache. In that case, the 4D attention mask attends to the newest tokens only.
- if attention_mask.shape[-2] < cache_position[0] + sequence_length:
- offset = cache_position[0]
- else:
- offset = 0
- mask_shape = attention_mask.shape
- mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype
- causal_mask[
- : mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]
- ] = mask_slice
-
- if (
- self.config._attn_implementation == "sdpa"
- and attention_mask is not None
- and attention_mask.device.type == "cuda"
- ):
- # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
- # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
- # Details: https://github.com/pytorch/pytorch/issues/110213
- causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
-
- return causal_mask
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with LLAMA->GEMMA,Llama->Gemma,llama->gemma
-class GemmaForCausalLM(GemmaPreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.model = GemmaModel(config)
- self.vocab_size = config.vocab_size
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.model.embed_tokens
-
- def set_input_embeddings(self, value):
- self.model.embed_tokens = value
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def set_decoder(self, decoder):
- self.model = decoder
-
- def get_decoder(self):
- return self.model
-
- # Ignore copy
- @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- Args:
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, GemmaForCausalLM
-
- >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
- >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
-
- >>> prompt = "What is your favorite condiment?"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "What is your favorite condiment?"
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- )
-
- hidden_states = outputs[0]
- logits = self.lm_head(hidden_states)
- logits = logits.float()
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- loss = loss_fct(shift_logits, shift_labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def prepare_inputs_for_generation(
- self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs
- ):
- # With static cache, the `past_key_values` is None
- # TODO joao: standardize interface for the different Cache classes and remove of this if
- has_static_cache = False
- if past_key_values is None:
- past_key_values = getattr(getattr(self.model.layers[0], "self_attn", {}), "past_key_value", None)
- has_static_cache = past_key_values is not None
-
- past_length = 0
- if past_key_values is not None:
- if isinstance(past_key_values, Cache):
- past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
- max_cache_length = (
- torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
- if past_key_values.get_max_length() is not None
- else None
- )
- cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
- # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
- else:
- cache_length = past_length = past_key_values[0][0].shape[2]
- max_cache_length = None
-
- # Keep only the unprocessed tokens:
- # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
- # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
- # input)
- if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
- input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
- # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
- # input_ids based on the past_length.
- elif past_length < input_ids.shape[1]:
- input_ids = input_ids[:, past_length:]
- # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
-
- # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
- if (
- max_cache_length is not None
- and attention_mask is not None
- and cache_length + input_ids.shape[1] > max_cache_length
- ):
- attention_mask = attention_mask[:, -max_cache_length:]
-
- position_ids = kwargs.get("position_ids", None)
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
- # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
- # TODO: use `next_tokens` directly instead.
- model_inputs = {"input_ids": input_ids.contiguous()}
-
- input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
- if cache_position is None:
- cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
- else:
- cache_position = cache_position[-input_length:]
-
- if has_static_cache:
- past_key_values = None
-
- model_inputs.update(
- {
- "position_ids": position_ids,
- "cache_position": cache_position,
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "attention_mask": attention_mask,
- }
- )
- return model_inputs
-
- @staticmethod
- def _reorder_cache(past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
- )
- return reordered_past
-
-
-@add_start_docstrings(
- """
- The Gemma Model transformer with a sequence classification head on top (linear layer).
-
- [`GemmaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT-2) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- GEMMA_START_DOCSTRING,
-)
-# Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with LLAMA->GEMMA,Llama->Gemma
-class GemmaForSequenceClassification(GemmaPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.model = GemmaModel(config)
- self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.model.embed_tokens
-
- def set_input_embeddings(self, value):
- self.model.embed_tokens = value
-
- @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.model(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size = input_ids.shape[0]
- else:
- batch_size = inputs_embeds.shape[0]
-
- if self.config.pad_token_id is None and batch_size != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
-
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
diff --git a/transformers/models/gemma/tokenization_gemma.py b/transformers/models/gemma/tokenization_gemma.py
deleted file mode 100644
index ab19ee23c794a2531ee5941e5844b362c9a09d04..0000000000000000000000000000000000000000
--- a/transformers/models/gemma/tokenization_gemma.py
+++ /dev/null
@@ -1,326 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tokenization classes for Gemma."""
-import os
-from shutil import copyfile
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
-
-import sentencepiece as spm
-
-from ...tokenization_utils import AddedToken, PreTrainedTokenizer
-from ...utils import logging
-
-
-if TYPE_CHECKING:
- pass
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
-
-SPIECE_UNDERLINE = "▁"
-
-
-class GemmaTokenizer(PreTrainedTokenizer):
- """
- Construct a Gemma tokenizer. Based on byte-level Byte-Pair-Encoding. The default padding token is unset as there is
- no padding token in the original model.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
- eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- The end of sequence token.
- pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by
- attention mechanisms or loss computation.
- sp_model_kwargs (`Dict[str, Any]`, `Optional`, *optional*):
- Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
- SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
- to set:
-
- - `enable_sampling`: Enable subword regularization.
- - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
-
- - `nbest_size = {0,1}`: No sampling is performed.
- - `nbest_size > 1`: samples from the nbest_size results.
- - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
- using forward-filtering-and-backward-sampling algorithm.
-
- - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
- BPE-dropout.
-
- add_bos_token (`bool`, *optional*, defaults to `True`):
- Whether or not to add an `bos_token` at the start of sequences.
- add_eos_token (`bool`, *optional*, defaults to `False`):
- Whether or not to add an `eos_token` at the end of sequences.
- clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
- Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
- extra spaces.
- use_default_system_prompt (`bool`, *optional*, defaults to `False`):
- Whether or not the default system prompt for Gemma should be used.
- spaces_between_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not to add spaces between special tokens.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- unk_token="",
- bos_token="",
- eos_token="",
- pad_token="",
- sp_model_kwargs: Optional[Dict[str, Any]] = None,
- add_bos_token=True,
- add_eos_token=False,
- clean_up_tokenization_spaces=False,
- use_default_system_prompt=False,
- spaces_between_special_tokens=False,
- **kwargs,
- ):
- self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
- bos_token = AddedToken(bos_token, normalized=False, special=True) if isinstance(bos_token, str) else bos_token
- eos_token = AddedToken(eos_token, normalized=False, special=True) if isinstance(eos_token, str) else eos_token
- unk_token = AddedToken(unk_token, normalized=False, special=True) if isinstance(unk_token, str) else unk_token
- pad_token = AddedToken(pad_token, normalized=False, special=True) if isinstance(pad_token, str) else pad_token
-
- self.vocab_file = vocab_file
- self.add_bos_token = add_bos_token
- self.add_eos_token = add_eos_token
- self.use_default_system_prompt = use_default_system_prompt
-
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(vocab_file)
-
- super().__init__(
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- pad_token=pad_token,
- add_bos_token=add_bos_token,
- add_eos_token=add_eos_token,
- sp_model_kwargs=self.sp_model_kwargs,
- clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- use_default_system_prompt=use_default_system_prompt,
- spaces_between_special_tokens=spaces_between_special_tokens,
- **kwargs,
- )
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.__getstate__
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sp_model"] = None
- state["sp_model_proto"] = self.sp_model.serialized_model_proto()
- return state
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.__setstate__
- def __setstate__(self, d):
- self.__dict__ = d
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.LoadFromSerializedProto(self.sp_model_proto)
-
- @property
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.vocab_size
- def vocab_size(self):
- """Returns vocab size"""
- return self.sp_model.get_piece_size()
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.get_vocab
- def get_vocab(self):
- """Returns vocab as a dict"""
- vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
- vocab.update(self.added_tokens_encoder)
- return vocab
-
- def _tokenize(self, text, **kwargs):
- """
- Returns a tokenized string. The Gemma tokenizer never adds a prefix space.
- """
- return self.sp_model.encode(text, out_type=str)
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.sp_model.piece_to_id(token)
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- token = self.sp_model.IdToPiece(index)
- return token
-
- def _decode(
- self,
- token_ids: List[int],
- skip_special_tokens: bool = False,
- spaces_between_special_tokens: bool = False,
- **kwargs,
- ) -> str:
- sub_texts = []
- current_sub_text = []
- for ids in token_ids:
- if skip_special_tokens and ids in self.all_special_ids:
- continue
- if ids in self._added_tokens_decoder:
- if current_sub_text:
- sub_texts.append(self.sp_model.decode(current_sub_text))
- sub_texts.append(self._added_tokens_decoder[ids].content)
- current_sub_text = []
- else:
- current_sub_text.append(ids)
- if current_sub_text:
- sub_texts.append(self.sp_model.decode(current_sub_text))
-
- if spaces_between_special_tokens:
- sub_texts = " ".join(sub_texts)
- else:
- sub_texts = "".join(sub_texts)
-
- return sub_texts
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- current_sub_tokens = []
- out_string = ""
- for token in tokens:
- # make sure that special tokens are not decoded using sentencepiece model
- if token in self._added_tokens_encoder:
- out_string += self.sp_model.decode(current_sub_tokens) + token
- current_sub_tokens = []
- else:
- current_sub_tokens.append(token)
- out_string += self.sp_model.decode(current_sub_tokens)
- return out_string
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
- """
- Save the vocabulary and special tokens file to a directory.
-
- Args:
- save_directory (`str`):
- The directory in which to save the vocabulary.
-
- Returns:
- `Tuple(str)`: Paths to the files saved.
- """
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
- elif not os.path.isfile(self.vocab_file):
- with open(out_vocab_file, "wb") as fi:
- content_spiece_model = self.sp_model.serialized_model_proto()
- fi.write(content_spiece_model)
-
- return (out_vocab_file,)
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- bos_token_id = [self.bos_token_id] if self.add_bos_token else []
- eos_token_id = [self.eos_token_id] if self.add_eos_token else []
-
- output = bos_token_id + token_ids_0 + eos_token_id
-
- if token_ids_1 is not None:
- output = output + bos_token_id + token_ids_1 + eos_token_id
-
- return output
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.get_special_tokens_mask
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- bos_token_id = [1] if self.add_bos_token else []
- eos_token_id = [1] if self.add_eos_token else []
-
- if token_ids_1 is None:
- return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
- return (
- bos_token_id
- + ([0] * len(token_ids_0))
- + eos_token_id
- + bos_token_id
- + ([0] * len(token_ids_1))
- + eos_token_id
- )
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.create_token_type_ids_from_sequences
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- """
- Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
- sequence pair mask has the following format:
-
- ```
- 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
- | first sequence | second sequence |
- ```
-
- if token_ids_1 is None, only returns the first portion of the mask (0s).
-
- Args:
- token_ids_0 (`List[int]`):
- List of ids.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
-
- Returns:
- `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
- """
- bos_token_id = [self.bos_token_id] if self.add_bos_token else []
- eos_token_id = [self.eos_token_id] if self.add_eos_token else []
-
- output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
-
- if token_ids_1 is not None:
- output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
-
- return output
diff --git a/transformers/models/gemma/tokenization_gemma_fast.py b/transformers/models/gemma/tokenization_gemma_fast.py
deleted file mode 100644
index fd7a979e8b7509cd03de5fe12879b3a0b5a49dfa..0000000000000000000000000000000000000000
--- a/transformers/models/gemma/tokenization_gemma_fast.py
+++ /dev/null
@@ -1,199 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import os
-from shutil import copyfile
-from typing import Optional, Tuple
-
-from tokenizers import processors
-
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import is_sentencepiece_available, logging
-from ...utils.versions import require_version
-
-
-require_version("tokenizers>=0.13.3")
-
-if is_sentencepiece_available():
- from .tokenization_gemma import GemmaTokenizer
-else:
- GemmaTokenizer = None
-
-logger = logging.get_logger(__name__)
-VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model", "tokenizer_file": "tokenizer.json"}
-
-
-class GemmaTokenizerFast(PreTrainedTokenizerFast):
- """
- Construct a Gemma tokenizer fast. Based on byte-level Byte-Pair-Encoding.
-
- This uses notably ByteFallback and no prefix space. Normalization is applied to replace `" "` with `"▁"`
-
- ```python
- >>> from transformers import GemmaTokenizerFast
-
- >>> tokenizer = GemmaTokenizerFast.from_pretrained("hf-internal-testing/dummy-gemma")
- >>> tokenizer.encode("Hello this is a test")
- [2, 4521, 736, 603, 476, 2121]
- ```
-
- If you want to change the `bos_token` or the `eos_token`, make sure to specify them when initializing the model, or
- call `tokenizer.update_post_processor()` to make sure that the post-processing is correctly done (otherwise the
- values of the first token and final token of an encoded sequence will not be correct). For more details, checkout
- [post-processors] (https://huggingface.co/docs/tokenizers/api/post-processors) documentation.
-
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`, *optional*):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .model extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- tokenizer_file (`str`, *optional*):
- [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that
- contains everything needed to load the tokenizer.
- clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
- Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
- extra spaces.
- unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
- eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`):
- The end of sequence token.
- pad_token (`str`, *optional*, defaults to `""`):
- The padding token
- add_bos_token (`bool`, *optional*, defaults to `True`):
- Whether or not to add an `bos_token` at the start of sequences.
- add_eos_token (`bool`, *optional*, defaults to `False`):
- Whether or not to add an `eos_token` at the end of sequences.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- slow_tokenizer_class = GemmaTokenizer
- padding_side = "left"
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file=None,
- tokenizer_file=None,
- clean_up_tokenization_spaces=False,
- unk_token="",
- bos_token="",
- eos_token="",
- pad_token="",
- add_bos_token=True,
- add_eos_token=False,
- **kwargs,
- ):
- super().__init__(
- vocab_file=vocab_file,
- tokenizer_file=tokenizer_file,
- clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- unk_token=unk_token,
- bos_token=bos_token,
- eos_token=eos_token,
- pad_token=pad_token,
- add_bos_token=add_bos_token,
- add_eos_token=add_eos_token,
- **kwargs,
- )
- self._add_bos_token = add_bos_token
- self._add_eos_token = add_eos_token
- self.update_post_processor()
- self.vocab_file = vocab_file
-
- @property
- def can_save_slow_tokenizer(self) -> bool:
- return os.path.isfile(self.vocab_file) if self.vocab_file else False
-
- # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.update_post_processor
- def update_post_processor(self):
- """
- Updates the underlying post processor with the current `bos_token` and `eos_token`.
- """
- bos = self.bos_token
- bos_token_id = self.bos_token_id
- if bos is None and self.add_bos_token:
- raise ValueError("add_bos_token = True but bos_token = None")
-
- eos = self.eos_token
- eos_token_id = self.eos_token_id
- if eos is None and self.add_eos_token:
- raise ValueError("add_eos_token = True but eos_token = None")
-
- single = f"{(bos+':0 ') if self.add_bos_token else ''}$A:0{(' '+eos+':0') if self.add_eos_token else ''}"
- pair = f"{single}{(' '+bos+':1') if self.add_bos_token else ''} $B:1{(' '+eos+':1') if self.add_eos_token else ''}"
-
- special_tokens = []
- if self.add_bos_token:
- special_tokens.append((bos, bos_token_id))
- if self.add_eos_token:
- special_tokens.append((eos, eos_token_id))
- self._tokenizer.post_processor = processors.TemplateProcessing(
- single=single, pair=pair, special_tokens=special_tokens
- )
-
- @property
- def add_eos_token(self):
- return self._add_eos_token
-
- @property
- def add_bos_token(self):
- return self._add_bos_token
-
- @add_eos_token.setter
- def add_eos_token(self, value):
- self._add_eos_token = value
- self.update_post_processor()
-
- @add_bos_token.setter
- def add_bos_token(self, value):
- self._add_bos_token = value
- self.update_post_processor()
-
- # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not self.can_save_slow_tokenizer:
- raise ValueError(
- "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
- "tokenizer."
- )
-
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
-
- return (out_vocab_file,)
-
- # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- bos_token_id = [self.bos_token_id] if self.add_bos_token else []
- eos_token_id = [self.eos_token_id] if self.add_eos_token else []
-
- output = bos_token_id + token_ids_0 + eos_token_id
-
- if token_ids_1 is not None:
- output = output + bos_token_id + token_ids_1 + eos_token_id
-
- return output
diff --git a/transformers/models/git/__init__.py b/transformers/models/git/__init__.py
deleted file mode 100644
index e234a4b01db188e83c4e21ba000d24f60b13b286..0000000000000000000000000000000000000000
--- a/transformers/models/git/__init__.py
+++ /dev/null
@@ -1,60 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
-
-
-_import_structure = {
- "configuration_git": ["GIT_PRETRAINED_CONFIG_ARCHIVE_MAP", "GitConfig", "GitVisionConfig"],
- "processing_git": ["GitProcessor"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_git"] = [
- "GIT_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GitForCausalLM",
- "GitModel",
- "GitPreTrainedModel",
- "GitVisionModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_git import GIT_PRETRAINED_CONFIG_ARCHIVE_MAP, GitConfig, GitVisionConfig
- from .processing_git import GitProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_git import (
- GIT_PRETRAINED_MODEL_ARCHIVE_LIST,
- GitForCausalLM,
- GitModel,
- GitPreTrainedModel,
- GitVisionModel,
- )
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/git/__pycache__/__init__.cpython-310.pyc b/transformers/models/git/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 4fd3ad381baf6db88686c10cea4afce26ed1fdad..0000000000000000000000000000000000000000
Binary files a/transformers/models/git/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/git/__pycache__/configuration_git.cpython-310.pyc b/transformers/models/git/__pycache__/configuration_git.cpython-310.pyc
deleted file mode 100644
index 6dbd79fcdacca1f212a7a81609934d78d934d9ca..0000000000000000000000000000000000000000
Binary files a/transformers/models/git/__pycache__/configuration_git.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/git/__pycache__/convert_git_to_pytorch.cpython-310.pyc b/transformers/models/git/__pycache__/convert_git_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 88b88a325202c92684fce0eaceef44ffa4fd7a79..0000000000000000000000000000000000000000
Binary files a/transformers/models/git/__pycache__/convert_git_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/git/__pycache__/modeling_git.cpython-310.pyc b/transformers/models/git/__pycache__/modeling_git.cpython-310.pyc
deleted file mode 100644
index 24a3db8786fa739f9ada296dc51aeb1bca367322..0000000000000000000000000000000000000000
Binary files a/transformers/models/git/__pycache__/modeling_git.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/git/__pycache__/processing_git.cpython-310.pyc b/transformers/models/git/__pycache__/processing_git.cpython-310.pyc
deleted file mode 100644
index 0e369d206f892395ec63cbf5a2439f0ec5fac84b..0000000000000000000000000000000000000000
Binary files a/transformers/models/git/__pycache__/processing_git.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/git/configuration_git.py b/transformers/models/git/configuration_git.py
deleted file mode 100644
index 0c28bbabff6b0b15d27644599fc9692e19aa1d9e..0000000000000000000000000000000000000000
--- a/transformers/models/git/configuration_git.py
+++ /dev/null
@@ -1,240 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-from typing import Union
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GIT_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GitVisionConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GitVisionModel`]. It is used to instantiate a GIT
- vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the vision encoder of the GIT
- [microsoft/git-base](https://huggingface.co/microsoft/git-base) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- num_hidden_layers (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- image_size (`int`, *optional*, defaults to 224):
- The size (resolution) of each image.
- patch_size (`int`, *optional*, defaults to 16):
- The size (resolution) of each patch.
- hidden_act (`str` or `function`, *optional*, defaults to `"quick_gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` ``"quick_gelu"` are supported.
- layer_norm_eps (`float`, *optional*, defaults to 1e-5):
- The epsilon used by the layer normalization layers.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-
- Example:
-
- ```python
- >>> from transformers import GitVisionConfig, GitVisionModel
-
- >>> # Initializing a GitVisionConfig with microsoft/git-base style configuration
- >>> configuration = GitVisionConfig()
-
- >>> # Initializing a GitVisionModel (with random weights) from the microsoft/git-base style configuration
- >>> model = GitVisionModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "git_vision_model"
-
- def __init__(
- self,
- hidden_size=768,
- intermediate_size=3072,
- num_hidden_layers=12,
- num_attention_heads=12,
- num_channels=3,
- image_size=224,
- patch_size=16,
- hidden_act="quick_gelu",
- layer_norm_eps=1e-5,
- attention_dropout=0.0,
- initializer_range=0.02,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.hidden_size = hidden_size
- self.intermediate_size = intermediate_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.num_channels = num_channels
- self.patch_size = patch_size
- self.image_size = image_size
- self.initializer_range = initializer_range
- self.attention_dropout = attention_dropout
- self.layer_norm_eps = layer_norm_eps
- self.hidden_act = hidden_act
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
- cls._set_token_in_kwargs(kwargs)
-
- config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
-
- # get the vision config dict if we are loading from GITConfig
- if config_dict.get("model_type") == "git":
- config_dict = config_dict["vision_config"]
-
- if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
- logger.warning(
- f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
- f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
- )
-
- return cls.from_dict(config_dict, **kwargs)
-
-
-class GitConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GitModel`]. It is used to instantiate a GIT model
- according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the GIT
- [microsoft/git-base](https://huggingface.co/microsoft/git-base) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- vision_config (`dict`, *optional*):
- Dictionary of configuration options used to initialize [`GitVisionConfig`].
- vocab_size (`int`, *optional*, defaults to 30522):
- Vocabulary size of the GIT model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`GitModel`].
- hidden_size (`int`, *optional*, defaults to 768):
- Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 6):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 3072):
- Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"silu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention probabilities.
- max_position_embeddings (`int`, *optional*, defaults to 1024):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
- Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
- positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
- [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
- For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
- with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
- num_image_with_embedding (`int`, *optional*):
- The number of temporal embeddings to add, in case the model is used for video captioning/VQA.
-
- Examples:
-
- ```python
- >>> from transformers import GitConfig, GitModel
-
- >>> # Initializing a GIT microsoft/git-base style configuration
- >>> configuration = GitConfig()
-
- >>> # Initializing a model (with random weights) from the microsoft/git-base style configuration
- >>> model = GitModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "git"
-
- def __init__(
- self,
- vision_config=None,
- vocab_size=30522,
- hidden_size=768,
- num_hidden_layers=6,
- num_attention_heads=12,
- intermediate_size=3072,
- hidden_act="gelu",
- hidden_dropout_prob=0.1,
- attention_probs_dropout_prob=0.1,
- max_position_embeddings=1024,
- initializer_range=0.02,
- layer_norm_eps=1e-12,
- pad_token_id=0,
- position_embedding_type="absolute",
- use_cache=True,
- tie_word_embeddings=False,
- bos_token_id=101,
- eos_token_id=102,
- num_image_with_embedding=None,
- **kwargs,
- ):
- super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, **kwargs)
-
- if vision_config is None:
- vision_config = {}
- logger.info("vision_config is None. initializing the GitVisionConfig with default values.")
-
- self.vision_config = GitVisionConfig(**vision_config)
- self.vocab_size = vocab_size
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.hidden_act = hidden_act
- self.intermediate_size = intermediate_size
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.max_position_embeddings = max_position_embeddings
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.position_embedding_type = position_embedding_type
- self.use_cache = use_cache
- self.tie_word_embeddings = tie_word_embeddings
- self.num_image_with_embedding = num_image_with_embedding
-
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
diff --git a/transformers/models/git/convert_git_to_pytorch.py b/transformers/models/git/convert_git_to_pytorch.py
deleted file mode 100644
index 4e3e8e7b317905f9108acdd9e3b5b97c1445b24b..0000000000000000000000000000000000000000
--- a/transformers/models/git/convert_git_to_pytorch.py
+++ /dev/null
@@ -1,428 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert GIT checkpoints from the original repository.
-
-URL: https://github.com/microsoft/GenerativeImage2Text/tree/main"""
-
-
-import argparse
-from pathlib import Path
-
-import numpy as np
-import requests
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
-
-from transformers import (
- AutoTokenizer,
- CLIPImageProcessor,
- GitConfig,
- GitForCausalLM,
- GitProcessor,
- GitVisionConfig,
- VideoMAEImageProcessor,
-)
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def get_git_config(model_name):
- if "base" in model_name and "vqa" in model_name:
- image_size = 480
- elif "large" in model_name and "vqa" in model_name:
- image_size = 420
- else:
- image_size = 224
-
- vision_config = GitVisionConfig(image_size=image_size)
-
- if "large" in model_name:
- vision_config.patch_size = 14
- vision_config.hidden_size = 1024
- vision_config.intermediate_size = 4096
- vision_config.num_hidden_layers = 24
- vision_config.num_attention_heads = 16
-
- is_video = "vatex" in model_name or "msrvtt" in model_name
- num_image_with_embedding = 6 if is_video else None
- config = GitConfig(vision_config=vision_config.to_dict(), num_image_with_embedding=num_image_with_embedding)
-
- return config, image_size, is_video
-
-
-# here we list all keys to be renamed (original name on the left, our name on the right)
-def create_rename_keys(config, prefix=""):
- rename_keys = []
-
- # image encoder
- # ftm: off
- rename_keys.append(
- (f"{prefix}image_encoder.class_embedding", "git.image_encoder.vision_model.embeddings.class_embedding")
- )
- rename_keys.append(
- (
- f"{prefix}image_encoder.positional_embedding",
- "git.image_encoder.vision_model.embeddings.position_embedding.weight",
- )
- )
- rename_keys.append(
- (f"{prefix}image_encoder.conv1.weight", "git.image_encoder.vision_model.embeddings.patch_embedding.weight")
- )
- rename_keys.append((f"{prefix}image_encoder.ln_pre.weight", "git.image_encoder.vision_model.pre_layrnorm.weight"))
- rename_keys.append((f"{prefix}image_encoder.ln_pre.bias", "git.image_encoder.vision_model.pre_layrnorm.bias"))
- rename_keys.append(
- (f"{prefix}image_encoder.ln_post.weight", "git.image_encoder.vision_model.post_layernorm.weight")
- )
- rename_keys.append((f"{prefix}image_encoder.ln_post.bias", "git.image_encoder.vision_model.post_layernorm.bias"))
- # fmt: on
- rename_keys.append((f"{prefix}image_encoder.proj", "git.image_encoder.visual_projection.weight"))
-
- # fmt: off
- for i in range(config.vision_config.num_hidden_layers):
- # image encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.attn.out_proj.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.out_proj.weight"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.attn.out_proj.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.out_proj.bias"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.ln_1.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.layer_norm1.weight"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.ln_1.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.layer_norm1.bias"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.mlp.c_fc.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.mlp.fc1.weight"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.mlp.c_fc.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.mlp.fc1.bias"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.mlp.c_proj.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.mlp.fc2.weight"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.mlp.c_proj.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.mlp.fc2.bias"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.ln_2.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.layer_norm2.weight"))
- rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.ln_2.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.layer_norm2.bias"))
- # fmt: on
-
- # text decoder
- # fmt: off
- rename_keys.append((f"{prefix}textual.embedding.words.weight", "git.embeddings.word_embeddings.weight"))
- rename_keys.append((f"{prefix}textual.embedding.positions.weight", "git.embeddings.position_embeddings.weight"))
- rename_keys.append((f"{prefix}textual.visual_projection.0.weight", "git.visual_projection.visual_projection.0.weight"))
- rename_keys.append((f"{prefix}textual.visual_projection.0.bias", "git.visual_projection.visual_projection.0.bias"))
- rename_keys.append((f"{prefix}textual.visual_projection.1.weight", "git.visual_projection.visual_projection.1.weight"))
- rename_keys.append((f"{prefix}textual.visual_projection.1.bias", "git.visual_projection.visual_projection.1.bias"))
-
- rename_keys.append((f"{prefix}textual.embedding.layer_norm.weight", "git.embeddings.LayerNorm.weight"))
- rename_keys.append((f"{prefix}textual.embedding.layer_norm.bias", "git.embeddings.LayerNorm.bias"))
- rename_keys.append((f"{prefix}textual.output.weight", "output.weight"))
- rename_keys.append((f"{prefix}textual.output.bias", "output.bias"))
- for i in range(config.num_hidden_layers):
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.query.weight", f"git.encoder.layer.{i}.attention.self.query.weight"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.query.bias", f"git.encoder.layer.{i}.attention.self.query.bias"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.key.weight", f"git.encoder.layer.{i}.attention.self.key.weight"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.key.bias", f"git.encoder.layer.{i}.attention.self.key.bias"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.value.weight", f"git.encoder.layer.{i}.attention.self.value.weight"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.value.bias", f"git.encoder.layer.{i}.attention.self.value.bias"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.output.dense.weight", f"git.encoder.layer.{i}.attention.output.dense.weight"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.output.dense.bias", f"git.encoder.layer.{i}.attention.output.dense.bias"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.output.LayerNorm.weight", f"git.encoder.layer.{i}.attention.output.LayerNorm.weight"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.output.LayerNorm.bias", f"git.encoder.layer.{i}.attention.output.LayerNorm.bias"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.intermediate.dense.weight", f"git.encoder.layer.{i}.intermediate.dense.weight"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.intermediate.dense.bias", f"git.encoder.layer.{i}.intermediate.dense.bias"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.output.dense.weight", f"git.encoder.layer.{i}.output.dense.weight"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.output.dense.bias", f"git.encoder.layer.{i}.output.dense.bias"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.output.LayerNorm.weight", f"git.encoder.layer.{i}.output.LayerNorm.weight"))
- rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.output.LayerNorm.bias", f"git.encoder.layer.{i}.output.LayerNorm.bias"))
- # fmt: on
-
- if config.num_image_with_embedding is not None:
- rename_keys.append(("img_temperal_embedding.0", "git.img_temperal_embedding.0"))
- rename_keys.append(("img_temperal_embedding.1", "git.img_temperal_embedding.1"))
- rename_keys.append(("img_temperal_embedding.2", "git.img_temperal_embedding.2"))
- rename_keys.append(("img_temperal_embedding.3", "git.img_temperal_embedding.3"))
- rename_keys.append(("img_temperal_embedding.4", "git.img_temperal_embedding.4"))
- rename_keys.append(("img_temperal_embedding.5", "git.img_temperal_embedding.5"))
-
- return rename_keys
-
-
-def rename_key(dct, old, new):
- val = dct.pop(old)
- dct[new] = val.T if "image_encoder.visual_projection" in new else val
-
-
-# we split up the matrix of each CLIP encoder layer into queries, keys and values
-def read_in_q_k_v(state_dict, config, prefix=""):
- dim = config.vision_config.hidden_size
- for i in range(config.vision_config.num_hidden_layers):
- # read in weights + bias of input projection layer (in the original implementation, this is a single matrix + bias)
- in_proj_weight = state_dict.pop(f"{prefix}image_encoder.transformer.resblocks.{i}.attn.in_proj_weight")
- in_proj_bias = state_dict.pop(f"{prefix}image_encoder.transformer.resblocks.{i}.attn.in_proj_bias")
- # next, add query, keys and values (in that order) to the state dict
- state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[
- :dim, :
- ]
- state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:dim]
- state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[
- dim : dim * 2, :
- ]
- state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[
- dim : dim * 2
- ]
- state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[
- -dim:, :
- ]
- state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-dim:]
-
-
-# We will verify our results on an image
-def prepare_img(model_name):
- if "textvqa" in model_name:
- filepath = hf_hub_download(repo_id="nielsr/textvqa-sample", filename="bus.png", repo_type="dataset")
- image = Image.open(filepath).convert("RGB")
- else:
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- image = Image.open(requests.get(url, stream=True).raw)
-
- return image
-
-
-def prepare_video():
- from decord import VideoReader, cpu
-
- # set seed for reproducability
- np.random.seed(0)
-
- def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
- """
- Sample a given number of frame indices from the video.
-
- Args:
- clip_len (`int`): Total number of frames to sample.
- frame_sample_rate (`int`): Sample every n-th frame.
- seg_len (`int`): Maximum allowed index of sample's last frame.
-
- Returns:
- indices (`List[int]`): List of sampled frame indices
- """
- converted_len = int(clip_len * frame_sample_rate)
- end_idx = np.random.randint(converted_len, seg_len)
- start_idx = end_idx - converted_len
- indices = np.linspace(start_idx, end_idx, num=clip_len)
- indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
- return indices
-
- # video clip consists of 300 frames (10 seconds at 30 FPS)
- file_path = hf_hub_download(repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset")
- videoreader = VideoReader(file_path, num_threads=1, ctx=cpu(0))
-
- # sample 6 frames
- videoreader.seek(0)
- indices = sample_frame_indices(clip_len=6, frame_sample_rate=4, seg_len=len(videoreader))
- video = videoreader.get_batch(indices).asnumpy()
-
- return video
-
-
-@torch.no_grad()
-def convert_git_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub=False):
- """
- Copy/paste/tweak model's weights to our GIT structure.
- """
-
- model_name_to_url = {
- "git-base": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE/snapshot/model.pt",
- "git-base-coco": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_COCO/snapshot/model.pt",
- "git-base-textcaps": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_TEXTCAPS/snapshot/model.pt",
- "git-base-vqav2": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_VQAv2/snapshot/model.pt",
- "git-base-textvqa": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_TEXTVQA/snapshot/model.pt", # todo
- "git-base-vatex": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_VATEX/snapshot/model.pt",
- "git-base-msrvtt-qa": (
- "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_MSRVTT_QA/snapshot/model.pt"
- ),
- "git-large": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE/snapshot/model.pt",
- "git-large-coco": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_COCO/snapshot/model.pt",
- "git-large-textcaps": (
- "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_TEXTCAPS/snapshot/model.pt"
- ),
- "git-large-vqav2": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_VQAv2/snapshot/model.pt",
- "git-large-textvqa": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_TEXTVQA/snapshot/model.pt",
- "git-large-vatex": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_VATEX/snapshot/model.pt",
- "git-large-msrvtt-qa": (
- "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_MSRVTT_QA/snapshot/model.pt"
- ),
- "git-large-r": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_R/snapshot/model.pt",
- "git-large-r-coco": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_R_COCO/snapshot/model.pt",
- "git-large-r-textcaps": (
- "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_R_TEXTCAPS/snapshot/model.pt"
- ),
- }
-
- model_name_to_path = {
- "git-large": "/Users/nielsrogge/Documents/GIT/git_large_model.pt",
- "git-large-coco": "/Users/nielsrogge/Documents/GIT/git_large_coco_model.pt",
- "git-large-textcaps": "/Users/nielsrogge/Documents/GIT/git_large_textcaps_model.pt",
- "git-large-vqav2": "/Users/nielsrogge/Documents/GIT/git_large_vqav2_model.pt",
- "git-large-textvqa": "/Users/nielsrogge/Documents/GIT/git_large_textvqa_model.pt",
- }
-
- # define GIT configuration based on model name
- config, image_size, is_video = get_git_config(model_name)
- if "large" in model_name and not is_video and "large-r" not in model_name:
- # large checkpoints take way too long to download
- checkpoint_path = model_name_to_path[model_name]
- state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
- else:
- checkpoint_url = model_name_to_url[model_name]
- state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu", file_name=model_name)[
- "model"
- ]
- # rename keys
- prefix = "module." if model_name == "git-base" else ""
- rename_keys = create_rename_keys(config, prefix=prefix)
- for src, dest in rename_keys:
- rename_key(state_dict, src, dest)
- read_in_q_k_v(state_dict, config, prefix=prefix)
-
- # load HuggingFace model
- model = GitForCausalLM(config)
- missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
- model.eval()
-
- print("Missing keys:", missing_keys)
- print("Unexpected keys:", unexpected_keys)
-
- assert missing_keys == ["git.embeddings.position_ids", "git.image_encoder.vision_model.embeddings.position_ids"]
- assert unexpected_keys == ["git.image_encoder.visual_projection.weight"]
-
- # verify results
- image_processor = (
- VideoMAEImageProcessor(
- size={"shortest_edge": image_size}, crop_size={"height": image_size, "width": image_size}
- )
- if is_video
- else CLIPImageProcessor(
- size={"shortest_edge": image_size}, crop_size={"height": image_size, "width": image_size}
- )
- )
- tokenizer = AutoTokenizer.from_pretrained(
- "google-bert/bert-base-uncased", model_input_names=["input_ids", "attention_mask"]
- )
- processor = GitProcessor(tokenizer=tokenizer, image_processor=image_processor)
-
- if is_video:
- video = prepare_video()
- pixel_values = processor(images=list(video), return_tensors="pt").pixel_values
- else:
- image = prepare_img(model_name)
- image_transforms = Compose(
- [
- Resize(image_size, interpolation=Image.BICUBIC),
- CenterCrop(image_size),
- ToTensor(),
- Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
- ]
- )
- original_pixel_values = image_transforms(image).unsqueeze(0)
- pixel_values = processor(images=image, return_tensors="pt").pixel_values
-
- assert torch.allclose(pixel_values, original_pixel_values)
-
- input_ids = torch.tensor([[101]])
- outputs = model(input_ids, pixel_values=pixel_values)
- logits = outputs.logits
- print("Logits:", logits[0, -1, :3])
-
- if model_name == "git-base":
- expected_slice_logits = torch.tensor([-1.2832, -1.2835, -1.2840])
- elif model_name == "git-base-coco":
- expected_slice_logits = torch.tensor([-0.9925, -0.9930, -0.9935])
- elif model_name == "git-base-textcaps":
- expected_slice_logits = torch.tensor([-1.2980, -1.2983, -1.2985])
- elif model_name == "git-base-vqav2":
- expected_slice_logits = torch.tensor([-0.8570, -0.8568, -0.8561])
- elif model_name == "git-base-textvqa":
- expected_slice_logits = torch.tensor([-1.4085, -1.4083, -1.4082])
- elif model_name == "git-base-vatex":
- expected_slice_logits = torch.tensor([-1.3451, -1.3447, -1.3447])
- elif model_name == "git-base-msrvtt-qa":
- expected_slice_logits = torch.tensor([-0.8554, -0.8550, -0.8540])
- elif model_name == "git-large":
- expected_slice_logits = torch.tensor([-1.1708, -1.1707, -1.1705])
- elif model_name == "git-large-coco":
- expected_slice_logits = torch.tensor([-1.0425, -1.0423, -1.0422])
- elif model_name == "git-large-textcaps":
- expected_slice_logits = torch.tensor([-1.2705, -1.2708, -1.2706])
- elif model_name == "git-large-vqav2":
- expected_slice_logits = torch.tensor([-0.7042, -0.7043, -0.7043])
- elif model_name == "git-large-textvqa":
- expected_slice_logits = torch.tensor([-0.8590, -0.8592, -0.8590])
- elif model_name == "git-large-vatex":
- expected_slice_logits = torch.tensor([-1.0113, -1.0114, -1.0113])
- elif model_name == "git-large-msrvtt-qa":
- expected_slice_logits = torch.tensor([0.0130, 0.0134, 0.0131])
- elif model_name == "git-large-r":
- expected_slice_logits = torch.tensor([-1.1283, -1.1285, -1.1286])
- elif model_name == "git-large-r-coco":
- expected_slice_logits = torch.tensor([-0.9641, -0.9641, -0.9641])
- elif model_name == "git-large-r-textcaps":
- expected_slice_logits = torch.tensor([-1.1121, -1.1120, -1.1124])
-
- assert torch.allclose(logits[0, -1, :3], expected_slice_logits, atol=1e-4)
- print("Looks ok!")
-
- prompt = ""
- if "textvqa" in model_name:
- prompt = "what does the front of the bus say at the top?"
- elif "msrvtt-qa" in model_name:
- prompt = "what does the woman eat?"
- elif "vqa" in model_name:
- prompt = "what are the cats doing?"
- input_ids = tokenizer(prompt, add_special_tokens=False).input_ids
- input_ids = [processor.tokenizer.cls_token_id] + input_ids
- input_ids = torch.tensor(input_ids).unsqueeze(0)
- print("Generating caption...")
- generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50)
- print("Generated caption:", processor.batch_decode(generated_ids, skip_special_tokens=True))
-
- if pytorch_dump_folder_path is not None:
- Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
- print(f"Saving model and processor of {model_name} to {pytorch_dump_folder_path}")
- model.save_pretrained(pytorch_dump_folder_path)
- processor.save_pretrained(pytorch_dump_folder_path)
-
- if push_to_hub:
- print(f"Pushing model and processor of {model_name} to the hub...")
- model.push_to_hub(f"microsoft/{model_name}")
- processor.push_to_hub(f"microsoft/{model_name}")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--model_name",
- default="git-base",
- type=str,
- help="Name of the model you'd like to convert.",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path",
- default=None,
- type=str,
- help="Path to the output PyTorch model directory.",
- )
- parser.add_argument(
- "--push_to_hub",
- action="store_true",
- help="Whether to push the model to the hub.",
- )
-
- args = parser.parse_args()
- convert_git_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/transformers/models/git/modeling_git.py b/transformers/models/git/modeling_git.py
deleted file mode 100644
index c8953d498428eac44d3b4f2d26673557ab48c30a..0000000000000000000000000000000000000000
--- a/transformers/models/git/modeling_git.py
+++ /dev/null
@@ -1,1543 +0,0 @@
-# coding=utf-8
-# Copyright 2022 Microsoft Research and The HuggingFace Inc. team.
-# All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch GIT model."""
-
-
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import CrossEntropyLoss
-
-from ...activations import ACT2FN
-from ...file_utils import ModelOutput
-from ...modeling_attn_mask_utils import _prepare_4d_attention_mask
-from ...modeling_outputs import (
- BaseModelOutput,
- BaseModelOutputWithPast,
- BaseModelOutputWithPooling,
- CausalLMOutputWithPast,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from .configuration_git import GitConfig, GitVisionConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "microsoft/git-base"
-_CONFIG_FOR_DOC = "GitConfig"
-
-
-from ..deprecated._archive_maps import GIT_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-@dataclass
-# Copied from transformers.models.clip.modeling_clip.CLIPVisionModelOutput with CLIP->Git
-class GitVisionModelOutput(ModelOutput):
- """
- Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.
-
- Args:
- image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
- The image embeddings obtained by applying the projection layer to the pooler_output.
- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Sequence of hidden-states at the output of the last layer of the model.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
- one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- image_embeds: Optional[torch.FloatTensor] = None
- last_hidden_state: torch.FloatTensor = None
- hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
- attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-
-
-class GitEmbeddings(nn.Module):
- """Construct the embeddings from word and position embeddings."""
-
- def __init__(self, config):
- super().__init__()
- self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
-
- # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
- # any TensorFlow checkpoint file
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
- # position_ids (1, len position emb) is contiguous in memory and exported when serialized
- self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
- self.register_buffer(
- "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
- )
-
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- past_key_values_length: int = 0,
- ) -> torch.Tensor:
- if input_ids is not None:
- input_shape = input_ids.size()
- else:
- input_shape = inputs_embeds.size()[:-1]
-
- seq_length = input_shape[1]
-
- if position_ids is None:
- position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
-
- if inputs_embeds is None:
- embeddings = self.word_embeddings(input_ids)
- else:
- embeddings = inputs_embeds
-
- if self.position_embedding_type == "absolute":
- position_embeddings = self.position_embeddings(position_ids)
- embeddings += position_embeddings
- embeddings = self.LayerNorm(embeddings)
- embeddings = self.dropout(embeddings)
- return embeddings
-
-
-class GitSelfAttention(nn.Module):
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
- raise ValueError(
- f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
- f"heads ({config.num_attention_heads})"
- )
-
- self.num_attention_heads = config.num_attention_heads
- self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
- self.image_patch_tokens = int((config.vision_config.image_size / config.vision_config.patch_size) ** 2 + 1)
- if config.num_image_with_embedding is not None:
- self.image_patch_tokens *= config.num_image_with_embedding
-
- self.query = nn.Linear(config.hidden_size, self.all_head_size)
- self.key = nn.Linear(config.hidden_size, self.all_head_size)
- self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
- self.position_embedding_type = position_embedding_type or getattr(
- config, "position_embedding_type", "absolute"
- )
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- self.max_position_embeddings = config.max_position_embeddings
- self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
-
- def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
- new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- x = x.view(new_x_shape)
- return x.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- pixel_values_present: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- mixed_query_layer = self.query(hidden_states)
-
- cutoff = self.image_patch_tokens if pixel_values_present else 0
- if past_key_value is not None:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
- key_layer = torch.cat([key_layer[:, :, :cutoff, :], past_key_value[0], key_layer[:, :, -1:, :]], dim=2)
- value_layer = torch.cat(
- [value_layer[:, :, :cutoff, :], past_key_value[1], value_layer[:, :, -1:, :]], dim=2
- )
- else:
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- query_layer = self.transpose_for_scores(mixed_query_layer)
-
- use_cache = past_key_value is not None
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- # NOTE: like in other caches, we store the text component. In GIT it means we discard the image component.
- past_key_value = (
- key_layer[:, :, cutoff:, :],
- value_layer[:, :, cutoff:, :],
- )
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
- query_length, key_length = query_layer.shape[2], key_layer.shape[2]
- if use_cache:
- position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(
- -1, 1
- )
- else:
- position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
- position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
- distance = position_ids_l - position_ids_r
-
- positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
- positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
-
- if self.position_embedding_type == "relative_key":
- relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores
- elif self.position_embedding_type == "relative_key_query":
- relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
- relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
- attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
- if attention_mask is not None:
- # Apply the attention mask is (precomputed for all layers in GitModel forward() function)
- attention_scores = attention_scores + attention_mask
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- # Mask heads if we want to
- if head_mask is not None:
- attention_probs = attention_probs * head_mask
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- outputs = outputs + (past_key_value,)
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
-class GitSelfOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class GitAttention(nn.Module):
- # Copied from transformers.models.bert.modeling_bert.BertAttention.__init__ with Bert->Git
- def __init__(self, config, position_embedding_type=None):
- super().__init__()
- self.self = GitSelfAttention(config, position_embedding_type=position_embedding_type)
- self.output = GitSelfOutput(config)
- self.pruned_heads = set()
-
- # Copied from transformers.models.bert.modeling_bert.BertAttention.prune_heads
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- pixel_values_present: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- self_outputs = self.self(
- hidden_states,
- attention_mask,
- head_mask,
- past_key_value,
- output_attentions,
- pixel_values_present,
- )
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.bert.modeling_bert.BertIntermediate
-class GitIntermediate(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.intermediate_act_fn(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.bert.modeling_bert.BertOutput
-class GitOutput(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
- self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.LayerNorm(hidden_states + input_tensor)
- return hidden_states
-
-
-class GitLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.chunk_size_feed_forward = config.chunk_size_feed_forward
- self.seq_len_dim = 1
- self.attention = GitAttention(config)
- self.intermediate = GitIntermediate(config)
- self.output = GitOutput(config)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- output_attentions: Optional[bool] = False,
- pixel_values_present: Optional[bool] = False,
- ) -> Tuple[torch.Tensor]:
- # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
- self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
- self_attention_outputs = self.attention(
- hidden_states,
- attention_mask,
- head_mask,
- output_attentions=output_attentions,
- past_key_value=self_attn_past_key_value,
- pixel_values_present=pixel_values_present,
- )
- attention_output = self_attention_outputs[0]
-
- # if decoder, the last output is tuple of self-attn cache
- outputs = self_attention_outputs[1:-1]
- present_key_value = self_attention_outputs[-1]
-
- layer_output = apply_chunking_to_forward(
- self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
- )
- outputs = (layer_output,) + outputs
-
- # if decoder, return the attn key/values as the last output
- outputs = outputs + (present_key_value,)
-
- return outputs
-
- def feed_forward_chunk(self, attention_output):
- intermediate_output = self.intermediate(attention_output)
- layer_output = self.output(intermediate_output, attention_output)
- return layer_output
-
-
-class GitEncoder(nn.Module):
- # Copied from transformers.models.bert.modeling_bert.BertEncoder.__init__ with Bert->Git
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.layer = nn.ModuleList([GitLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = False,
- output_hidden_states: Optional[bool] = False,
- pixel_values_present: Optional[bool] = False,
- return_dict: Optional[bool] = True,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- next_decoder_cache = () if use_cache else None
- for i, layer_module in enumerate(self.layer):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- layer_head_mask = head_mask[i] if head_mask is not None else None
- past_key_value = past_key_values[i] if past_key_values is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- layer_module.__call__,
- hidden_states,
- attention_mask,
- layer_head_mask,
- past_key_value,
- output_attentions,
- )
- else:
- layer_outputs = layer_module(
- hidden_states,
- attention_mask,
- layer_head_mask,
- past_key_value,
- output_attentions,
- pixel_values_present,
- )
-
- hidden_states = layer_outputs[0]
- if use_cache:
- next_decoder_cache += (layer_outputs[-1],)
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- next_decoder_cache,
- all_hidden_states,
- all_self_attentions,
- ]
- if v is not None
- )
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=next_decoder_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-class GitPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GitConfig
- base_model_prefix = "git"
- supports_gradient_checkpointing = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, GitVisionEmbeddings):
- nn.init.normal_(module.class_embedding, mean=0.0, std=self.config.initializer_range)
- nn.init.normal_(module.patch_embedding.weight, std=self.config.initializer_range)
- nn.init.normal_(module.position_embedding.weight, std=self.config.initializer_range)
- if isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-GIT_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`GitConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GIT_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
-
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
-
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
- [`CLIPImageProcessor.__call__`] for details.
-
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-# Copied from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings with CLIP->Git
-class GitVisionEmbeddings(nn.Module):
- def __init__(self, config: GitVisionConfig):
- super().__init__()
- self.config = config
- self.embed_dim = config.hidden_size
- self.image_size = config.image_size
- self.patch_size = config.patch_size
-
- self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
-
- self.patch_embedding = nn.Conv2d(
- in_channels=config.num_channels,
- out_channels=self.embed_dim,
- kernel_size=self.patch_size,
- stride=self.patch_size,
- bias=False,
- )
-
- self.num_patches = (self.image_size // self.patch_size) ** 2
- self.num_positions = self.num_patches + 1
- self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
- self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
-
- def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
- batch_size = pixel_values.shape[0]
- target_dtype = self.patch_embedding.weight.dtype
- patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid]
- patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
-
- class_embeds = self.class_embedding.expand(batch_size, 1, -1)
- embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
- embeddings = embeddings + self.position_embedding(self.position_ids)
- return embeddings
-
-
-# Copied from transformers.models.clip.modeling_clip.CLIPMLP
-class GitVisionMLP(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.activation_fn = ACT2FN[config.hidden_act]
- self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
- self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- hidden_states = self.fc1(hidden_states)
- hidden_states = self.activation_fn(hidden_states)
- hidden_states = self.fc2(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.clip.modeling_clip.CLIPAttention
-class GitVisionAttention(nn.Module):
- """Multi-headed attention from 'Attention Is All You Need' paper"""
-
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_heads
- if self.head_dim * self.num_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {self.num_heads})."
- )
- self.scale = self.head_dim**-0.5
- self.dropout = config.attention_dropout
-
- self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
- self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
- self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
- self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
-
- def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
- return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- causal_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel"""
-
- bsz, tgt_len, embed_dim = hidden_states.size()
-
- # get query proj
- query_states = self.q_proj(hidden_states) * self.scale
- key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
- value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-
- proj_shape = (bsz * self.num_heads, -1, self.head_dim)
- query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
- key_states = key_states.view(*proj_shape)
- value_states = value_states.view(*proj_shape)
-
- src_len = key_states.size(1)
- attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-
- if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
- raise ValueError(
- f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- # apply the causal_attention_mask first
- if causal_attention_mask is not None:
- if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
- raise ValueError(
- f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
- f" {causal_attention_mask.size()}"
- )
- attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- if attention_mask is not None:
- if attention_mask.size() != (bsz, 1, tgt_len, src_len):
- raise ValueError(
- f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
- )
- attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if output_attentions:
- # this operation is a bit akward, but it's required to
- # make sure that attn_weights keeps its gradient.
- # In order to do so, attn_weights have to reshaped
- # twice and have to be reused in the following
- attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
- attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-
- attn_output = torch.bmm(attn_probs, value_states)
-
- if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
- attn_output = attn_output.transpose(1, 2)
- attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped
-
-
-# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->GitVision
-class GitVisionEncoderLayer(nn.Module):
- def __init__(self, config: GitVisionConfig):
- super().__init__()
- self.embed_dim = config.hidden_size
- self.self_attn = GitVisionAttention(config)
- self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
- self.mlp = GitVisionMLP(config)
- self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: torch.Tensor,
- causal_attention_mask: torch.Tensor,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[torch.FloatTensor]:
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`): attention mask of size
- `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
- `(config.encoder_attention_heads,)`.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- """
- residual = hidden_states
-
- hidden_states = self.layer_norm1(hidden_states)
- hidden_states, attn_weights = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- causal_attention_mask=causal_attention_mask,
- output_attentions=output_attentions,
- )
- hidden_states = residual + hidden_states
-
- residual = hidden_states
- hidden_states = self.layer_norm2(hidden_states)
- hidden_states = self.mlp(hidden_states)
- hidden_states = residual + hidden_states
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
-
-# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->GitVision, CLIPConfig
-class GitVisionEncoder(nn.Module):
- """
- Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
- [`GitVisionEncoderLayer`].
-
- Args:
- config: GitVisionConfig
- """
-
- def __init__(self, config: GitVisionConfig):
- super().__init__()
- self.config = config
- self.layers = nn.ModuleList([GitVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
- self.gradient_checkpointing = False
-
- def forward(
- self,
- inputs_embeds,
- attention_mask: Optional[torch.Tensor] = None,
- causal_attention_mask: Optional[torch.Tensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- r"""
- Args:
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
- This is useful if you want more control over how to convert `input_ids` indices into associated vectors
- than the model's internal embedding lookup matrix.
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Causal mask for the text model. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
- for more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- encoder_states = () if output_hidden_states else None
- all_attentions = () if output_attentions else None
-
- hidden_states = inputs_embeds
- for idx, encoder_layer in enumerate(self.layers):
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
- if self.gradient_checkpointing and self.training:
- layer_outputs = self._gradient_checkpointing_func(
- encoder_layer.__call__,
- hidden_states,
- attention_mask,
- causal_attention_mask,
- output_attentions,
- )
- else:
- layer_outputs = encoder_layer(
- hidden_states,
- attention_mask,
- causal_attention_mask,
- output_attentions=output_attentions,
- )
-
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions = all_attentions + (layer_outputs[1],)
-
- if output_hidden_states:
- encoder_states = encoder_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
- )
-
-
-GIT_VISION_INPUTS_DOCSTRING = r"""
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
- [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class GitVisionTransformer(nn.Module):
- # Copied from transformers.models.clip.modeling_clip.CLIPVisionTransformer.__init__ with CLIPEncoder->GitVisionEncoder, CLIP->Git
- def __init__(self, config: GitVisionConfig):
- super().__init__()
- self.config = config
- embed_dim = config.hidden_size
-
- self.embeddings = GitVisionEmbeddings(config)
- self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
- self.encoder = GitVisionEncoder(config)
- self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
-
- @add_start_docstrings_to_model_forward(GIT_VISION_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BaseModelOutput, config_class=GitVisionConfig)
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- r"""
- Returns:
-
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if pixel_values is None:
- raise ValueError("You have to specify pixel_values")
-
- hidden_states = self.embeddings(pixel_values)
- hidden_states = self.pre_layrnorm(hidden_states)
-
- encoder_outputs = self.encoder(
- inputs_embeds=hidden_states,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- last_hidden_state = encoder_outputs[0]
-
- last_hidden_state = self.post_layernorm(last_hidden_state)
-
- if not return_dict:
- return (last_hidden_state,) + encoder_outputs[1:]
-
- return BaseModelOutput(
- last_hidden_state=last_hidden_state,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """The vision model from CLIP, used in GIT, without any head or projection on top.""",
- GIT_START_DOCSTRING,
-)
-class GitVisionModel(GitPreTrainedModel):
- config_class = GitVisionConfig
- main_input_name = "pixel_values"
-
- # Copied from transformers.models.clip.modeling_clip.CLIPVisionModel.__init__ with CLIP->Git
- def __init__(self, config: GitVisionConfig):
- super().__init__(config)
- self.vision_model = GitVisionTransformer(config)
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self) -> nn.Module:
- return self.vision_model.embeddings.patch_embedding
-
- @add_start_docstrings_to_model_forward(GIT_VISION_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=BaseModelOutput, config_class=GitVisionConfig)
- def forward(
- self,
- pixel_values: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- r"""
- Returns:
-
- Examples:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoProcessor, GitVisionModel
-
- >>> processor = AutoProcessor.from_pretrained("microsoft/git-base")
- >>> model = GitVisionModel.from_pretrained("microsoft/git-base")
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> inputs = processor(images=image, return_tensors="pt")
-
- >>> outputs = model(**inputs)
- >>> last_hidden_state = outputs.last_hidden_state
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- return self.vision_model(
- pixel_values=pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
-
-class GitProjection(nn.Module):
- def __init__(self, config: GitConfig):
- super().__init__()
- self.config = config
- self.visual_projection = nn.Sequential(
- nn.Linear(config.vision_config.hidden_size, config.hidden_size),
- nn.LayerNorm(config.hidden_size, eps=config.vision_config.layer_norm_eps),
- )
-
- def forward(self, embeddings: torch.Tensor) -> torch.Tensor:
- return self.visual_projection(embeddings)
-
-
-@add_start_docstrings(
- "The bare GIT Model transformer consisting of a CLIP image encoder and text decoder outputting raw hidden-states"
- " without any specific head on top.",
- GIT_START_DOCSTRING,
-)
-class GitModel(GitPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.config = config
-
- self.embeddings = GitEmbeddings(config)
- self.image_encoder = GitVisionModel(config.vision_config)
- self.encoder = GitEncoder(config)
-
- self.visual_projection = GitProjection(config)
-
- if config.num_image_with_embedding is not None:
- self.img_temperal_embedding = nn.ParameterList(
- nn.Parameter(torch.zeros(1, 1, config.vision_config.hidden_size))
- for _ in range(config.num_image_with_embedding)
- )
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embeddings.word_embeddings
-
- def set_input_embeddings(self, value):
- self.embeddings.word_embeddings = value
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- def _generate_future_mask(self, size: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
- # Default mask is for forward direction. Flip for backward direction.
- mask = torch.triu(torch.ones(size, size, device=device, dtype=dtype), diagonal=1)
- mask = mask.masked_fill(mask == 1, float("-inf"))
- return mask
-
- def create_attention_mask(self, tgt, memory, tgt_mask, past_key_values_length, memory_key_padding_mask=None):
- num_tgt = tgt.shape[1]
- num_memory = memory.shape[1]
- device = tgt.device
- dtype = tgt.dtype
- top_left = torch.zeros((num_memory, num_memory), device=device, dtype=dtype)
- top_right = torch.full(
- (num_memory, num_tgt + past_key_values_length),
- float("-inf"),
- device=tgt.device,
- dtype=dtype,
- )
- bottom_left = torch.zeros(
- (num_tgt, num_memory),
- dtype=dtype,
- device=tgt_mask.device,
- )
-
- if past_key_values_length > 0:
- tgt_mask = torch.zeros(
- (tgt_mask.shape[0], tgt_mask.shape[0] + past_key_values_length),
- dtype=dtype,
- device=tgt_mask.device,
- )
-
- left = torch.cat((top_left, bottom_left), dim=0)
- right = torch.cat((top_right, tgt_mask.to(dtype)), dim=0)
-
- full_attention_mask = torch.cat((left, right), dim=1)[None, :]
-
- if memory_key_padding_mask is None:
- memory_key_padding_mask = torch.full((memory.shape[0], memory.shape[1]), fill_value=False, device=device)
- # if it is False, it means valid. That is, it is not a padding
- if memory_key_padding_mask.dtype != torch.bool:
- raise ValueError("Memory key padding mask must be a boolean tensor.")
- zero_negative_infinity = torch.zeros_like(memory_key_padding_mask, dtype=tgt.dtype)
- zero_negative_infinity[memory_key_padding_mask] = float("-inf")
- full_attention_mask = full_attention_mask.expand(
- (memory_key_padding_mask.shape[0], num_memory + num_tgt, num_memory + past_key_values_length + num_tgt)
- )
- full_attention_mask = full_attention_mask.clone()
- origin_left = full_attention_mask[:, :, :num_memory]
- update = zero_negative_infinity[:, None, :]
- full_attention_mask[:, :, :num_memory] = origin_left + update
-
- # add axis for multi-head
- full_attention_mask = full_attention_mask[:, None, :, :]
-
- return full_attention_mask
-
- @add_start_docstrings_to_model_forward(GIT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- pixel_values: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:
- r"""
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoProcessor, AutoModel
- >>> import requests
- >>> from PIL import Image
-
- >>> processor = AutoProcessor.from_pretrained("microsoft/git-base")
- >>> model = AutoModel.from_pretrained("microsoft/git-base")
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> text = "this is an image of two cats"
-
- >>> inputs = processor(text, images=image, return_tensors="pt")
-
- >>> outputs = model(**inputs)
- >>> last_hidden_state = outputs.last_hidden_state
- ```"""
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- seq_length = input_shape[1]
-
- # past_key_values_length
- past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- projected_visual_features = None
- if pixel_values is not None:
- if pixel_values.ndim == 4:
- # here we assume pixel_values is of shape (batch_size, num_channels, height, width)
- visual_features = self.image_encoder(pixel_values).last_hidden_state
-
- elif pixel_values.ndim == 5:
- # here we assume pixel_values is of shape (batch_size, num_frames, num_channels, height, width)
- visual_features = []
- for frame_idx in range(pixel_values.shape[1]):
- visual_features_frame = self.image_encoder(pixel_values[:, frame_idx, :, :]).last_hidden_state
- visual_features_frame += self.img_temperal_embedding[frame_idx]
- visual_features.append(visual_features_frame)
-
- # finally, concatenate all features along sequence dimension
- visual_features = torch.cat(visual_features, dim=1)
-
- else:
- raise ValueError("pixel_values must be of rank 4 or 5")
-
- projected_visual_features = self.visual_projection(visual_features)
-
- embedding_output = self.embeddings(
- input_ids=input_ids,
- position_ids=position_ids,
- inputs_embeds=inputs_embeds,
- past_key_values_length=past_key_values_length,
- )
-
- if projected_visual_features is None:
- projected_visual_features = torch.zeros(
- (embedding_output.shape[0], 0, embedding_output.shape[2]),
- dtype=embedding_output.dtype,
- device=embedding_output.device,
- )
-
- # Repeat visual features to match embedding batch size.
- projected_visual_features = projected_visual_features.repeat(
- embedding_output.size(0) // projected_visual_features.size(0), 1, 1
- )
-
- # concatenate patch token and text token embeddings
- hidden_states = torch.cat((projected_visual_features, embedding_output), dim=1)
-
- # By default, an additive causal mask is created
- # for masking the future (one direction).
- tgt_mask = self._generate_future_mask(seq_length, embedding_output.dtype, embedding_output.device)
-
- # Create an attention mask of shape (batch_size, 1, tgt_seq_len, src_seq_len)
- combined_attention_mask = self.create_attention_mask(
- tgt=embedding_output,
- memory=projected_visual_features,
- tgt_mask=tgt_mask,
- past_key_values_length=past_key_values_length,
- )
-
- if attention_mask is not None:
- # if the user provides an attention mask, we add it to the default one
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
- expanded_attn_mask = _prepare_4d_attention_mask(
- attention_mask, embedding_output.dtype, tgt_len=input_shape[-1]
- ).to(embedding_output.device)
- if past_key_values_length > 0:
- expanded_attn_mask = expanded_attn_mask[:, :, -past_key_values_length:, :]
- else:
- combined_attention_mask[:, :, -input_shape[1] :, -input_shape[1] :] += expanded_attn_mask
-
- encoder_outputs = self.encoder(
- hidden_states,
- attention_mask=combined_attention_mask,
- head_mask=head_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- pixel_values_present=pixel_values is not None,
- )
- sequence_output = encoder_outputs[0]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[1:]
-
- return BaseModelOutputWithPast(
- last_hidden_state=sequence_output,
- past_key_values=encoder_outputs.past_key_values,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """GIT Model with a `language modeling` head on top for autoregressive language modeling.""", GIT_START_DOCSTRING
-)
-class GitForCausalLM(GitPreTrainedModel):
- _tied_weights_keys = ["output.weight"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.git = GitModel(config)
- self.output = nn.Linear(config.hidden_size, config.vocab_size)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.output
-
- def set_output_embeddings(self, new_embeddings):
- self.output = new_embeddings
-
- @add_start_docstrings_to_model_forward(GIT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- pixel_values: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.Tensor]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
- `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
- ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
-
- Returns:
-
- Examples:
-
- Image captioning example:
-
- ```python
- >>> from transformers import AutoProcessor, AutoModelForCausalLM
- >>> import requests
- >>> from PIL import Image
-
- >>> processor = AutoProcessor.from_pretrained("microsoft/git-base-coco")
- >>> model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-coco")
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> pixel_values = processor(images=image, return_tensors="pt").pixel_values
-
- >>> generated_ids = model.generate(pixel_values=pixel_values, max_length=50)
- >>> generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
- >>> print(generated_caption)
- two cats sleeping on a pink blanket next to remotes.
- ```
-
- Visual question answering (VQA) example:
-
- ```python
- >>> from transformers import AutoProcessor, AutoModelForCausalLM
- >>> from huggingface_hub import hf_hub_download
- >>> from PIL import Image
-
- >>> processor = AutoProcessor.from_pretrained("microsoft/git-base-textvqa")
- >>> model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-textvqa")
-
- >>> file_path = hf_hub_download(repo_id="nielsr/textvqa-sample", filename="bus.png", repo_type="dataset")
- >>> image = Image.open(file_path).convert("RGB")
-
- >>> pixel_values = processor(images=image, return_tensors="pt").pixel_values
-
- >>> question = "what does the front of the bus say at the top?"
-
- >>> input_ids = processor(text=question, add_special_tokens=False).input_ids
- >>> input_ids = [processor.tokenizer.cls_token_id] + input_ids
- >>> input_ids = torch.tensor(input_ids).unsqueeze(0)
-
- >>> generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50)
- >>> print(processor.batch_decode(generated_ids, skip_special_tokens=True))
- ['what does the front of the bus say at the top? special']
- ```
-
- Video captioning example:
-
- ```python
- >>> import av
- >>> import numpy as np
- >>> from PIL import Image
- >>> from huggingface_hub import hf_hub_download
- >>> from transformers import AutoProcessor, AutoModelForCausalLM
-
- >>> processor = AutoProcessor.from_pretrained("microsoft/git-base-vatex")
- >>> model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-vatex")
-
- >>> # set seed for reproducability
- >>> np.random.seed(45)
-
-
- >>> def read_video_pyav(container, indices):
- ... '''
- ... Decode the video with PyAV decoder.
- ... Args:
- ... container (`av.container.input.InputContainer`): PyAV container.
- ... indices (`List[int]`): List of frame indices to decode.
- ... Returns:
- ... result (np.ndarray): np array of decoded frames of shape (num_frames, height, width, 3).
- ... '''
- ... frames = []
- ... container.seek(0)
- ... start_index = indices[0]
- ... end_index = indices[-1]
- ... for i, frame in enumerate(container.decode(video=0)):
- ... if i > end_index:
- ... break
- ... if i >= start_index and i in indices:
- ... frames.append(frame)
- ... return np.stack([x.to_ndarray(format="rgb24") for x in frames])
-
-
- >>> def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
- ... '''
- ... Sample a given number of frame indices from the video.
- ... Args:
- ... clip_len (`int`): Total number of frames to sample.
- ... frame_sample_rate (`int`): Sample every n-th frame.
- ... seg_len (`int`): Maximum allowed index of sample's last frame.
- ... Returns:
- ... indices (`List[int]`): List of sampled frame indices
- ... '''
- ... converted_len = int(clip_len * frame_sample_rate)
- ... end_idx = np.random.randint(converted_len, seg_len)
- ... start_idx = end_idx - converted_len
- ... indices = np.linspace(start_idx, end_idx, num=clip_len)
- ... indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
- ... return indices
-
-
- >>> # load video
- >>> file_path = hf_hub_download(
- ... repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset"
- ... )
- >>> container = av.open(file_path)
-
- >>> # sample frames
- >>> num_frames = model.config.num_image_with_embedding
- >>> indices = sample_frame_indices(
- ... clip_len=num_frames, frame_sample_rate=4, seg_len=container.streams.video[0].frames
- ... )
- >>> frames = read_video_pyav(container, indices)
-
- >>> pixel_values = processor(images=list(frames), return_tensors="pt").pixel_values
-
- >>> generated_ids = model.generate(pixel_values=pixel_values, max_length=50)
-
- >>> print("Generated caption:", processor.batch_decode(generated_ids, skip_special_tokens=True))
- Generated caption: ['a woman is sitting at a table and she is talking about the food she is holding.']
- ```
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- if labels is not None:
- use_cache = False
-
- outputs = self.git(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- pixel_values=pixel_values,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
- logits = self.output(sequence_output)
-
- loss = None
- if labels is not None:
- # we are doing next-token prediction; shift prediction scores and input ids by one
- num_image_tokens = self.git.encoder.layer[0].attention.self.image_patch_tokens
- shifted_logits = logits[:, num_image_tokens:-1, :].contiguous()
- labels = labels[:, 1:].contiguous()
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(shifted_logits.view(-1, self.config.vocab_size), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def prepare_inputs_for_generation(
- self, input_ids, past_key_values=None, attention_mask=None, use_cache=None, **kwargs
- ):
- # cut decoder_input_ids if past_key_values is used
- if past_key_values is not None:
- input_ids = input_ids[:, -1:]
-
- # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
- input_shape = input_ids.shape
- if attention_mask is None:
- attention_mask = input_ids.new_ones(input_shape)
-
- return {
- "input_ids": input_ids,
- "attention_mask": attention_mask,
- "pixel_values": kwargs.get("pixel_values", None),
- "past_key_values": past_key_values,
- "use_cache": use_cache,
- }
-
- def _reorder_cache(self, past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
- )
- return reordered_past
diff --git a/transformers/models/git/processing_git.py b/transformers/models/git/processing_git.py
deleted file mode 100644
index 79f26f3bf24b14116c43a66b7e9ebb682ee667a2..0000000000000000000000000000000000000000
--- a/transformers/models/git/processing_git.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Image/Text processor class for GIT
-"""
-
-from ...processing_utils import ProcessorMixin
-from ...tokenization_utils_base import BatchEncoding
-
-
-class GitProcessor(ProcessorMixin):
- r"""
- Constructs a GIT processor which wraps a CLIP image processor and a BERT tokenizer into a single processor.
-
- [`GitProcessor`] offers all the functionalities of [`CLIPImageProcessor`] and [`BertTokenizerFast`]. See the
- [`~GitProcessor.__call__`] and [`~GitProcessor.decode`] for more information.
-
- Args:
- image_processor ([`AutoImageProcessor`]):
- The image processor is a required input.
- tokenizer ([`AutoTokenizer`]):
- The tokenizer is a required input.
- """
-
- attributes = ["image_processor", "tokenizer"]
- image_processor_class = "AutoImageProcessor"
- tokenizer_class = "AutoTokenizer"
-
- def __init__(self, image_processor, tokenizer):
- super().__init__(image_processor, tokenizer)
- self.current_processor = self.image_processor
-
- def __call__(self, text=None, images=None, return_tensors=None, **kwargs):
- """
- Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
- and `kwargs` arguments to BertTokenizerFast's [`~BertTokenizerFast.__call__`] if `text` is not `None` to encode
- the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to
- CLIPImageProcessor's [`~CLIPImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
- of the above two methods for more information.
-
- Args:
- text (`str`, `List[str]`, `List[List[str]]`):
- The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
- (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
- `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
- images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
- The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
- tensor. Both channels-first and channels-last formats are supported.
-
- return_tensors (`str` or [`~utils.TensorType`], *optional*):
- If set, will return tensors of a particular framework. Acceptable values are:
-
- - `'tf'`: Return TensorFlow `tf.constant` objects.
- - `'pt'`: Return PyTorch `torch.Tensor` objects.
- - `'np'`: Return NumPy `np.ndarray` objects.
- - `'jax'`: Return JAX `jnp.ndarray` objects.
-
- Returns:
- [`BatchEncoding`]: A [`BatchEncoding`] with the following fields:
-
- - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
- - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
- `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
- `None`).
- - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
- """
-
- if text is None and images is None:
- raise ValueError("You have to specify either text or images. Both cannot be none.")
-
- if text is not None:
- encoding = self.tokenizer(text, return_tensors=return_tensors, **kwargs)
-
- if images is not None:
- image_features = self.image_processor(images, return_tensors=return_tensors, **kwargs)
-
- if text is not None and images is not None:
- encoding["pixel_values"] = image_features.pixel_values
- return encoding
- elif text is not None:
- return encoding
- else:
- return BatchEncoding(data=dict(**image_features), tensor_type=return_tensors)
-
- def batch_decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to BertTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
- refer to the docstring of this method for more information.
- """
- return self.tokenizer.batch_decode(*args, **kwargs)
-
- def decode(self, *args, **kwargs):
- """
- This method forwards all its arguments to BertTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
- the docstring of this method for more information.
- """
- return self.tokenizer.decode(*args, **kwargs)
-
- @property
- def model_input_names(self):
- return ["input_ids", "attention_mask", "pixel_values"]
diff --git a/transformers/models/glpn/__init__.py b/transformers/models/glpn/__init__.py
deleted file mode 100644
index 94788dcb85e76faa2f312df8d13f5577c21a88d1..0000000000000000000000000000000000000000
--- a/transformers/models/glpn/__init__.py
+++ /dev/null
@@ -1,75 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available
-
-
-_import_structure = {"configuration_glpn": ["GLPN_PRETRAINED_CONFIG_ARCHIVE_MAP", "GLPNConfig"]}
-
-try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["feature_extraction_glpn"] = ["GLPNFeatureExtractor"]
- _import_structure["image_processing_glpn"] = ["GLPNImageProcessor"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_glpn"] = [
- "GLPN_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GLPNForDepthEstimation",
- "GLPNLayer",
- "GLPNModel",
- "GLPNPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_glpn import GLPN_PRETRAINED_CONFIG_ARCHIVE_MAP, GLPNConfig
-
- try:
- if not is_vision_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .feature_extraction_glpn import GLPNFeatureExtractor
- from .image_processing_glpn import GLPNImageProcessor
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_glpn import (
- GLPN_PRETRAINED_MODEL_ARCHIVE_LIST,
- GLPNForDepthEstimation,
- GLPNLayer,
- GLPNModel,
- GLPNPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/glpn/__pycache__/__init__.cpython-310.pyc b/transformers/models/glpn/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 764f172bfcc28fa9db8673137c5d7cc23d38c428..0000000000000000000000000000000000000000
Binary files a/transformers/models/glpn/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/glpn/__pycache__/configuration_glpn.cpython-310.pyc b/transformers/models/glpn/__pycache__/configuration_glpn.cpython-310.pyc
deleted file mode 100644
index cb9fa0777bf59f8750ed45f6fa4ce3813a31e483..0000000000000000000000000000000000000000
Binary files a/transformers/models/glpn/__pycache__/configuration_glpn.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/glpn/__pycache__/convert_glpn_to_pytorch.cpython-310.pyc b/transformers/models/glpn/__pycache__/convert_glpn_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 50de59bdcfa12ed306803ac241a234cddaac1694..0000000000000000000000000000000000000000
Binary files a/transformers/models/glpn/__pycache__/convert_glpn_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/glpn/__pycache__/feature_extraction_glpn.cpython-310.pyc b/transformers/models/glpn/__pycache__/feature_extraction_glpn.cpython-310.pyc
deleted file mode 100644
index 7fc19fed1259e8096b6ec33c3a98e507108f9803..0000000000000000000000000000000000000000
Binary files a/transformers/models/glpn/__pycache__/feature_extraction_glpn.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/glpn/__pycache__/image_processing_glpn.cpython-310.pyc b/transformers/models/glpn/__pycache__/image_processing_glpn.cpython-310.pyc
deleted file mode 100644
index 6c9ae113b3fe86f74c3653ece9fb80a61b6a57c9..0000000000000000000000000000000000000000
Binary files a/transformers/models/glpn/__pycache__/image_processing_glpn.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/glpn/__pycache__/modeling_glpn.cpython-310.pyc b/transformers/models/glpn/__pycache__/modeling_glpn.cpython-310.pyc
deleted file mode 100644
index 1bc4bd848fed37b1bbd5917f77a732116d462aaf..0000000000000000000000000000000000000000
Binary files a/transformers/models/glpn/__pycache__/modeling_glpn.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/glpn/configuration_glpn.py b/transformers/models/glpn/configuration_glpn.py
deleted file mode 100644
index c3341192169aa09597f6c51ad07ff5305ef1a4e6..0000000000000000000000000000000000000000
--- a/transformers/models/glpn/configuration_glpn.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# coding=utf-8
-# Copyright 2022 KAIST and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GLPN model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GLPN_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GLPNConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GLPNModel`]. It is used to instantiate an GLPN
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the GLPN
- [vinvino02/glpn-kitti](https://huggingface.co/vinvino02/glpn-kitti) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Args:
- num_channels (`int`, *optional*, defaults to 3):
- The number of input channels.
- num_encoder_blocks (`int`, *optional*, defaults to 4):
- The number of encoder blocks (i.e. stages in the Mix Transformer encoder).
- depths (`List[int]`, *optional*, defaults to `[2, 2, 2, 2]`):
- The number of layers in each encoder block.
- sr_ratios (`List[int]`, *optional*, defaults to `[8, 4, 2, 1]`):
- Sequence reduction ratios in each encoder block.
- hidden_sizes (`List[int]`, *optional*, defaults to `[32, 64, 160, 256]`):
- Dimension of each of the encoder blocks.
- patch_sizes (`List[int]`, *optional*, defaults to `[7, 3, 3, 3]`):
- Patch size before each encoder block.
- strides (`List[int]`, *optional*, defaults to `[4, 2, 2, 2]`):
- Stride before each encoder block.
- num_attention_heads (`List[int]`, *optional*, defaults to `[1, 2, 5, 8]`):
- Number of attention heads for each attention layer in each block of the Transformer encoder.
- mlp_ratios (`List[int]`, *optional*, defaults to `[4, 4, 4, 4]`):
- Ratio of the size of the hidden layer compared to the size of the input layer of the Mix FFNs in the
- encoder blocks.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- drop_path_rate (`float`, *optional*, defaults to 0.1):
- The dropout probability for stochastic depth, used in the blocks of the Transformer encoder.
- layer_norm_eps (`float`, *optional*, defaults to 1e-06):
- The epsilon used by the layer normalization layers.
- decoder_hidden_size (`int`, *optional*, defaults to 64):
- The dimension of the decoder.
- max_depth (`int`, *optional*, defaults to 10):
- The maximum depth of the decoder.
- head_in_index (`int`, *optional*, defaults to -1):
- The index of the features to use in the head.
-
- Example:
-
- ```python
- >>> from transformers import GLPNModel, GLPNConfig
-
- >>> # Initializing a GLPN vinvino02/glpn-kitti style configuration
- >>> configuration = GLPNConfig()
-
- >>> # Initializing a model from the vinvino02/glpn-kitti style configuration
- >>> model = GLPNModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "glpn"
-
- def __init__(
- self,
- num_channels=3,
- num_encoder_blocks=4,
- depths=[2, 2, 2, 2],
- sr_ratios=[8, 4, 2, 1],
- hidden_sizes=[32, 64, 160, 256],
- patch_sizes=[7, 3, 3, 3],
- strides=[4, 2, 2, 2],
- num_attention_heads=[1, 2, 5, 8],
- mlp_ratios=[4, 4, 4, 4],
- hidden_act="gelu",
- hidden_dropout_prob=0.0,
- attention_probs_dropout_prob=0.0,
- initializer_range=0.02,
- drop_path_rate=0.1,
- layer_norm_eps=1e-6,
- decoder_hidden_size=64,
- max_depth=10,
- head_in_index=-1,
- **kwargs,
- ):
- super().__init__(**kwargs)
-
- self.num_channels = num_channels
- self.num_encoder_blocks = num_encoder_blocks
- self.depths = depths
- self.sr_ratios = sr_ratios
- self.hidden_sizes = hidden_sizes
- self.patch_sizes = patch_sizes
- self.strides = strides
- self.mlp_ratios = mlp_ratios
- self.num_attention_heads = num_attention_heads
- self.hidden_act = hidden_act
- self.hidden_dropout_prob = hidden_dropout_prob
- self.attention_probs_dropout_prob = attention_probs_dropout_prob
- self.initializer_range = initializer_range
- self.drop_path_rate = drop_path_rate
- self.layer_norm_eps = layer_norm_eps
- self.decoder_hidden_size = decoder_hidden_size
- self.max_depth = max_depth
- self.head_in_index = head_in_index
diff --git a/transformers/models/glpn/convert_glpn_to_pytorch.py b/transformers/models/glpn/convert_glpn_to_pytorch.py
deleted file mode 100644
index 5f0183783ec812f69766d9220efb58652a21cb87..0000000000000000000000000000000000000000
--- a/transformers/models/glpn/convert_glpn_to_pytorch.py
+++ /dev/null
@@ -1,219 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert GLPN checkpoints."""
-
-
-import argparse
-from collections import OrderedDict
-from pathlib import Path
-
-import requests
-import torch
-from PIL import Image
-
-from transformers import GLPNConfig, GLPNForDepthEstimation, GLPNImageProcessor
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-logger = logging.get_logger(__name__)
-
-
-def rename_keys(state_dict):
- new_state_dict = OrderedDict()
- for key, value in state_dict.items():
- if key.startswith("module.encoder"):
- key = key.replace("module.encoder", "glpn.encoder")
- if key.startswith("module.decoder"):
- key = key.replace("module.decoder", "decoder.stages")
- if "patch_embed" in key:
- # replace for example patch_embed1 by patch_embeddings.0
- idx = key[key.find("patch_embed") + len("patch_embed")]
- key = key.replace(f"patch_embed{idx}", f"patch_embeddings.{int(idx)-1}")
- if "norm" in key:
- key = key.replace("norm", "layer_norm")
- if "glpn.encoder.layer_norm" in key:
- # replace for example layer_norm1 by layer_norm.0
- idx = key[key.find("glpn.encoder.layer_norm") + len("glpn.encoder.layer_norm")]
- key = key.replace(f"layer_norm{idx}", f"layer_norm.{int(idx)-1}")
- if "layer_norm1" in key:
- key = key.replace("layer_norm1", "layer_norm_1")
- if "layer_norm2" in key:
- key = key.replace("layer_norm2", "layer_norm_2")
- if "block" in key:
- # replace for example block1 by block.0
- idx = key[key.find("block") + len("block")]
- key = key.replace(f"block{idx}", f"block.{int(idx)-1}")
- if "attn.q" in key:
- key = key.replace("attn.q", "attention.self.query")
- if "attn.proj" in key:
- key = key.replace("attn.proj", "attention.output.dense")
- if "attn" in key:
- key = key.replace("attn", "attention.self")
- if "fc1" in key:
- key = key.replace("fc1", "dense1")
- if "fc2" in key:
- key = key.replace("fc2", "dense2")
- if "linear_pred" in key:
- key = key.replace("linear_pred", "classifier")
- if "linear_fuse" in key:
- key = key.replace("linear_fuse.conv", "linear_fuse")
- key = key.replace("linear_fuse.bn", "batch_norm")
- if "linear_c" in key:
- # replace for example linear_c4 by linear_c.3
- idx = key[key.find("linear_c") + len("linear_c")]
- key = key.replace(f"linear_c{idx}", f"linear_c.{int(idx)-1}")
- if "bot_conv" in key:
- key = key.replace("bot_conv", "0.convolution")
- if "skip_conv1" in key:
- key = key.replace("skip_conv1", "1.convolution")
- if "skip_conv2" in key:
- key = key.replace("skip_conv2", "2.convolution")
- if "fusion1" in key:
- key = key.replace("fusion1", "1.fusion")
- if "fusion2" in key:
- key = key.replace("fusion2", "2.fusion")
- if "fusion3" in key:
- key = key.replace("fusion3", "3.fusion")
- if "fusion" in key and "conv" in key:
- key = key.replace("conv", "convolutional_layer")
- if key.startswith("module.last_layer_depth"):
- key = key.replace("module.last_layer_depth", "head.head")
- new_state_dict[key] = value
-
- return new_state_dict
-
-
-def read_in_k_v(state_dict, config):
- # for each of the encoder blocks:
- for i in range(config.num_encoder_blocks):
- for j in range(config.depths[i]):
- # read in weights + bias of keys and values (which is a single matrix in the original implementation)
- kv_weight = state_dict.pop(f"glpn.encoder.block.{i}.{j}.attention.self.kv.weight")
- kv_bias = state_dict.pop(f"glpn.encoder.block.{i}.{j}.attention.self.kv.bias")
- # next, add keys and values (in that order) to the state dict
- state_dict[f"glpn.encoder.block.{i}.{j}.attention.self.key.weight"] = kv_weight[
- : config.hidden_sizes[i], :
- ]
- state_dict[f"glpn.encoder.block.{i}.{j}.attention.self.key.bias"] = kv_bias[: config.hidden_sizes[i]]
- state_dict[f"glpn.encoder.block.{i}.{j}.attention.self.value.weight"] = kv_weight[
- config.hidden_sizes[i] :, :
- ]
- state_dict[f"glpn.encoder.block.{i}.{j}.attention.self.value.bias"] = kv_bias[config.hidden_sizes[i] :]
-
-
-# We will verify our results on a COCO image
-def prepare_img():
- url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- image = Image.open(requests.get(url, stream=True).raw)
-
- return image
-
-
-@torch.no_grad()
-def convert_glpn_checkpoint(checkpoint_path, pytorch_dump_folder_path, push_to_hub=False, model_name=None):
- """
- Copy/paste/tweak model's weights to our GLPN structure.
- """
-
- # load GLPN configuration (Segformer-B4 size)
- config = GLPNConfig(hidden_sizes=[64, 128, 320, 512], decoder_hidden_size=64, depths=[3, 8, 27, 3])
-
- # load image processor (only resize + rescale)
- image_processor = GLPNImageProcessor()
-
- # prepare image
- image = prepare_img()
- pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
-
- logger.info("Converting model...")
-
- # load original state dict
- state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))
-
- # rename keys
- state_dict = rename_keys(state_dict)
-
- # key and value matrices need special treatment
- read_in_k_v(state_dict, config)
-
- # create HuggingFace model and load state dict
- model = GLPNForDepthEstimation(config)
- model.load_state_dict(state_dict)
- model.eval()
-
- # forward pass
- outputs = model(pixel_values)
- predicted_depth = outputs.predicted_depth
-
- # verify output
- if model_name is not None:
- if "nyu" in model_name:
- expected_slice = torch.tensor(
- [[4.4147, 4.0873, 4.0673], [3.7890, 3.2881, 3.1525], [3.7674, 3.5423, 3.4913]]
- )
- elif "kitti" in model_name:
- expected_slice = torch.tensor(
- [[3.4291, 2.7865, 2.5151], [3.2841, 2.7021, 2.3502], [3.1147, 2.4625, 2.2481]]
- )
- else:
- raise ValueError(f"Unknown model name: {model_name}")
-
- expected_shape = torch.Size([1, 480, 640])
-
- assert predicted_depth.shape == expected_shape
- assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4)
- print("Looks ok!")
-
- # finally, push to hub if required
- if push_to_hub:
- logger.info("Pushing model and image processor to the hub...")
- model.push_to_hub(
- repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
- organization="nielsr",
- commit_message="Add model",
- use_temp_dir=True,
- )
- image_processor.push_to_hub(
- repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
- organization="nielsr",
- commit_message="Add image processor",
- use_temp_dir=True,
- )
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- "--checkpoint_path",
- default=None,
- type=str,
- help="Path to the original PyTorch checkpoint (.pth file).",
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
- )
- parser.add_argument(
- "--push_to_hub", action="store_true", help="Whether to upload the model to the HuggingFace hub."
- )
- parser.add_argument(
- "--model_name",
- default="glpn-kitti",
- type=str,
- help="Name of the model in case you're pushing to the hub.",
- )
- args = parser.parse_args()
- convert_glpn_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name)
diff --git a/transformers/models/glpn/feature_extraction_glpn.py b/transformers/models/glpn/feature_extraction_glpn.py
deleted file mode 100644
index 314268225d2af41f3cc6af55af4e21aebe087b60..0000000000000000000000000000000000000000
--- a/transformers/models/glpn/feature_extraction_glpn.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Feature extractor class for GLPN."""
-
-import warnings
-
-from ...utils import logging
-from .image_processing_glpn import GLPNImageProcessor
-
-
-logger = logging.get_logger(__name__)
-
-
-class GLPNFeatureExtractor(GLPNImageProcessor):
- def __init__(self, *args, **kwargs) -> None:
- warnings.warn(
- "The class GLPNFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
- " use GLPNImageProcessor instead.",
- FutureWarning,
- )
- super().__init__(*args, **kwargs)
diff --git a/transformers/models/glpn/image_processing_glpn.py b/transformers/models/glpn/image_processing_glpn.py
deleted file mode 100644
index 7577b4eeb3d0c20b9d023bc488f8bf3c6bb39fdd..0000000000000000000000000000000000000000
--- a/transformers/models/glpn/image_processing_glpn.py
+++ /dev/null
@@ -1,233 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for GLPN."""
-
-from typing import List, Optional, Union
-
-import numpy as np
-import PIL.Image
-
-from ...image_processing_utils import BaseImageProcessor, BatchFeature
-from ...image_transforms import resize, to_channel_dimension_format
-from ...image_utils import (
- ChannelDimension,
- PILImageResampling,
- get_image_size,
- infer_channel_dimension_format,
- is_scaled_image,
- make_list_of_images,
- to_numpy_array,
- valid_images,
- validate_kwargs,
- validate_preprocess_arguments,
-)
-from ...utils import TensorType, logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class GLPNImageProcessor(BaseImageProcessor):
- r"""
- Constructs a GLPN image processor.
-
- Args:
- do_resize (`bool`, *optional*, defaults to `True`):
- Whether to resize the image's (height, width) dimensions, rounding them down to the closest multiple of
- `size_divisor`. Can be overridden by `do_resize` in `preprocess`.
- size_divisor (`int`, *optional*, defaults to 32):
- When `do_resize` is `True`, images are resized so their height and width are rounded down to the closest
- multiple of `size_divisor`. Can be overridden by `size_divisor` in `preprocess`.
- resample (`PIL.Image` resampling filter, *optional*, defaults to `Resampling.BILINEAR`):
- Resampling filter to use if resizing the image. Can be overridden by `resample` in `preprocess`.
- do_rescale (`bool`, *optional*, defaults to `True`):
- Whether or not to apply the scaling factor (to make pixel values floats between 0. and 1.). Can be
- overridden by `do_rescale` in `preprocess`.
- """
-
- model_input_names = ["pixel_values"]
-
- def __init__(
- self,
- do_resize: bool = True,
- size_divisor: int = 32,
- resample=PILImageResampling.BILINEAR,
- do_rescale: bool = True,
- **kwargs,
- ) -> None:
- self.do_resize = do_resize
- self.do_rescale = do_rescale
- self.size_divisor = size_divisor
- self.resample = resample
- super().__init__(**kwargs)
- self._valid_processor_keys = [
- "images",
- "do_resize",
- "size_divisor",
- "resample",
- "do_rescale",
- "return_tensors",
- "data_format",
- "input_data_format",
- ]
-
- def resize(
- self,
- image: np.ndarray,
- size_divisor: int,
- resample: PILImageResampling = PILImageResampling.BILINEAR,
- data_format: Optional[ChannelDimension] = None,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> np.ndarray:
- """
- Resize the image, rounding the (height, width) dimensions down to the closest multiple of size_divisor.
-
- If the image is of dimension (3, 260, 170) and size_divisor is 32, the image will be resized to (3, 256, 160).
-
- Args:
- image (`np.ndarray`):
- The image to resize.
- size_divisor (`int`):
- The image is resized so its height and width are rounded down to the closest multiple of
- `size_divisor`.
- resample:
- `PIL.Image` resampling filter to use when resizing the image e.g. `PILImageResampling.BILINEAR`.
- data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the output image. If `None`, the channel dimension format of the input
- image is used. Can be one of:
- - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format of the input image. If not set, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
-
- Returns:
- `np.ndarray`: The resized image.
- """
- height, width = get_image_size(image, channel_dim=input_data_format)
- # Rounds the height and width down to the closest multiple of size_divisor
- new_h = height // size_divisor * size_divisor
- new_w = width // size_divisor * size_divisor
- image = resize(
- image,
- (new_h, new_w),
- resample=resample,
- data_format=data_format,
- input_data_format=input_data_format,
- **kwargs,
- )
- return image
-
- def preprocess(
- self,
- images: Union["PIL.Image.Image", TensorType, List["PIL.Image.Image"], List[TensorType]],
- do_resize: Optional[bool] = None,
- size_divisor: Optional[int] = None,
- resample=None,
- do_rescale: Optional[bool] = None,
- return_tensors: Optional[Union[TensorType, str]] = None,
- data_format: ChannelDimension = ChannelDimension.FIRST,
- input_data_format: Optional[Union[str, ChannelDimension]] = None,
- **kwargs,
- ) -> BatchFeature:
- """
- Preprocess the given images.
-
- Args:
- images (`PIL.Image.Image` or `TensorType` or `List[np.ndarray]` or `List[TensorType]`):
- Images to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
- passing in images with pixel values between 0 and 1, set `do_normalize=False`.
- do_resize (`bool`, *optional*, defaults to `self.do_resize`):
- Whether to resize the input such that the (height, width) dimensions are a multiple of `size_divisor`.
- size_divisor (`int`, *optional*, defaults to `self.size_divisor`):
- When `do_resize` is `True`, images are resized so their height and width are rounded down to the
- closest multiple of `size_divisor`.
- resample (`PIL.Image` resampling filter, *optional*, defaults to `self.resample`):
- `PIL.Image` resampling filter to use if resizing the image e.g. `PILImageResampling.BILINEAR`. Only has
- an effect if `do_resize` is set to `True`.
- do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
- Whether or not to apply the scaling factor (to make pixel values floats between 0. and 1.).
- return_tensors (`str` or `TensorType`, *optional*):
- The type of tensors to return. Can be one of:
- - `None`: Return a list of `np.ndarray`.
- - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
- - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
- - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
- - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
- The channel dimension format for the output image. Can be one of:
- - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- input_data_format (`ChannelDimension` or `str`, *optional*):
- The channel dimension format for the input image. If unset, the channel dimension format is inferred
- from the input image. Can be one of:
- - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
- - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
- - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
- """
- do_resize = do_resize if do_resize is not None else self.do_resize
- do_rescale = do_rescale if do_rescale is not None else self.do_rescale
- size_divisor = size_divisor if size_divisor is not None else self.size_divisor
- resample = resample if resample is not None else self.resample
-
- images = make_list_of_images(images)
-
- validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
-
- if not valid_images(images):
- raise ValueError(
- "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
- "torch.Tensor, tf.Tensor or jax.ndarray."
- )
-
- # Here, the rescale() method uses a constant rescale_factor. It does not need to be validated
- # with a rescale_factor.
- validate_preprocess_arguments(
- do_resize=do_resize,
- size=size_divisor, # Here, size_divisor is used as a parameter for optimal resizing instead of size.
- resample=resample,
- )
-
- # All transformations expect numpy arrays.
- images = [to_numpy_array(img) for img in images]
-
- if is_scaled_image(images[0]) and do_rescale:
- logger.warning_once(
- "It looks like you are trying to rescale already rescaled images. If the input"
- " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
- )
-
- if input_data_format is None:
- # We assume that all images have the same channel dimension format.
- input_data_format = infer_channel_dimension_format(images[0])
-
- if do_resize:
- images = [
- self.resize(image, size_divisor=size_divisor, resample=resample, input_data_format=input_data_format)
- for image in images
- ]
-
- if do_rescale:
- images = [self.rescale(image, scale=1 / 255, input_data_format=input_data_format) for image in images]
-
- images = [
- to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
- ]
-
- data = {"pixel_values": images}
- return BatchFeature(data=data, tensor_type=return_tensors)
diff --git a/transformers/models/glpn/modeling_glpn.py b/transformers/models/glpn/modeling_glpn.py
deleted file mode 100644
index e5d30b62720c9d72eaadb3d5125d55f6ac4664ab..0000000000000000000000000000000000000000
--- a/transformers/models/glpn/modeling_glpn.py
+++ /dev/null
@@ -1,778 +0,0 @@
-# coding=utf-8
-# Copyright 2022 KAIST and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch GLPN model."""
-
-
-import math
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from ...activations import ACT2FN
-from ...modeling_outputs import BaseModelOutput, DepthEstimatorOutput
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_glpn import GLPNConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-# General docstring
-_CONFIG_FOR_DOC = "GLPNConfig"
-
-# Base docstring
-_CHECKPOINT_FOR_DOC = "vinvino02/glpn-kitti"
-_EXPECTED_OUTPUT_SHAPE = [1, 512, 15, 20]
-
-
-from ..deprecated._archive_maps import GLPN_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.beit.modeling_beit.drop_path
-def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
- """
- Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
- Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
- however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
- See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
- layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
- argument.
- """
- if drop_prob == 0.0 or not training:
- return input
- keep_prob = 1 - drop_prob
- shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
- random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
- random_tensor.floor_() # binarize
- output = input.div(keep_prob) * random_tensor
- return output
-
-
-# Copied from transformers.models.segformer.modeling_segformer.SegformerDropPath
-class GLPNDropPath(nn.Module):
- """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
-
- def __init__(self, drop_prob: Optional[float] = None) -> None:
- super().__init__()
- self.drop_prob = drop_prob
-
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
- return drop_path(hidden_states, self.drop_prob, self.training)
-
- def extra_repr(self) -> str:
- return "p={}".format(self.drop_prob)
-
-
-# Copied from transformers.models.segformer.modeling_segformer.SegformerOverlapPatchEmbeddings
-class GLPNOverlapPatchEmbeddings(nn.Module):
- """Construct the overlapping patch embeddings."""
-
- def __init__(self, patch_size, stride, num_channels, hidden_size):
- super().__init__()
- self.proj = nn.Conv2d(
- num_channels,
- hidden_size,
- kernel_size=patch_size,
- stride=stride,
- padding=patch_size // 2,
- )
-
- self.layer_norm = nn.LayerNorm(hidden_size)
-
- def forward(self, pixel_values):
- embeddings = self.proj(pixel_values)
- _, _, height, width = embeddings.shape
- # (batch_size, num_channels, height, width) -> (batch_size, num_channels, height*width) -> (batch_size, height*width, num_channels)
- # this can be fed to a Transformer layer
- embeddings = embeddings.flatten(2).transpose(1, 2)
- embeddings = self.layer_norm(embeddings)
- return embeddings, height, width
-
-
-# Copied from transformers.models.segformer.modeling_segformer.SegformerEfficientSelfAttention
-class GLPNEfficientSelfAttention(nn.Module):
- """SegFormer's efficient self-attention mechanism. Employs the sequence reduction process introduced in the [PvT
- paper](https://arxiv.org/abs/2102.12122)."""
-
- def __init__(self, config, hidden_size, num_attention_heads, sequence_reduction_ratio):
- super().__init__()
- self.hidden_size = hidden_size
- self.num_attention_heads = num_attention_heads
-
- if self.hidden_size % self.num_attention_heads != 0:
- raise ValueError(
- f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
- f"heads ({self.num_attention_heads})"
- )
-
- self.attention_head_size = int(self.hidden_size / self.num_attention_heads)
- self.all_head_size = self.num_attention_heads * self.attention_head_size
-
- self.query = nn.Linear(self.hidden_size, self.all_head_size)
- self.key = nn.Linear(self.hidden_size, self.all_head_size)
- self.value = nn.Linear(self.hidden_size, self.all_head_size)
-
- self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
- self.sr_ratio = sequence_reduction_ratio
- if sequence_reduction_ratio > 1:
- self.sr = nn.Conv2d(
- hidden_size, hidden_size, kernel_size=sequence_reduction_ratio, stride=sequence_reduction_ratio
- )
- self.layer_norm = nn.LayerNorm(hidden_size)
-
- def transpose_for_scores(self, hidden_states):
- new_shape = hidden_states.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
- hidden_states = hidden_states.view(new_shape)
- return hidden_states.permute(0, 2, 1, 3)
-
- def forward(
- self,
- hidden_states,
- height,
- width,
- output_attentions=False,
- ):
- query_layer = self.transpose_for_scores(self.query(hidden_states))
-
- if self.sr_ratio > 1:
- batch_size, seq_len, num_channels = hidden_states.shape
- # Reshape to (batch_size, num_channels, height, width)
- hidden_states = hidden_states.permute(0, 2, 1).reshape(batch_size, num_channels, height, width)
- # Apply sequence reduction
- hidden_states = self.sr(hidden_states)
- # Reshape back to (batch_size, seq_len, num_channels)
- hidden_states = hidden_states.reshape(batch_size, num_channels, -1).permute(0, 2, 1)
- hidden_states = self.layer_norm(hidden_states)
-
- key_layer = self.transpose_for_scores(self.key(hidden_states))
- value_layer = self.transpose_for_scores(self.value(hidden_states))
-
- # Take the dot product between "query" and "key" to get the raw attention scores.
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-
- attention_scores = attention_scores / math.sqrt(self.attention_head_size)
-
- # Normalize the attention scores to probabilities.
- attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
- # This is actually dropping out entire tokens to attend to, which might
- # seem a bit unusual, but is taken from the original Transformer paper.
- attention_probs = self.dropout(attention_probs)
-
- context_layer = torch.matmul(attention_probs, value_layer)
-
- context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
- new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
- context_layer = context_layer.view(new_context_layer_shape)
-
- outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
-
- return outputs
-
-
-# Copied from transformers.models.segformer.modeling_segformer.SegformerSelfOutput
-class GLPNSelfOutput(nn.Module):
- def __init__(self, config, hidden_size):
- super().__init__()
- self.dense = nn.Linear(hidden_size, hidden_size)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states, input_tensor):
- hidden_states = self.dense(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.segformer.modeling_segformer.SegformerAttention with Segformer->GLPN
-class GLPNAttention(nn.Module):
- def __init__(self, config, hidden_size, num_attention_heads, sequence_reduction_ratio):
- super().__init__()
- self.self = GLPNEfficientSelfAttention(
- config=config,
- hidden_size=hidden_size,
- num_attention_heads=num_attention_heads,
- sequence_reduction_ratio=sequence_reduction_ratio,
- )
- self.output = GLPNSelfOutput(config, hidden_size=hidden_size)
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(
- heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
- )
-
- # Prune linear layers
- self.self.query = prune_linear_layer(self.self.query, index)
- self.self.key = prune_linear_layer(self.self.key, index)
- self.self.value = prune_linear_layer(self.self.value, index)
- self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
-
- # Update hyper params and store pruned heads
- self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
- self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def forward(self, hidden_states, height, width, output_attentions=False):
- self_outputs = self.self(hidden_states, height, width, output_attentions)
-
- attention_output = self.output(self_outputs[0], hidden_states)
- outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
- return outputs
-
-
-# Copied from transformers.models.segformer.modeling_segformer.SegformerDWConv
-class GLPNDWConv(nn.Module):
- def __init__(self, dim=768):
- super().__init__()
- self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
-
- def forward(self, hidden_states, height, width):
- batch_size, seq_len, num_channels = hidden_states.shape
- hidden_states = hidden_states.transpose(1, 2).view(batch_size, num_channels, height, width)
- hidden_states = self.dwconv(hidden_states)
- hidden_states = hidden_states.flatten(2).transpose(1, 2)
-
- return hidden_states
-
-
-# Copied from transformers.models.segformer.modeling_segformer.SegformerMixFFN with Segformer->GLPN
-class GLPNMixFFN(nn.Module):
- def __init__(self, config, in_features, hidden_features=None, out_features=None):
- super().__init__()
- out_features = out_features or in_features
- self.dense1 = nn.Linear(in_features, hidden_features)
- self.dwconv = GLPNDWConv(hidden_features)
- if isinstance(config.hidden_act, str):
- self.intermediate_act_fn = ACT2FN[config.hidden_act]
- else:
- self.intermediate_act_fn = config.hidden_act
- self.dense2 = nn.Linear(hidden_features, out_features)
- self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
- def forward(self, hidden_states, height, width):
- hidden_states = self.dense1(hidden_states)
- hidden_states = self.dwconv(hidden_states, height, width)
- hidden_states = self.intermediate_act_fn(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.dense2(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.segformer.modeling_segformer.SegformerLayer with Segformer->GLPN
-class GLPNLayer(nn.Module):
- """This corresponds to the Block class in the original implementation."""
-
- def __init__(self, config, hidden_size, num_attention_heads, drop_path, sequence_reduction_ratio, mlp_ratio):
- super().__init__()
- self.layer_norm_1 = nn.LayerNorm(hidden_size)
- self.attention = GLPNAttention(
- config,
- hidden_size=hidden_size,
- num_attention_heads=num_attention_heads,
- sequence_reduction_ratio=sequence_reduction_ratio,
- )
- self.drop_path = GLPNDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
- self.layer_norm_2 = nn.LayerNorm(hidden_size)
- mlp_hidden_size = int(hidden_size * mlp_ratio)
- self.mlp = GLPNMixFFN(config, in_features=hidden_size, hidden_features=mlp_hidden_size)
-
- def forward(self, hidden_states, height, width, output_attentions=False):
- self_attention_outputs = self.attention(
- self.layer_norm_1(hidden_states), # in GLPN, layernorm is applied before self-attention
- height,
- width,
- output_attentions=output_attentions,
- )
-
- attention_output = self_attention_outputs[0]
- outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
-
- # first residual connection (with stochastic depth)
- attention_output = self.drop_path(attention_output)
- hidden_states = attention_output + hidden_states
-
- mlp_output = self.mlp(self.layer_norm_2(hidden_states), height, width)
-
- # second residual connection (with stochastic depth)
- mlp_output = self.drop_path(mlp_output)
- layer_output = mlp_output + hidden_states
-
- outputs = (layer_output,) + outputs
-
- return outputs
-
-
-class GLPNEncoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
-
- # stochastic depth decay rule
- dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
-
- # patch embeddings
- embeddings = []
- for i in range(config.num_encoder_blocks):
- embeddings.append(
- GLPNOverlapPatchEmbeddings(
- patch_size=config.patch_sizes[i],
- stride=config.strides[i],
- num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1],
- hidden_size=config.hidden_sizes[i],
- )
- )
- self.patch_embeddings = nn.ModuleList(embeddings)
-
- # Transformer blocks
- blocks = []
- cur = 0
- for i in range(config.num_encoder_blocks):
- # each block consists of layers
- layers = []
- if i != 0:
- cur += config.depths[i - 1]
- for j in range(config.depths[i]):
- layers.append(
- GLPNLayer(
- config,
- hidden_size=config.hidden_sizes[i],
- num_attention_heads=config.num_attention_heads[i],
- drop_path=dpr[cur + j],
- sequence_reduction_ratio=config.sr_ratios[i],
- mlp_ratio=config.mlp_ratios[i],
- )
- )
- blocks.append(nn.ModuleList(layers))
-
- self.block = nn.ModuleList(blocks)
-
- # Layer norms
- self.layer_norm = nn.ModuleList(
- [nn.LayerNorm(config.hidden_sizes[i]) for i in range(config.num_encoder_blocks)]
- )
-
- def forward(
- self,
- pixel_values,
- output_attentions=False,
- output_hidden_states=False,
- return_dict=True,
- ):
- all_hidden_states = () if output_hidden_states else None
- all_self_attentions = () if output_attentions else None
-
- batch_size = pixel_values.shape[0]
-
- hidden_states = pixel_values
- for idx, x in enumerate(zip(self.patch_embeddings, self.block, self.layer_norm)):
- embedding_layer, block_layer, norm_layer = x
- # first, obtain patch embeddings
- hidden_states, height, width = embedding_layer(hidden_states)
- # second, send embeddings through blocks
- for i, blk in enumerate(block_layer):
- layer_outputs = blk(hidden_states, height, width, output_attentions)
- hidden_states = layer_outputs[0]
- if output_attentions:
- all_self_attentions = all_self_attentions + (layer_outputs[1],)
- # third, apply layer norm
- hidden_states = norm_layer(hidden_states)
- # fourth, optionally reshape back to (batch_size, num_channels, height, width)
- hidden_states = hidden_states.reshape(batch_size, height, width, -1).permute(0, 3, 1, 2).contiguous()
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
- return BaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-class GLPNPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GLPNConfig
- base_model_prefix = "glpn"
- main_input_name = "pixel_values"
-
- # Copied from transformers.models.segformer.modeling_segformer.SegformerPreTrainedModel._init_weights
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-GLPN_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`GLPNConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GLPN_INPUTS_DOCSTRING = r"""
-
- Args:
- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
- Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
- [`AutoImageProcessor`]. See [`GLPNImageProcessor.__call__`] for details.
-
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare GLPN encoder (Mix-Transformer) outputting raw hidden-states without any specific head on top.",
- GLPN_START_DOCSTRING,
-)
-class GLPNModel(GLPNPreTrainedModel):
- # Copied from transformers.models.segformer.modeling_segformer.SegformerModel.__init__ with Segformer->GLPN
- def __init__(self, config):
- super().__init__(config)
- self.config = config
-
- # hierarchical Transformer encoder
- self.encoder = GLPNEncoder(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
- class PreTrainedModel
- """
- for layer, heads in heads_to_prune.items():
- self.encoder.layer[layer].attention.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(GLPN_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutput,
- config_class=_CONFIG_FOR_DOC,
- modality="vision",
- expected_output=_EXPECTED_OUTPUT_SHAPE,
- )
- # Copied from transformers.models.segformer.modeling_segformer.SegformerModel.forward
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutput]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- encoder_outputs = self.encoder(
- pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- sequence_output = encoder_outputs[0]
-
- if not return_dict:
- return (sequence_output,) + encoder_outputs[1:]
-
- return BaseModelOutput(
- last_hidden_state=sequence_output,
- hidden_states=encoder_outputs.hidden_states,
- attentions=encoder_outputs.attentions,
- )
-
-
-class GLPNSelectiveFeatureFusion(nn.Module):
- """
- Selective Feature Fusion module, as explained in the [paper](https://arxiv.org/abs/2201.07436) (section 3.4). This
- module adaptively selects and integrates local and global features by attaining an attention map for each feature.
- """
-
- def __init__(self, in_channel=64):
- super().__init__()
-
- self.convolutional_layer1 = nn.Sequential(
- nn.Conv2d(in_channels=int(in_channel * 2), out_channels=in_channel, kernel_size=3, stride=1, padding=1),
- nn.BatchNorm2d(in_channel),
- nn.ReLU(),
- )
-
- self.convolutional_layer2 = nn.Sequential(
- nn.Conv2d(in_channels=in_channel, out_channels=int(in_channel / 2), kernel_size=3, stride=1, padding=1),
- nn.BatchNorm2d(int(in_channel / 2)),
- nn.ReLU(),
- )
-
- self.convolutional_layer3 = nn.Conv2d(
- in_channels=int(in_channel / 2), out_channels=2, kernel_size=3, stride=1, padding=1
- )
-
- self.sigmoid = nn.Sigmoid()
-
- def forward(self, local_features, global_features):
- # concatenate features along the channel dimension
- features = torch.cat((local_features, global_features), dim=1)
- # pass through convolutional layers
- features = self.convolutional_layer1(features)
- features = self.convolutional_layer2(features)
- features = self.convolutional_layer3(features)
- # apply sigmoid to get two-channel attention map
- attn = self.sigmoid(features)
- # construct hybrid features by adding element-wise
- hybrid_features = local_features * attn[:, 0, :, :].unsqueeze(1) + global_features * attn[
- :, 1, :, :
- ].unsqueeze(1)
-
- return hybrid_features
-
-
-class GLPNDecoderStage(nn.Module):
- def __init__(self, in_channels, out_channels):
- super().__init__()
- should_skip = in_channels == out_channels
- self.convolution = nn.Conv2d(in_channels, out_channels, kernel_size=1) if not should_skip else nn.Identity()
- self.fusion = GLPNSelectiveFeatureFusion(out_channels)
- self.upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
-
- def forward(self, hidden_state, residual=None):
- hidden_state = self.convolution(hidden_state)
- if residual is not None:
- hidden_state = self.fusion(hidden_state, residual)
- hidden_state = self.upsample(hidden_state)
-
- return hidden_state
-
- hidden_state = self.upsample(hidden_state)
- return hidden_state
-
-
-class GLPNDecoder(nn.Module):
- def __init__(self, config):
- super().__init__()
- # we use features from end -> start
- reserved_hidden_sizes = config.hidden_sizes[::-1]
- out_channels = config.decoder_hidden_size
-
- self.stages = nn.ModuleList(
- [GLPNDecoderStage(hidden_size, out_channels) for hidden_size in reserved_hidden_sizes]
- )
- # don't fuse in first stage
- self.stages[0].fusion = None
-
- self.final_upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
-
- def forward(self, hidden_states: List[torch.Tensor]) -> List[torch.Tensor]:
- stage_hidden_states = []
- stage_hidden_state = None
- for hidden_state, stage in zip(hidden_states[::-1], self.stages):
- stage_hidden_state = stage(hidden_state, stage_hidden_state)
- stage_hidden_states.append(stage_hidden_state)
-
- stage_hidden_states[-1] = self.final_upsample(stage_hidden_state)
-
- return stage_hidden_states
-
-
-class SiLogLoss(nn.Module):
- r"""
- Implements the Scale-invariant log scale loss [Eigen et al., 2014](https://arxiv.org/abs/1406.2283).
-
- $$L=\frac{1}{n} \sum_{i} d_{i}^{2}-\frac{1}{2 n^{2}}\left(\sum_{i} d_{i}^{2}\right)$$ where $d_{i}=\log y_{i}-\log
- y_{i}^{*}$.
-
- """
-
- def __init__(self, lambd=0.5):
- super().__init__()
- self.lambd = lambd
-
- def forward(self, pred, target):
- valid_mask = (target > 0).detach()
- diff_log = torch.log(target[valid_mask]) - torch.log(pred[valid_mask])
- loss = torch.sqrt(torch.pow(diff_log, 2).mean() - self.lambd * torch.pow(diff_log.mean(), 2))
-
- return loss
-
-
-class GLPNDepthEstimationHead(nn.Module):
- def __init__(self, config):
- super().__init__()
-
- self.config = config
-
- channels = config.decoder_hidden_size
- self.head = nn.Sequential(
- nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1),
- nn.ReLU(inplace=False),
- nn.Conv2d(channels, 1, kernel_size=3, stride=1, padding=1),
- )
-
- def forward(self, hidden_states: List[torch.Tensor]) -> torch.Tensor:
- # use last features of the decoder
- hidden_states = hidden_states[self.config.head_in_index]
-
- hidden_states = self.head(hidden_states)
-
- predicted_depth = torch.sigmoid(hidden_states) * self.config.max_depth
- predicted_depth = predicted_depth.squeeze(dim=1)
-
- return predicted_depth
-
-
-@add_start_docstrings(
- """GLPN Model transformer with a lightweight depth estimation head on top e.g. for KITTI, NYUv2.""",
- GLPN_START_DOCSTRING,
-)
-class GLPNForDepthEstimation(GLPNPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.glpn = GLPNModel(config)
- self.decoder = GLPNDecoder(config)
- self.head = GLPNDepthEstimationHead(config)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GLPN_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=DepthEstimatorOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- pixel_values: torch.FloatTensor,
- labels: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], DepthEstimatorOutput]:
- r"""
- labels (`torch.FloatTensor` of shape `(batch_size, height, width)`, *optional*):
- Ground truth depth estimation maps for computing the loss.
-
- Returns:
-
- Examples:
-
- ```python
- >>> from transformers import AutoImageProcessor, GLPNForDepthEstimation
- >>> import torch
- >>> import numpy as np
- >>> from PIL import Image
- >>> import requests
-
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> image_processor = AutoImageProcessor.from_pretrained("vinvino02/glpn-kitti")
- >>> model = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-kitti")
-
- >>> # prepare image for the model
- >>> inputs = image_processor(images=image, return_tensors="pt")
-
- >>> with torch.no_grad():
- ... outputs = model(**inputs)
- ... predicted_depth = outputs.predicted_depth
-
- >>> # interpolate to original size
- >>> prediction = torch.nn.functional.interpolate(
- ... predicted_depth.unsqueeze(1),
- ... size=image.size[::-1],
- ... mode="bicubic",
- ... align_corners=False,
- ... )
-
- >>> # visualize the prediction
- >>> output = prediction.squeeze().cpu().numpy()
- >>> formatted = (output * 255 / np.max(output)).astype("uint8")
- >>> depth = Image.fromarray(formatted)
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
-
- outputs = self.glpn(
- pixel_values,
- output_attentions=output_attentions,
- output_hidden_states=True, # we need the intermediate hidden states
- return_dict=return_dict,
- )
-
- hidden_states = outputs.hidden_states if return_dict else outputs[1]
-
- out = self.decoder(hidden_states)
- predicted_depth = self.head(out)
-
- loss = None
- if labels is not None:
- loss_fct = SiLogLoss()
- loss = loss_fct(predicted_depth, labels)
-
- if not return_dict:
- if output_hidden_states:
- output = (predicted_depth,) + outputs[1:]
- else:
- output = (predicted_depth,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return DepthEstimatorOutput(
- loss=loss,
- predicted_depth=predicted_depth,
- hidden_states=outputs.hidden_states if output_hidden_states else None,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/gpt2/__init__.py b/transformers/models/gpt2/__init__.py
deleted file mode 100644
index e99658ac1e885e1e80aa4ea9e52b7050f74abd1e..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/__init__.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# Copyright 2020 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_flax_available,
- is_keras_nlp_available,
- is_tensorflow_text_available,
- is_tf_available,
- is_tokenizers_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_gpt2": ["GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPT2Config", "GPT2OnnxConfig"],
- "tokenization_gpt2": ["GPT2Tokenizer"],
-}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_gpt2_fast"] = ["GPT2TokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_gpt2"] = [
- "GPT2_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GPT2DoubleHeadsModel",
- "GPT2ForQuestionAnswering",
- "GPT2ForSequenceClassification",
- "GPT2ForTokenClassification",
- "GPT2LMHeadModel",
- "GPT2Model",
- "GPT2PreTrainedModel",
- "load_tf_weights_in_gpt2",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_gpt2"] = [
- "TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST",
- "TFGPT2DoubleHeadsModel",
- "TFGPT2ForSequenceClassification",
- "TFGPT2LMHeadModel",
- "TFGPT2MainLayer",
- "TFGPT2Model",
- "TFGPT2PreTrainedModel",
- ]
-
-try:
- if not is_keras_nlp_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_gpt2_tf"] = ["TFGPT2Tokenizer"]
-
-try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flax_gpt2"] = ["FlaxGPT2LMHeadModel", "FlaxGPT2Model", "FlaxGPT2PreTrainedModel"]
-
-if TYPE_CHECKING:
- from .configuration_gpt2 import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPT2Config, GPT2OnnxConfig
- from .tokenization_gpt2 import GPT2Tokenizer
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_gpt2_fast import GPT2TokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_gpt2 import (
- GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
- GPT2DoubleHeadsModel,
- GPT2ForQuestionAnswering,
- GPT2ForSequenceClassification,
- GPT2ForTokenClassification,
- GPT2LMHeadModel,
- GPT2Model,
- GPT2PreTrainedModel,
- load_tf_weights_in_gpt2,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_gpt2 import (
- TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
- TFGPT2DoubleHeadsModel,
- TFGPT2ForSequenceClassification,
- TFGPT2LMHeadModel,
- TFGPT2MainLayer,
- TFGPT2Model,
- TFGPT2PreTrainedModel,
- )
-
- try:
- if not is_keras_nlp_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_gpt2_tf import TFGPT2Tokenizer
-
- try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flax_gpt2 import FlaxGPT2LMHeadModel, FlaxGPT2Model, FlaxGPT2PreTrainedModel
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gpt2/__pycache__/__init__.cpython-310.pyc b/transformers/models/gpt2/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 487aa779ecc00d60fe7613730e487c091ab56490..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/__pycache__/configuration_gpt2.cpython-310.pyc b/transformers/models/gpt2/__pycache__/configuration_gpt2.cpython-310.pyc
deleted file mode 100644
index ed2da891e95df42e84b7e446cc60455c50d3e5c2..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/configuration_gpt2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/__pycache__/convert_gpt2_original_tf_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/gpt2/__pycache__/convert_gpt2_original_tf_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 39a9fb0eaf47a32df1450035538c9278294348d7..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/convert_gpt2_original_tf_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/__pycache__/modeling_flax_gpt2.cpython-310.pyc b/transformers/models/gpt2/__pycache__/modeling_flax_gpt2.cpython-310.pyc
deleted file mode 100644
index 3d7b438341c5e212177e867414be0a41c9fc8a4e..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/modeling_flax_gpt2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/__pycache__/modeling_gpt2.cpython-310.pyc b/transformers/models/gpt2/__pycache__/modeling_gpt2.cpython-310.pyc
deleted file mode 100644
index ff254f2564f59408a54f3e3dbafdff58d66b2657..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/modeling_gpt2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/__pycache__/modeling_tf_gpt2.cpython-310.pyc b/transformers/models/gpt2/__pycache__/modeling_tf_gpt2.cpython-310.pyc
deleted file mode 100644
index fbf188affca1b9f91c84f0bbf47aeb840ac47422..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/modeling_tf_gpt2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/__pycache__/tokenization_gpt2.cpython-310.pyc b/transformers/models/gpt2/__pycache__/tokenization_gpt2.cpython-310.pyc
deleted file mode 100644
index 3b399f0ca043ff9d736c1cbbfc803050ff8d258f..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/tokenization_gpt2.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/__pycache__/tokenization_gpt2_fast.cpython-310.pyc b/transformers/models/gpt2/__pycache__/tokenization_gpt2_fast.cpython-310.pyc
deleted file mode 100644
index b6ae4d5f80991499fa440a73088a5b6dea153d81..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/tokenization_gpt2_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/__pycache__/tokenization_gpt2_tf.cpython-310.pyc b/transformers/models/gpt2/__pycache__/tokenization_gpt2_tf.cpython-310.pyc
deleted file mode 100644
index 4b8729d69f0868aab375b663a52164926ca71ac9..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt2/__pycache__/tokenization_gpt2_tf.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt2/configuration_gpt2.py b/transformers/models/gpt2/configuration_gpt2.py
deleted file mode 100644
index 45495c0012fdd811342cf091a8f2df2cd568d7aa..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/configuration_gpt2.py
+++ /dev/null
@@ -1,272 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" OpenAI GPT-2 configuration"""
-from collections import OrderedDict
-from typing import Any, List, Mapping, Optional
-
-from ... import PreTrainedTokenizer, TensorType, is_torch_available
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfigWithPast, PatchingSpec
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GPT2Config(PretrainedConfig):
- """
- This is the configuration class to store the configuration of a [`GPT2Model`] or a [`TFGPT2Model`]. It is used to
- instantiate a GPT-2 model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the GPT-2
- [openai-community/gpt2](https://huggingface.co/openai-community/gpt2) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 50257):
- Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`GPT2Model`] or [`TFGPT2Model`].
- n_positions (`int`, *optional*, defaults to 1024):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- n_embd (`int`, *optional*, defaults to 768):
- Dimensionality of the embeddings and hidden states.
- n_layer (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- n_head (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- n_inner (`int`, *optional*):
- Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
- activation_function (`str`, *optional*, defaults to `"gelu_new"`):
- Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new"]`.
- resid_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- embd_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the embeddings.
- attn_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention.
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
- The epsilon to use in the layer normalization layers.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- summary_type (`string`, *optional*, defaults to `"cls_index"`):
- Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and
- [`TFGPT2DoubleHeadsModel`].
-
- Has to be one of the following options:
-
- - `"last"`: Take the last token hidden state (like XLNet).
- - `"first"`: Take the first token hidden state (like BERT).
- - `"mean"`: Take the mean of all tokens hidden states.
- - `"cls_index"`: Supply a Tensor of classification token position (like GPT/GPT-2).
- - `"attn"`: Not implemented now, use multi-head attention.
- summary_use_proj (`bool`, *optional*, defaults to `True`):
- Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and
- [`TFGPT2DoubleHeadsModel`].
-
- Whether or not to add a projection after the vector extraction.
- summary_activation (`str`, *optional*):
- Argument used when doing sequence summary. Used in for the multiple choice head in
- [`GPT2DoubleHeadsModel`].
-
- Pass `"tanh"` for a tanh activation to the output, any other value will result in no activation.
- summary_proj_to_labels (`bool`, *optional*, defaults to `True`):
- Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and
- [`TFGPT2DoubleHeadsModel`].
-
- Whether the projection outputs should have `config.num_labels` or `config.hidden_size` classes.
- summary_first_dropout (`float`, *optional*, defaults to 0.1):
- Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and
- [`TFGPT2DoubleHeadsModel`].
-
- The dropout ratio to be used after the projection and activation.
- scale_attn_weights (`bool`, *optional*, defaults to `True`):
- Scale attention weights by dividing by sqrt(hidden_size)..
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
- bos_token_id (`int`, *optional*, defaults to 50256):
- Id of the beginning of sentence token in the vocabulary.
- eos_token_id (`int`, *optional*, defaults to 50256):
- Id of the end of sentence token in the vocabulary.
- scale_attn_by_inverse_layer_idx (`bool`, *optional*, defaults to `False`):
- Whether to additionally scale attention weights by `1 / layer_idx + 1`.
- reorder_and_upcast_attn (`bool`, *optional*, defaults to `False`):
- Whether to scale keys (K) prior to computing attention (dot-product) and upcast attention
- dot-product/softmax to float() when training with mixed precision.
-
- Example:
-
- ```python
- >>> from transformers import GPT2Config, GPT2Model
-
- >>> # Initializing a GPT2 configuration
- >>> configuration = GPT2Config()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = GPT2Model(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "gpt2"
- keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {
- "hidden_size": "n_embd",
- "max_position_embeddings": "n_positions",
- "num_attention_heads": "n_head",
- "num_hidden_layers": "n_layer",
- }
-
- def __init__(
- self,
- vocab_size=50257,
- n_positions=1024,
- n_embd=768,
- n_layer=12,
- n_head=12,
- n_inner=None,
- activation_function="gelu_new",
- resid_pdrop=0.1,
- embd_pdrop=0.1,
- attn_pdrop=0.1,
- layer_norm_epsilon=1e-5,
- initializer_range=0.02,
- summary_type="cls_index",
- summary_use_proj=True,
- summary_activation=None,
- summary_proj_to_labels=True,
- summary_first_dropout=0.1,
- scale_attn_weights=True,
- use_cache=True,
- bos_token_id=50256,
- eos_token_id=50256,
- scale_attn_by_inverse_layer_idx=False,
- reorder_and_upcast_attn=False,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.n_positions = n_positions
- self.n_embd = n_embd
- self.n_layer = n_layer
- self.n_head = n_head
- self.n_inner = n_inner
- self.activation_function = activation_function
- self.resid_pdrop = resid_pdrop
- self.embd_pdrop = embd_pdrop
- self.attn_pdrop = attn_pdrop
- self.layer_norm_epsilon = layer_norm_epsilon
- self.initializer_range = initializer_range
- self.summary_type = summary_type
- self.summary_use_proj = summary_use_proj
- self.summary_activation = summary_activation
- self.summary_first_dropout = summary_first_dropout
- self.summary_proj_to_labels = summary_proj_to_labels
- self.scale_attn_weights = scale_attn_weights
- self.use_cache = use_cache
- self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx
- self.reorder_and_upcast_attn = reorder_and_upcast_attn
-
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
-
- super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
-
-
-class GPT2OnnxConfig(OnnxConfigWithPast):
- def __init__(
- self,
- config: PretrainedConfig,
- task: str = "default",
- patching_specs: List[PatchingSpec] = None,
- use_past: bool = False,
- ):
- super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past)
- if not getattr(self._config, "pad_token_id", None):
- # TODO: how to do that better?
- self._config.pad_token_id = 0
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
- if self.use_past:
- self.fill_with_past_key_values_(common_inputs, direction="inputs")
- common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"}
- else:
- common_inputs["attention_mask"] = {0: "batch", 1: "sequence"}
-
- return common_inputs
-
- @property
- def num_layers(self) -> int:
- return self._config.n_layer
-
- @property
- def num_attention_heads(self) -> int:
- return self._config.n_head
-
- def generate_dummy_inputs(
- self,
- tokenizer: PreTrainedTokenizer,
- batch_size: int = -1,
- seq_length: int = -1,
- is_pair: bool = False,
- framework: Optional[TensorType] = None,
- ) -> Mapping[str, Any]:
- common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
- tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
- )
-
- # We need to order the input in the way they appears in the forward()
- ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]})
-
- # Need to add the past_keys
- if self.use_past:
- if not is_torch_available():
- raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
- else:
- import torch
-
- batch, seqlen = common_inputs["input_ids"].shape
- # Not using the same length for past_key_values
- past_key_values_length = seqlen + 2
- past_shape = (
- batch,
- self.num_attention_heads,
- past_key_values_length,
- self._config.hidden_size // self.num_attention_heads,
- )
- ordered_inputs["past_key_values"] = [
- (torch.zeros(past_shape), torch.zeros(past_shape)) for _ in range(self.num_layers)
- ]
-
- ordered_inputs["attention_mask"] = common_inputs["attention_mask"]
- if self.use_past:
- mask_dtype = ordered_inputs["attention_mask"].dtype
- ordered_inputs["attention_mask"] = torch.cat(
- [ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1
- )
-
- return ordered_inputs
-
- @property
- def default_onnx_opset(self) -> int:
- return 13
diff --git a/transformers/models/gpt2/convert_gpt2_original_tf_checkpoint_to_pytorch.py b/transformers/models/gpt2/convert_gpt2_original_tf_checkpoint_to_pytorch.py
deleted file mode 100644
index 066ba06503affdb8fa2ef1b2df1a4cf6539efc22..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/convert_gpt2_original_tf_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,69 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert OpenAI GPT checkpoint."""
-
-
-import argparse
-
-import torch
-
-from transformers import GPT2Config, GPT2Model, load_tf_weights_in_gpt2
-from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging
-
-
-logging.set_verbosity_info()
-
-
-def convert_gpt2_checkpoint_to_pytorch(gpt2_checkpoint_path, gpt2_config_file, pytorch_dump_folder_path):
- # Construct model
- if gpt2_config_file == "":
- config = GPT2Config()
- else:
- config = GPT2Config.from_json_file(gpt2_config_file)
- model = GPT2Model(config)
-
- # Load weights from numpy
- load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path)
-
- # Save pytorch-model
- pytorch_weights_dump_path = pytorch_dump_folder_path + "/" + WEIGHTS_NAME
- pytorch_config_dump_path = pytorch_dump_folder_path + "/" + CONFIG_NAME
- print(f"Save PyTorch model to {pytorch_weights_dump_path}")
- torch.save(model.state_dict(), pytorch_weights_dump_path)
- print(f"Save configuration file to {pytorch_config_dump_path}")
- with open(pytorch_config_dump_path, "w", encoding="utf-8") as f:
- f.write(config.to_json_string())
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--gpt2_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
- )
- parser.add_argument(
- "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
- parser.add_argument(
- "--gpt2_config_file",
- default="",
- type=str,
- help=(
- "An optional config json file corresponding to the pre-trained OpenAI model. \n"
- "This specifies the model architecture."
- ),
- )
- args = parser.parse_args()
- convert_gpt2_checkpoint_to_pytorch(args.gpt2_checkpoint_path, args.gpt2_config_file, args.pytorch_dump_folder_path)
diff --git a/transformers/models/gpt2/modeling_flax_gpt2.py b/transformers/models/gpt2/modeling_flax_gpt2.py
deleted file mode 100644
index c3ef377642a3c5b043ae36acb6443041dcf75742..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/modeling_flax_gpt2.py
+++ /dev/null
@@ -1,779 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The Google Flax Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Any, Optional, Tuple
-
-import flax.linen as nn
-import jax
-import jax.numpy as jnp
-from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
-from flax.linen import combine_masks, make_causal_mask
-from flax.linen.attention import dot_product_attention_weights
-from flax.traverse_util import flatten_dict, unflatten_dict
-from jax import lax
-
-from ...modeling_flax_outputs import (
- FlaxBaseModelOutputWithPastAndCrossAttentions,
- FlaxCausalLMOutputWithCrossAttentions,
-)
-from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_gpt2 import GPT2Config
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "openai-community/gpt2"
-_CONFIG_FOR_DOC = "GPT2Config"
-
-
-GPT2_START_DOCSTRING = r"""
-
- This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a Flax Linen
- [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a
- regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.
-
- Finally, this model supports inherent JAX features such as:
-
- - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
-
- Parameters:
- config ([`GPT2Config`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.
- dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
- The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
- `jax.numpy.bfloat16` (on TPUs).
-
- This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
- specified all the computation will be performed with the given `dtype`.
-
- **Note that this only specifies the dtype of the computation and does not influence the dtype of model
- parameters.**
-
- If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and
- [`~FlaxPreTrainedModel.to_bf16`].
-"""
-
-GPT2_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`numpy.ndarray` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length`. Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
- past_key_values (`Dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):
- Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
- auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class FlaxConv1D(nn.Module):
- features: int
- use_bias: bool = True
- dtype: Any = jnp.float32
- precision: Any = None
-
- @nn.compact
- def __call__(self, inputs):
- inputs = jnp.asarray(inputs, self.dtype)
- kernel = self.param("kernel", jax.nn.initializers.normal(stddev=0.02), (self.features, inputs.shape[-1]))
- kernel = jnp.asarray(kernel.transpose(), self.dtype)
- y = lax.dot_general(inputs, kernel, (((inputs.ndim - 1,), (0,)), ((), ())), precision=self.precision)
- if self.use_bias:
- bias = self.param("bias", jax.nn.initializers.zeros, (self.features,))
- bias = jnp.asarray(bias, self.dtype)
- y = y + bias
- return y
-
-
-class FlaxGPT2Attention(nn.Module):
- config: GPT2Config
- dtype: jnp.dtype = jnp.float32
- causal: bool = True
- is_cross_attention: bool = False
-
- def setup(self):
- config = self.config
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_heads
-
- if self.is_cross_attention:
- self.c_attn = FlaxConv1D(2 * self.embed_dim, dtype=self.dtype)
- self.q_attn = FlaxConv1D(self.embed_dim, dtype=self.dtype)
- else:
- self.c_attn = FlaxConv1D(3 * self.embed_dim, dtype=self.dtype)
- self.c_proj = FlaxConv1D(self.embed_dim, dtype=self.dtype)
-
- self.resid_dropout = nn.Dropout(rate=config.resid_pdrop)
-
- if self.causal:
- self.causal_mask = make_causal_mask(
- jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool"
- )
-
- def _split_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))
-
- def _merge_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))
-
- @nn.compact
- def _concatenate_to_cache(self, key, value, query, attention_mask):
- """
- This function takes projected key, value states from a single input token and concatenates the states to cached
- states from previous steps. This function is slighly adapted from the official Flax repository:
- https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252
- """
- # detect if we're initializing by absence of existing cache data.
- is_initialized = self.has_variable("cache", "cached_key")
- cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype)
- cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype)
- cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
-
- if is_initialized:
- *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape
- # update key, value caches with our new 1d spatial slices
- cur_index = cache_index.value
- indices = (0,) * len(batch_dims) + (cur_index, 0, 0)
- key = lax.dynamic_update_slice(cached_key.value, key, indices)
- value = lax.dynamic_update_slice(cached_value.value, value, indices)
- cached_key.value = key
- cached_value.value = value
- num_updated_cache_vectors = query.shape[1]
- cache_index.value = cache_index.value + num_updated_cache_vectors
- # causal mask for cached decoder self-attention: our single query position should only attend to those key positions that have already been generated and cached, not the remaining zero elements.
- pad_mask = jnp.broadcast_to(
- jnp.arange(max_length) < cur_index + num_updated_cache_vectors,
- tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),
- )
- attention_mask = combine_masks(pad_mask, attention_mask)
- return key, value, attention_mask
-
- def __call__(
- self,
- hidden_states,
- key_value_states: Optional[jnp.ndarray] = None,
- attention_mask=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- # if key_value_states are provided this layer is used as a cross-attention layer
- # for the decoder
- is_cross_attention = key_value_states is not None
- batch_size = hidden_states.shape[0]
-
- if not is_cross_attention:
- qkv_out = self.c_attn(hidden_states)
- query, key, value = jnp.split(qkv_out, 3, axis=2)
- else:
- q_out = self.q_attn(hidden_states)
- (query,) = jnp.split(q_out, 1, axis=2)
- kv_out = self.c_attn(key_value_states)
- key, value = jnp.split(kv_out, 2, axis=2)
-
- query = self._split_heads(query)
- key = self._split_heads(key)
- value = self._split_heads(value)
-
- query_length, key_length = query.shape[1], key.shape[1]
-
- if self.causal:
- if self.has_variable("cache", "cached_key"):
- mask_shift = self.variables["cache"]["cache_index"]
- max_decoder_length = self.variables["cache"]["cached_key"].shape[1]
- causal_mask = lax.dynamic_slice(
- self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)
- )
- else:
- causal_mask = self.causal_mask[:, :, :query_length, :key_length]
- causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])
-
- # combine masks if needed
- if attention_mask is not None and self.causal:
- attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
- attention_mask = combine_masks(attention_mask, causal_mask)
- elif self.causal:
- attention_mask = causal_mask
- elif attention_mask is not None:
- attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))
-
- dropout_rng = None
- if not deterministic and self.config.attn_pdrop > 0.0:
- dropout_rng = self.make_rng("dropout")
-
- # During fast autoregressive decoding, we feed one position at a time,
- # and cache the keys and values step by step.
- if self.causal and (self.has_variable("cache", "cached_key") or init_cache):
- key, value, attention_mask = self._concatenate_to_cache(key, value, query, attention_mask)
-
- # transform boolean mask into float mask
- if attention_mask is not None:
- attention_bias = lax.select(
- attention_mask > 0,
- jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
- jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
- )
- else:
- attention_bias = None
-
- # usual dot product attention
- attn_weights = dot_product_attention_weights(
- query,
- key,
- bias=attention_bias,
- dropout_rng=dropout_rng,
- dropout_rate=self.config.attn_pdrop,
- deterministic=deterministic,
- dtype=self.dtype,
- precision=None,
- )
-
- attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value)
- attn_output = self._merge_heads(attn_output)
- attn_output = self.c_proj(attn_output)
- attn_output = self.resid_dropout(attn_output, deterministic=deterministic)
-
- outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)
- return outputs
-
-
-class FlaxGPT2MLP(nn.Module):
- config: GPT2Config
- intermediate_size: int
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- embed_dim = self.config.hidden_size
- self.c_fc = FlaxConv1D(self.intermediate_size, dtype=self.dtype)
- self.c_proj = FlaxConv1D(embed_dim, dtype=self.dtype)
- self.act = ACT2FN[self.config.activation_function]
- self.dropout = nn.Dropout(rate=self.config.resid_pdrop)
-
- def __call__(self, hidden_states, deterministic: bool = True):
- hidden_states = self.c_fc(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.c_proj(hidden_states)
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- return hidden_states
-
-
-class FlaxGPT2Block(nn.Module):
- config: GPT2Config
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- hidden_size = self.config.hidden_size
- inner_dim = self.config.n_inner if self.config.n_inner is not None else 4 * hidden_size
-
- self.ln_1 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
- self.attn = FlaxGPT2Attention(self.config, dtype=self.dtype)
- self.ln_2 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
-
- if self.config.add_cross_attention:
- self.crossattention = FlaxGPT2Attention(
- config=self.config, dtype=self.dtype, causal=False, is_cross_attention=True
- )
- self.ln_cross_attn = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
-
- self.mlp = FlaxGPT2MLP(self.config, inner_dim, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states,
- attention_mask=attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
- # residual connection
- attn_output = attn_outputs[0] # output_attn: a, (attentions)
- outputs = attn_outputs[1:]
- # residual connection
- hidden_states = attn_output + residual
-
- # Cross-Attention Block
- if encoder_hidden_states is not None:
- # add one self-attention block for cross-attention
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
- "cross-attention layers by setting `config.add_cross_attention=True`"
- )
- residual = hidden_states
- hidden_states = self.ln_cross_attn(hidden_states)
- cross_attn_outputs = self.crossattention(
- hidden_states,
- key_value_states=encoder_hidden_states,
- attention_mask=encoder_attention_mask,
- deterministic=deterministic,
- output_attentions=output_attentions,
- )
- attn_output = cross_attn_outputs[0]
- # residual connection
- hidden_states = residual + attn_output
- outputs = outputs + cross_attn_outputs[1:] # add cross attentions if we output attention weights
-
- residual = hidden_states
- hidden_states = self.ln_2(hidden_states)
- feed_forward_hidden_states = self.mlp(hidden_states, deterministic=deterministic)
- # residual connection
- hidden_states = residual + feed_forward_hidden_states
-
- outputs = (hidden_states,) + outputs
-
- return outputs
-
-
-class FlaxGPT2PreTrainedModel(FlaxPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPT2Config
- base_model_prefix = "transformer"
- module_class: nn.Module = None
-
- def __init__(
- self,
- config: GPT2Config,
- input_shape: Tuple = (1, 1),
- seed: int = 0,
- dtype: jnp.dtype = jnp.float32,
- _do_init: bool = True,
- **kwargs,
- ):
- module = self.module_class(config=config, dtype=dtype, **kwargs)
- super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
-
- def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
- # init input tensors
- input_ids = jnp.zeros(input_shape, dtype="i4")
- attention_mask = jnp.ones_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape)
- params_rng, dropout_rng = jax.random.split(rng)
- rngs = {"params": params_rng, "dropout": dropout_rng}
-
- if self.config.add_cross_attention:
- encoder_hidden_states = jnp.zeros(input_shape + (self.config.n_embd,))
- encoder_attention_mask = attention_mask
- module_init_outputs = self.module.init(
- rngs,
- input_ids,
- attention_mask,
- position_ids,
- encoder_hidden_states,
- encoder_attention_mask,
- return_dict=False,
- )
- else:
- module_init_outputs = self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False)
-
- random_params = module_init_outputs["params"]
-
- if params is not None:
- random_params = flatten_dict(unfreeze(random_params))
- params = flatten_dict(unfreeze(params))
- for missing_key in self._missing_keys:
- params[missing_key] = random_params[missing_key]
- self._missing_keys = set()
- return freeze(unflatten_dict(params))
- else:
- return random_params
-
- def init_cache(self, batch_size, max_length):
- r"""
- Args:
- batch_size (`int`):
- batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
- max_length (`int`):
- maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
- cache.
- """
- # init input variables to retrieve cache
- input_ids = jnp.ones((batch_size, max_length))
- attention_mask = jnp.ones_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
-
- init_variables = self.module.init(
- jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True
- )
- return unfreeze(init_variables["cache"])
-
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- position_ids=None,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- params: dict = None,
- past_key_values: dict = None,
- dropout_rng: jax.random.PRNGKey = None,
- train: bool = False,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ):
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- if encoder_hidden_states is not None and encoder_attention_mask is None:
- batch_size, sequence_length = encoder_hidden_states.shape[:2]
- encoder_attention_mask = jnp.ones((batch_size, sequence_length))
-
- batch_size, sequence_length = input_ids.shape
-
- if position_ids is None:
- if past_key_values is not None:
- raise ValueError("Make sure to provide `position_ids` when passing `past_key_values`.")
-
- position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
-
- if attention_mask is None:
- attention_mask = jnp.ones((batch_size, sequence_length))
-
- # Handle any PRNG if needed
- rngs = {}
- if dropout_rng is not None:
- rngs["dropout"] = dropout_rng
-
- inputs = {"params": params or self.params}
-
- # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be changed by FlaxGPT2Attention module
- if past_key_values:
- inputs["cache"] = past_key_values
- mutable = ["cache"]
- else:
- mutable = False
-
- outputs = self.module.apply(
- inputs,
- jnp.array(input_ids, dtype="i4"),
- jnp.array(attention_mask, dtype="i4"),
- jnp.array(position_ids, dtype="i4"),
- encoder_hidden_states,
- encoder_attention_mask,
- not train,
- False,
- output_attentions,
- output_hidden_states,
- return_dict,
- rngs=rngs,
- mutable=mutable,
- )
-
- # add updated cache to model output
- if past_key_values is not None and return_dict:
- outputs, past_key_values = outputs
- outputs["past_key_values"] = unfreeze(past_key_values["cache"])
- return outputs
- elif past_key_values is not None and not return_dict:
- outputs, past_key_values = outputs
- outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:]
-
- return outputs
-
-
-class FlaxGPT2BlockCollection(nn.Module):
- config: GPT2Config
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.blocks = [
- FlaxGPT2Block(self.config, name=str(i), dtype=self.dtype) for i in range(self.config.num_hidden_layers)
- ]
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- all_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
- all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
-
- for block in self.blocks:
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- layer_outputs = block(
- hidden_states,
- attention_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions += (layer_outputs[1],)
-
- if encoder_hidden_states is not None:
- all_cross_attentions += (layer_outputs[2],)
-
- # this contains possible `None` values - `FlaxGPT2Module` will filter them out
- outputs = (hidden_states, all_hidden_states, all_attentions, all_cross_attentions)
-
- return outputs
-
-
-class FlaxGPT2Module(nn.Module):
- config: GPT2Config
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.embed_dim = self.config.hidden_size
-
- self.wte = nn.Embed(
- self.config.vocab_size,
- self.embed_dim,
- embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- dtype=self.dtype,
- )
- self.wpe = nn.Embed(
- self.config.max_position_embeddings,
- self.embed_dim,
- embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- dtype=self.dtype,
- )
- self.dropout = nn.Dropout(rate=self.config.embd_pdrop)
- self.h = FlaxGPT2BlockCollection(self.config, dtype=self.dtype)
- self.ln_f = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- position_ids,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- deterministic=True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- input_embeds = self.wte(input_ids.astype("i4"))
- position_embeds = self.wpe(position_ids.astype("i4"))
-
- hidden_states = input_embeds + position_embeds
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
-
- outputs = self.h(
- hidden_states,
- attention_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- hidden_states = self.ln_f(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = outputs[1] + (hidden_states,)
- outputs = (hidden_states, all_hidden_states) + outputs[2:]
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- if not return_dict:
- return tuple(v for v in outputs if v is not None)
-
- return FlaxBaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- hidden_states=outputs[1],
- attentions=outputs[2],
- cross_attentions=outputs[3],
- )
-
-
-@add_start_docstrings(
- "The bare GPT2 Model transformer outputting raw hidden-states without any specific head on top.",
- GPT2_START_DOCSTRING,
-)
-class FlaxGPT2Model(FlaxGPT2PreTrainedModel):
- module_class = FlaxGPT2Module
-
-
-append_call_sample_docstring(
- FlaxGPT2Model,
- _CHECKPOINT_FOR_DOC,
- FlaxBaseModelOutputWithPastAndCrossAttentions,
- _CONFIG_FOR_DOC,
-)
-
-
-class FlaxGPT2LMHeadModule(nn.Module):
- config: GPT2Config
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.transformer = FlaxGPT2Module(self.config, dtype=self.dtype)
- self.lm_head = nn.Dense(
- self.config.vocab_size,
- use_bias=False,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- position_ids,
- encoder_hidden_states: Optional[jnp.ndarray] = None,
- encoder_attention_mask: Optional[jnp.ndarray] = None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- outputs = self.transformer(
- input_ids,
- attention_mask,
- position_ids,
- encoder_hidden_states,
- encoder_attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
-
- if self.config.tie_word_embeddings:
- shared_kernel = self.transformer.variables["params"]["wte"]["embedding"].T
- lm_logits = self.lm_head.apply({"params": {"kernel": shared_kernel}}, hidden_states)
- else:
- lm_logits = self.lm_head(hidden_states)
-
- if not return_dict:
- return (lm_logits,) + outputs[1:]
-
- return FlaxCausalLMOutputWithCrossAttentions(
- logits=lm_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- cross_attentions=outputs.cross_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- GPT2_START_DOCSTRING,
-)
-class FlaxGPT2LMHeadModel(FlaxGPT2PreTrainedModel):
- module_class = FlaxGPT2LMHeadModule
-
- def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jax.Array] = None):
- # initializing the cache
- batch_size, seq_length = input_ids.shape
-
- past_key_values = self.init_cache(batch_size, max_length)
- # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
- # But since GPT2 uses a causal mask, those positions are masked anyways.
- # Thus we can create a single static attention_mask here, which is more efficient for compilation
- extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
- if attention_mask is not None:
- position_ids = attention_mask.cumsum(axis=-1) - 1
- extended_attention_mask = lax.dynamic_update_slice(
- extended_attention_mask, attention_mask.astype("i4"), (0, 0)
- )
- else:
- position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))
-
- return {
- "past_key_values": past_key_values,
- "attention_mask": extended_attention_mask,
- "position_ids": position_ids,
- }
-
- def update_inputs_for_generation(self, model_outputs, model_kwargs):
- model_kwargs["past_key_values"] = model_outputs.past_key_values
- model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1
- return model_kwargs
-
-
-append_call_sample_docstring(
- FlaxGPT2LMHeadModel,
- _CHECKPOINT_FOR_DOC,
- FlaxCausalLMOutputWithCrossAttentions,
- _CONFIG_FOR_DOC,
-)
diff --git a/transformers/models/gpt2/modeling_gpt2.py b/transformers/models/gpt2/modeling_gpt2.py
deleted file mode 100644
index 1409a3fc3f0fcbc16fb17d7fdbe1d166f2bc38ce..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/modeling_gpt2.py
+++ /dev/null
@@ -1,1944 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch OpenAI GPT-2 model."""
-
-import math
-import os
-import warnings
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch import nn
-from torch.cuda.amp import autocast
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutputWithPastAndCrossAttentions,
- CausalLMOutputWithCrossAttentions,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutputWithPast,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel, SequenceSummary
-from ...pytorch_utils import Conv1D, find_pruneable_heads_and_indices, prune_conv1d_layer
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- logging,
- replace_return_docstrings,
-)
-from ...utils.model_parallel_utils import assert_device_map, get_device_map
-from .configuration_gpt2 import GPT2Config
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "openai-community/gpt2"
-_CONFIG_FOR_DOC = "GPT2Config"
-
-
-from ..deprecated._archive_maps import GPT2_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-def load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path):
- """Load tf checkpoints in a pytorch model"""
- try:
- import re
-
- import tensorflow as tf
- except ImportError:
- logger.error(
- "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
- "https://www.tensorflow.org/install/ for installation instructions."
- )
- raise
- tf_path = os.path.abspath(gpt2_checkpoint_path)
- logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
- # Load weights from TF model
- init_vars = tf.train.list_variables(tf_path)
- names = []
- arrays = []
- for name, shape in init_vars:
- logger.info(f"Loading TF weight {name} with shape {shape}")
- array = tf.train.load_variable(tf_path, name)
- names.append(name)
- arrays.append(array.squeeze())
-
- for name, array in zip(names, arrays):
- name = name[6:] # skip "model/"
- name = name.split("/")
- pointer = model
- for m_name in name:
- if re.fullmatch(r"[A-Za-z]+\d+", m_name):
- scope_names = re.split(r"(\d+)", m_name)
- else:
- scope_names = [m_name]
- if scope_names[0] == "w" or scope_names[0] == "g":
- pointer = getattr(pointer, "weight")
- elif scope_names[0] == "b":
- pointer = getattr(pointer, "bias")
- elif scope_names[0] == "wpe" or scope_names[0] == "wte":
- pointer = getattr(pointer, scope_names[0])
- pointer = getattr(pointer, "weight")
- else:
- pointer = getattr(pointer, scope_names[0])
- if len(scope_names) >= 2:
- num = int(scope_names[1])
- pointer = pointer[num]
- try:
- if pointer.shape != array.shape:
- raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
- except ValueError as e:
- e.args += (pointer.shape, array.shape)
- raise
- logger.info(f"Initialize PyTorch weight {name}")
- pointer.data = torch.from_numpy(array)
- return model
-
-
-class GPT2Attention(nn.Module):
- def __init__(self, config, is_cross_attention=False, layer_idx=None):
- super().__init__()
- self.config = config
- max_positions = config.max_position_embeddings
- self.register_buffer(
- "bias",
- torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
- 1, 1, max_positions, max_positions
- ),
- persistent=False,
- )
- self.register_buffer("masked_bias", torch.tensor(-1e4), persistent=False)
-
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_heads
- self.split_size = self.embed_dim
- if self.head_dim * self.num_heads != self.embed_dim:
- raise ValueError(
- f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {self.num_heads})."
- )
-
- self.scale_attn_weights = config.scale_attn_weights
- self.is_cross_attention = is_cross_attention
-
- # Layer-wise attention scaling, reordering, and upcasting
- self.scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx
- self.layer_idx = layer_idx
- self.reorder_and_upcast_attn = config.reorder_and_upcast_attn
-
- if self.is_cross_attention:
- self.c_attn = Conv1D(2 * self.embed_dim, self.embed_dim)
- self.q_attn = Conv1D(self.embed_dim, self.embed_dim)
- else:
- self.c_attn = Conv1D(3 * self.embed_dim, self.embed_dim)
- self.c_proj = Conv1D(self.embed_dim, self.embed_dim)
-
- self.attn_dropout = nn.Dropout(config.attn_pdrop)
- self.resid_dropout = nn.Dropout(config.resid_pdrop)
- self.is_causal = True
-
- self.pruned_heads = set()
-
- def prune_heads(self, heads):
- if len(heads) == 0:
- return
- heads, index = find_pruneable_heads_and_indices(heads, self.num_heads, self.head_dim, self.pruned_heads)
- index_attn = torch.cat([index, index + self.split_size, index + (2 * self.split_size)])
-
- # Prune conv1d layers
- self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
- self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0)
-
- # Update hyper params
- self.split_size = (self.split_size // self.num_heads) * (self.num_heads - len(heads))
- self.num_heads = self.num_heads - len(heads)
- self.pruned_heads = self.pruned_heads.union(heads)
-
- def _attn(self, query, key, value, attention_mask=None, head_mask=None):
- attn_weights = torch.matmul(query, key.transpose(-1, -2))
-
- if self.scale_attn_weights:
- attn_weights = attn_weights / torch.full(
- [], value.size(-1) ** 0.5, dtype=attn_weights.dtype, device=attn_weights.device
- )
-
- # Layer-wise attention scaling
- if self.scale_attn_by_inverse_layer_idx:
- attn_weights = attn_weights / float(self.layer_idx + 1)
-
- if not self.is_cross_attention:
- # if only "normal" attention layer implements causal mask
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
- mask_value = torch.finfo(attn_weights.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.full([], mask_value, dtype=attn_weights.dtype, device=attn_weights.device)
- attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value)
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise
- attn_weights = attn_weights.type(value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, head_mask=None):
- # Use `torch.baddbmm` (a bit more efficient w/ alpha param for scaling -- from Megatron-LM)
- bsz, num_heads, q_seq_len, dk = query.size()
- _, _, k_seq_len, _ = key.size()
-
- # Preallocate attn_weights for `baddbmm`
- attn_weights = torch.empty(bsz * num_heads, q_seq_len, k_seq_len, dtype=torch.float32, device=query.device)
-
- # Compute Scale Factor
- scale_factor = 1.0
- if self.scale_attn_weights:
- scale_factor /= float(value.size(-1)) ** 0.5
-
- if self.scale_attn_by_inverse_layer_idx:
- scale_factor /= float(self.layer_idx + 1)
-
- # Upcast (turn off autocast) and reorder (Scale K by 1 / root(dk))
- with autocast(enabled=False):
- q, k = query.reshape(-1, q_seq_len, dk), key.transpose(-1, -2).reshape(-1, dk, k_seq_len)
- attn_weights = torch.baddbmm(attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor)
- attn_weights = attn_weights.reshape(bsz, num_heads, q_seq_len, k_seq_len)
-
- if not self.is_cross_attention:
- # if only "normal" attention layer implements causal mask
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
- mask_value = torch.finfo(attn_weights.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
- attn_weights = torch.where(causal_mask, attn_weights, mask_value)
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op if otherwise
- if attn_weights.dtype != torch.float32:
- raise RuntimeError("Error with upcasting, attn_weights does not have dtype torch.float32")
- attn_weights = attn_weights.type(value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def _split_heads(self, tensor, num_heads, attn_head_size):
- """
- Splits hidden_size dim into attn_head_size and num_heads
- """
- new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
- tensor = tensor.view(new_shape)
- return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)
-
- def _merge_heads(self, tensor, num_heads, attn_head_size):
- """
- Merges attn_head_size dim and num_attn_heads dim into hidden_size
- """
- tensor = tensor.permute(0, 2, 1, 3).contiguous()
- new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
- return tensor.view(new_shape)
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
- if encoder_hidden_states is not None:
- if not hasattr(self, "q_attn"):
- raise ValueError(
- "If class is used as cross attention, the weights `q_attn` have to be defined. "
- "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
- )
-
- query = self.q_attn(hidden_states)
- key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
- attention_mask = encoder_attention_mask
- else:
- query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
-
- query = self._split_heads(query, self.num_heads, self.head_dim)
- key = self._split_heads(key, self.num_heads, self.head_dim)
- value = self._split_heads(value, self.num_heads, self.head_dim)
-
- if layer_past is not None:
- past_key, past_value = layer_past
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- if use_cache is True:
- present = (key, value)
- else:
- present = None
-
- if self.reorder_and_upcast_attn:
- attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
- else:
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
- attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
- attn_output = self.c_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs # a, present, (attentions)
-
-
-class GPT2FlashAttention2(GPT2Attention):
- """
- GPT2 flash attention module. This module inherits from `GPT2Attention` as the weights of the module stays
- untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
- flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
- bsz, _, _ = hidden_states.size()
- if encoder_hidden_states is not None:
- if not hasattr(self, "q_attn"):
- raise ValueError(
- "If class is used as cross attention, the weights `q_attn` have to be defined. "
- "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
- )
-
- query = self.q_attn(hidden_states)
- key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
- attention_mask = encoder_attention_mask
- else:
- query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
-
- query = self._split_heads(query, self.num_heads, self.head_dim)
- key = self._split_heads(key, self.num_heads, self.head_dim)
- value = self._split_heads(value, self.num_heads, self.head_dim)
-
- if layer_past is not None:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- present = None
- if use_cache is True:
- present = (key, value)
-
- query_length = query.shape[2]
- tgt_len = key.shape[2]
-
- # Flash attention requires the input to have the shape
- # batch_size x seq_length x head_dim x hidden_dim
- query = query.transpose(1, 2).view(bsz, query_length, self.num_heads, self.head_dim)
- key = key.transpose(1, 2).view(bsz, tgt_len, self.num_heads, self.head_dim)
- value = value.transpose(1, 2).view(bsz, tgt_len, self.num_heads, self.head_dim)
-
- attn_dropout = self.attn_dropout.p if self.training else 0.0
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in the correct dtype just to be sure everything works as expected.
- # This might slowdown training & inference so it is recommended to not cast the LayerNorms
- # in fp32. (LlamaRMSNorm handles it correctly)
-
- if query.dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.c_proj.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
-
- query = query.to(target_dtype)
- key = key.to(target_dtype)
- value = value.to(target_dtype)
-
- attn_output = self._flash_attention_forward(
- query, key, value, attention_mask, query_length, dropout=attn_dropout
- )
-
- attn_weights_reshaped = attn_output.reshape(bsz, query_length, self.num_heads * self.head_dim)
- attn_output = self.c_proj(attn_weights_reshaped)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights_reshaped,)
-
- return outputs
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-class GPT2MLP(nn.Module):
- def __init__(self, intermediate_size, config):
- super().__init__()
- embed_dim = config.hidden_size
- self.c_fc = Conv1D(intermediate_size, embed_dim)
- self.c_proj = Conv1D(embed_dim, intermediate_size)
- self.act = ACT2FN[config.activation_function]
- self.dropout = nn.Dropout(config.resid_pdrop)
-
- def forward(self, hidden_states: Optional[Tuple[torch.FloatTensor]]) -> torch.FloatTensor:
- hidden_states = self.c_fc(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.c_proj(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-GPT2_ATTENTION_CLASSES = {
- "eager": GPT2Attention,
- "flash_attention_2": GPT2FlashAttention2,
-}
-
-
-class GPT2Block(nn.Module):
- def __init__(self, config, layer_idx=None):
- super().__init__()
- hidden_size = config.hidden_size
- inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
- attention_class = GPT2_ATTENTION_CLASSES[config._attn_implementation]
-
- self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
- self.attn = attention_class(config=config, layer_idx=layer_idx)
- self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
- if config.add_cross_attention:
- self.crossattention = attention_class(config=config, is_cross_attention=True, layer_idx=layer_idx)
- self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
- self.mlp = GPT2MLP(inner_dim, config)
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
- outputs = attn_outputs[1:]
- # residual connection
- hidden_states = attn_output + residual
-
- if encoder_hidden_states is not None:
- # add one self-attention block for cross-attention
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
- "cross-attention layers by setting `config.add_cross_attention=True`"
- )
- residual = hidden_states
- hidden_states = self.ln_cross_attn(hidden_states)
- cross_attn_outputs = self.crossattention(
- hidden_states,
- attention_mask=attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- )
- attn_output = cross_attn_outputs[0]
- # residual connection
- hidden_states = residual + attn_output
- outputs = outputs + cross_attn_outputs[2:] # add cross attentions if we output attention weights
-
- residual = hidden_states
- hidden_states = self.ln_2(hidden_states)
- feed_forward_hidden_states = self.mlp(hidden_states)
- # residual connection
- hidden_states = residual + feed_forward_hidden_states
-
- if use_cache:
- outputs = (hidden_states,) + outputs
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- return outputs # hidden_states, present, (attentions, cross_attentions)
-
-
-class GPT2PreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPT2Config
- load_tf_weights = load_tf_weights_in_gpt2
- base_model_prefix = "transformer"
- is_parallelizable = True
- supports_gradient_checkpointing = True
- _no_split_modules = ["GPT2Block"]
- _skip_keys_device_placement = "past_key_values"
- _supports_flash_attn_2 = True
-
- def __init__(self, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, (nn.Linear, Conv1D)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
- # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
- # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
- # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
- # > -- GPT-2 :: https://openai.com/blog/better-language-models/
- #
- # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
- for name, p in module.named_parameters():
- if name == "c_proj.weight":
- # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
- p.data.normal_(mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.n_layer)))
-
-
-@dataclass
-class GPT2DoubleHeadsModelOutput(ModelOutput):
- """
- Base class for outputs of models predicting if two sentences are consecutive or not.
-
- Args:
- loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
- Language modeling loss.
- mc_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mc_labels` is provided):
- Multiple choice classification loss.
- logits (`torch.FloatTensor` of shape `(batch_size, num_choices, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
- mc_logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`):
- Prediction scores of the multiple choice classification head (scores for each choice before SoftMax).
- past_key_values (`Tuple[Tuple[torch.Tensor]]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- Tuple of length `config.n_layers`, containing tuples of tensors of shape `(batch_size, num_heads,
- sequence_length, embed_size_per_head)`).
-
- Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
- `past_key_values` input) to speed up sequential decoding.
- hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
- shape `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- GPT2Attentions weights after the attention softmax, used to compute the weighted average in the
- self-attention heads.
- """
-
- loss: Optional[torch.FloatTensor] = None
- mc_loss: Optional[torch.FloatTensor] = None
- logits: torch.FloatTensor = None
- mc_logits: torch.FloatTensor = None
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
- hidden_states: Optional[Tuple[torch.FloatTensor]] = None
- attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-
-GPT2_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`GPT2Config`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPT2_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
- `past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input
- sequence tokens in the vocabulary.
-
- If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
- `input_ids`.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`):
- Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
- their past given to this model should not be passed as `input_ids` as they have already been computed.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- If `past_key_values` is used, `attention_mask` needs to contain the masking strategy that was used for
- `past_key_values`. In other words, the `attention_mask` always has to have the length:
- `len(past_key_values) + len(input_ids)`
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
-
- If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see
- `past_key_values`).
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-PARALLELIZE_DOCSTRING = r"""
- This is an experimental feature and is a subject to change at a moment's notice.
-
- Uses a device map to distribute attention modules of the model across several devices. If no device map is given,
- it will evenly distribute blocks across all devices.
-
- Args:
- device_map (`Dict[int, list]`, optional, defaults to None):
- A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always
- automatically mapped to the first device (for esoteric reasons). That means that the first device should
- have fewer attention modules mapped to it than other devices. For reference, the gpt2 models have the
- following number of attention modules:
-
- - openai-community/gpt2: 12
- - openai-community/gpt2-medium: 24
- - openai-community/gpt2-large: 36
- - openai-community/gpt2-xl: 48
-
- Example:
-
- ```python
- # Here is an example of a device map on a machine with 4 GPUs using gpt2-xl, which has a total of 48 attention modules:
- model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2-xl")
- device_map = {
- 0: [0, 1, 2, 3, 4, 5, 6, 7, 8],
- 1: [9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
- 2: [22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34],
- 3: [35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47],
- }
- model.parallelize(device_map)
- ```
-"""
-DEPARALLELIZE_DOCSTRING = r"""
- Moves the model to cpu from a model parallel state.
-
- Example:
-
- ```python
- # On a 4 GPU machine with openai-community/gpt2-large:
- model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2-large")
- device_map = {
- 0: [0, 1, 2, 3, 4, 5, 6, 7],
- 1: [8, 9, 10, 11, 12, 13, 14, 15],
- 2: [16, 17, 18, 19, 20, 21, 22, 23],
- 3: [24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35],
- }
- model.parallelize(device_map) # Splits the model across several devices
- model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache()
- ```
-"""
-
-
-@add_start_docstrings(
- "The bare GPT2 Model transformer outputting raw hidden-states without any specific head on top.",
- GPT2_START_DOCSTRING,
-)
-class GPT2Model(GPT2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.embed_dim = config.hidden_size
-
- self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
- self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
-
- self.drop = nn.Dropout(config.embd_pdrop)
- self.h = nn.ModuleList([GPT2Block(config, layer_idx=i) for i in range(config.num_hidden_layers)])
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
- self.gradient_checkpointing = False
- self._attn_implementation = config._attn_implementation
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings(PARALLELIZE_DOCSTRING)
- def parallelize(self, device_map=None):
- # Check validity of device_map
- warnings.warn(
- "`GPT2Model.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your"
- " model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own"
- " `device_map` but it needs to be a dictionary module_name to device, so for instance {'h.0': 0, 'h.1': 1,"
- " ...}",
- FutureWarning,
- )
- self.device_map = (
- get_device_map(len(self.h), range(torch.cuda.device_count())) if device_map is None else device_map
- )
- assert_device_map(self.device_map, len(self.h))
- self.model_parallel = True
- self.first_device = "cpu" if "cpu" in self.device_map.keys() else "cuda:" + str(min(self.device_map.keys()))
- self.last_device = "cuda:" + str(max(self.device_map.keys()))
- self.wte = self.wte.to(self.first_device)
- self.wpe = self.wpe.to(self.first_device)
- # Load onto devices
- for k, v in self.device_map.items():
- for block in v:
- cuda_device = "cuda:" + str(k)
- self.h[block] = self.h[block].to(cuda_device)
- # ln_f to last
- self.ln_f = self.ln_f.to(self.last_device)
-
- @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
- def deparallelize(self):
- warnings.warn(
- "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.",
- FutureWarning,
- )
- self.model_parallel = False
- self.device_map = None
- self.first_device = "cpu"
- self.last_device = "cpu"
- self.wte = self.wte.to("cpu")
- self.wpe = self.wpe.to("cpu")
- for index in range(len(self.h)):
- self.h[index] = self.h[index].to("cpu")
- self.ln_f = self.ln_f.to("cpu")
- torch.cuda.empty_cache()
-
- def get_input_embeddings(self):
- return self.wte
-
- def set_input_embeddings(self, new_embeddings):
- self.wte = new_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
- """
- for layer, heads in heads_to_prune.items():
- self.h[layer].attn.prune_heads(heads)
-
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPastAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- input_ids = input_ids.view(-1, input_shape[-1])
- batch_size = input_ids.shape[0]
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- batch_size = inputs_embeds.shape[0]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if token_type_ids is not None:
- token_type_ids = token_type_ids.view(-1, input_shape[-1])
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * len(self.h))
- else:
- past_length = past_key_values[0][0].size(-2)
- if position_ids is None:
- position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0)
-
- # Attention mask.
- if attention_mask is not None:
- attention_mask = attention_mask.view(batch_size, -1)
- if self._attn_implementation == "flash_attention_2":
- attention_mask = attention_mask if 0 in attention_mask else None
- else:
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask = attention_mask[:, None, None, :]
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and the dtype's smallest value for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
- attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
- # If a 2D or 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- if self.config.add_cross_attention and encoder_hidden_states is not None:
- encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
- encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
- if encoder_attention_mask is None:
- encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
- if self._attn_implementation != "flash_attention_2":
- encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
- else:
- encoder_attention_mask = None
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # head_mask has shape n_layer x batch x n_heads x N x N
- head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-
- if inputs_embeds is None:
- inputs_embeds = self.wte(input_ids)
- position_embeds = self.wpe(position_ids)
- hidden_states = inputs_embeds + position_embeds
-
- if token_type_ids is not None:
- token_type_embeds = self.wte(token_type_ids)
- hidden_states = hidden_states + token_type_embeds
-
- hidden_states = self.drop(hidden_states)
-
- output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- presents = () if use_cache else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
- all_hidden_states = () if output_hidden_states else None
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- # Model parallel
- if self.model_parallel:
- torch.cuda.set_device(hidden_states.device)
- # Ensure layer_past is on same device as hidden_states (might not be correct)
- if layer_past is not None:
- layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past)
- # Ensure that attention_mask is always on the same device as hidden_states
- if attention_mask is not None:
- attention_mask = attention_mask.to(hidden_states.device)
- if isinstance(head_mask, torch.Tensor):
- head_mask = head_mask.to(hidden_states.device)
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- None,
- attention_mask,
- head_mask[i],
- encoder_hidden_states,
- encoder_attention_mask,
- use_cache,
- output_attentions,
- )
- else:
- outputs = block(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask[i],
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
-
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],)
-
- # Model Parallel: If it's the last layer for that device, put things on the next device
- if self.model_parallel:
- for k, v in self.device_map.items():
- if i == v[-1] and "cuda:" + str(k) != self.last_device:
- hidden_states = hidden_states.to("cuda:" + str(k + 1))
-
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = hidden_states.view(output_shape)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions]
- if v is not None
- )
-
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- GPT2_START_DOCSTRING,
-)
-class GPT2LMHeadModel(GPT2PreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.transformer = GPT2Model(config)
- self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings(PARALLELIZE_DOCSTRING)
- def parallelize(self, device_map=None):
- warnings.warn(
- "`GPT2LMHeadModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should load"
- " your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own"
- " `device_map` but it needs to be a dictionary module_name to device, so for instance {'transformer.h.0':"
- " 0, 'transformer.h.1': 1, ...}",
- FutureWarning,
- )
- self.device_map = (
- get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
- if device_map is None
- else device_map
- )
- assert_device_map(self.device_map, len(self.transformer.h))
- self.transformer.parallelize(self.device_map)
- self.lm_head = self.lm_head.to(self.transformer.first_device)
- self.model_parallel = True
-
- @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
- def deparallelize(self):
- warnings.warn(
- "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.",
- FutureWarning,
- )
- self.transformer.deparallelize()
- self.transformer = self.transformer.to("cpu")
- self.lm_head = self.lm_head.to("cpu")
- self.model_parallel = False
- torch.cuda.empty_cache()
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # Omit tokens covered by past_key_values
- if past_key_values:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
- if token_type_ids is not None:
- token_type_ids = token_type_ids[:, -input_ids.shape[1] :]
-
- attention_mask = kwargs.get("attention_mask", None)
- position_ids = kwargs.get("position_ids", None)
-
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
- else:
- position_ids = None
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- model_inputs = {"input_ids": input_ids}
-
- model_inputs.update(
- {
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "position_ids": position_ids,
- "attention_mask": attention_mask,
- "token_type_ids": token_type_ids,
- }
- )
-
- return model_inputs
-
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutputWithCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
-
- # Set device for model parallelism
- if self.model_parallel:
- torch.cuda.set_device(self.transformer.first_device)
- hidden_states = hidden_states.to(self.lm_head.weight.device)
-
- lm_logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
- if not return_dict:
- output = (lm_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithCrossAttentions(
- loss=loss,
- logits=lm_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- cross_attentions=transformer_outputs.cross_attentions,
- )
-
- @staticmethod
- def _reorder_cache(
- past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
- ) -> Tuple[Tuple[torch.Tensor]]:
- """
- This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
- [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
- beam_idx at every generation step.
- """
- return tuple(
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
- for layer_past in past_key_values
- )
-
-
-@add_start_docstrings(
- """
-The GPT2 Model transformer with a language modeling and a multiple-choice classification head on top e.g. for
-RocStories/SWAG tasks. The two heads are two linear layers. The language modeling head has its weights tied to the
-input embeddings, the classification head takes as input the input of a specified classification token index in the
-input sequence).
-""",
- GPT2_START_DOCSTRING,
-)
-class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- config.num_labels = 1
- self.transformer = GPT2Model(config)
- self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
- self.multiple_choice_head = SequenceSummary(config)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings(PARALLELIZE_DOCSTRING)
- def parallelize(self, device_map=None):
- warnings.warn(
- "`GPT2DoubleHeadsModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should"
- " load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your"
- " own `device_map` but it needs to be a dictionary module_name to device, so for instance"
- " {'transformer.h.0': 0, 'transformer.h.1': 1, ...}",
- FutureWarning,
- )
- self.device_map = (
- get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
- if device_map is None
- else device_map
- )
- assert_device_map(self.device_map, len(self.transformer.h))
- self.transformer.parallelize(self.device_map)
- self.lm_head = self.lm_head.to(self.transformer.first_device)
- self.multiple_choice_head = self.multiple_choice_head.to(self.transformer.first_device)
- self.model_parallel = True
-
- @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
- def deparallelize(self):
- warnings.warn(
- "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.",
- FutureWarning,
- )
- self.transformer.deparallelize()
- self.transformer = self.transformer.to("cpu")
- self.lm_head = self.lm_head.to("cpu")
- self.multiple_choice_head = self.multiple_choice_head.to("cpu")
- self.model_parallel = False
- torch.cuda.empty_cache()
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # Omit tokens covered by past_key_values
- if past_key_values:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
- if token_type_ids is not None:
- token_type_ids = token_type_ids[:, -input_ids.shape[1] :]
-
- attention_mask = kwargs.get("attention_mask", None)
- position_ids = kwargs.get("position_ids", None)
-
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
- else:
- position_ids = None
-
- return {
- "input_ids": input_ids,
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "position_ids": position_ids,
- "attention_mask": attention_mask,
- "token_type_ids": token_type_ids,
- }
-
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=GPT2DoubleHeadsModelOutput, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- mc_token_ids: Optional[torch.LongTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- mc_labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- **kwargs,
- ) -> Union[Tuple, GPT2DoubleHeadsModelOutput]:
- r"""
- mc_token_ids (`torch.LongTensor` of shape `(batch_size, num_choices)`, *optional*, default to index of the last token of the input):
- Index of the classification token in each input sequence. Selected in the range `[0, input_ids.size(-1) -
- 1]`.
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids`. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`. All labels set to
- `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size - 1]`
- mc_labels (`torch.LongTensor` of shape `(batch_size)`, *optional*):
- Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]`
- where *num_choices* is the size of the second dimension of the input tensors. (see *input_ids* above)
-
- Return:
-
- Example:
-
- ```python
- >>> import torch
- >>> from transformers import AutoTokenizer, GPT2DoubleHeadsModel
-
- >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
- >>> model = GPT2DoubleHeadsModel.from_pretrained("openai-community/gpt2")
-
- >>> # Add a [CLS] to the vocabulary (we should train it also!)
- >>> num_added_tokens = tokenizer.add_special_tokens({"cls_token": "[CLS]"})
- >>> # Update the model embeddings with the new vocabulary size
- >>> embedding_layer = model.resize_token_embeddings(len(tokenizer))
-
- >>> choices = ["Hello, my dog is cute [CLS]", "Hello, my cat is cute [CLS]"]
- >>> encoded_choices = [tokenizer.encode(s) for s in choices]
- >>> cls_token_location = [tokens.index(tokenizer.cls_token_id) for tokens in encoded_choices]
-
- >>> input_ids = torch.tensor(encoded_choices).unsqueeze(0) # Batch size: 1, number of choices: 2
- >>> mc_token_ids = torch.tensor([cls_token_location]) # Batch size: 1
-
- >>> outputs = model(input_ids, mc_token_ids=mc_token_ids)
- >>> lm_logits = outputs.logits
- >>> mc_logits = outputs.mc_logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = transformer_outputs[0]
-
- # Set device for model parallelism
- if self.model_parallel:
- torch.cuda.set_device(self.transformer.first_device)
- hidden_states = hidden_states.to(self.lm_head.weight.device)
-
- lm_logits = self.lm_head(hidden_states)
- mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids).squeeze(-1)
-
- mc_loss = None
- if mc_labels is not None:
- loss_fct = CrossEntropyLoss()
- mc_loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1))
- lm_loss = None
- if labels is not None:
- labels = labels.to(lm_logits.device)
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- loss_fct = CrossEntropyLoss()
- lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
- if not return_dict:
- output = (lm_logits, mc_logits) + transformer_outputs[1:]
- if mc_loss is not None:
- output = (mc_loss,) + output
- return ((lm_loss,) + output) if lm_loss is not None else output
-
- return GPT2DoubleHeadsModelOutput(
- loss=lm_loss,
- mc_loss=mc_loss,
- logits=lm_logits,
- mc_logits=mc_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- @staticmethod
- def _reorder_cache(
- past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
- ) -> Tuple[Tuple[torch.Tensor]]:
- """
- This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
- [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
- beam_idx at every generation step.
- """
- return tuple(
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
- for layer_past in past_key_values
- )
-
-
-@add_start_docstrings(
- """
- The GPT2 Model transformer with a sequence classification head on top (linear layer).
-
- [`GPT2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT-1) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- GPT2_START_DOCSTRING,
-)
-class GPT2ForSequenceClassification(GPT2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = GPT2Model(config)
- self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint="microsoft/DialogRPT-updown",
- output_type=SequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size, sequence_length = input_ids.shape[:2]
- else:
- batch_size, sequence_length = inputs_embeds.shape[:2]
-
- assert (
- self.config.pad_token_id is not None or batch_size == 1
- ), "Cannot handle batch sizes > 1 if no padding token is defined."
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
-
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- GPT2 Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- GPT2_START_DOCSTRING,
-)
-class GPT2ForTokenClassification(GPT2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.transformer = GPT2Model(config)
- if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
- classifier_dropout = config.classifier_dropout
- elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
- classifier_dropout = config.hidden_dropout
- else:
- classifier_dropout = 0.1
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- # fmt: off
- @add_code_sample_docstrings(
- checkpoint="brad1141/gpt2-finetuned-comp2",
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_loss=0.25,
- expected_output=[
- "Lead",
- "Lead",
- "Lead",
- "Position",
- "Lead",
- "Lead",
- "Lead",
- "Lead",
- "Lead",
- "Lead",
- "Lead",
- "Lead",
- ],
- )
- # fmt: on
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = transformer_outputs[0]
- hidden_states = self.dropout(hidden_states)
- logits = self.classifier(hidden_states)
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + transformer_outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT-2 Model transformer with a span classification head on top for extractive question-answering tasks like
- SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- GPT2_START_DOCSTRING,
-)
-class GPT2ForQuestionAnswering(GPT2PreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = GPT2Model(config)
- self.qa_outputs = nn.Linear(config.hidden_size, 2)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- real_checkpoint=_CHECKPOINT_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- start_positions: Optional[torch.LongTensor] = None,
- end_positions: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1).to(start_logits.device)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1).to(end_logits.device)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/gpt2/modeling_tf_gpt2.py b/transformers/models/gpt2/modeling_tf_gpt2.py
deleted file mode 100644
index 26a4e7a398ae8d598e350cdfc1c7b17532936151..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/modeling_tf_gpt2.py
+++ /dev/null
@@ -1,1238 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 OpenAI GPT-2 model."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...modeling_tf_outputs import (
- TFBaseModelOutputWithPastAndCrossAttentions,
- TFCausalLMOutputWithCrossAttentions,
- TFSequenceClassifierOutputWithPast,
-)
-from ...modeling_tf_utils import (
- TFCausalLanguageModelingLoss,
- TFConv1D,
- TFModelInputType,
- TFPreTrainedModel,
- TFSequenceClassificationLoss,
- TFSequenceSummary,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import (
- ModelOutput,
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- logging,
- replace_return_docstrings,
-)
-from .configuration_gpt2 import GPT2Config
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "openai-community/gpt2"
-_CONFIG_FOR_DOC = "GPT2Config"
-
-
-from ..deprecated._archive_maps import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class TFAttention(keras.layers.Layer):
- def __init__(self, nx, config, scale=False, is_cross_attention=False, **kwargs):
- super().__init__(**kwargs)
-
- n_state = nx # in Attention: n_state=768 (nx=n_embd)
- # [switch nx => n_state from Block to Attention to keep identical to TF implementation]
- assert n_state % config.n_head == 0
- self.n_head = config.n_head
- self.split_size = n_state
- self.scale = scale
- self.output_attentions = config.output_attentions
-
- self.is_cross_attention = is_cross_attention
-
- if self.is_cross_attention:
- self.c_attn = TFConv1D(n_state * 2, nx, initializer_range=config.initializer_range, name="c_attn")
- self.q_attn = TFConv1D(n_state, nx, initializer_range=config.initializer_range, name="q_attn")
- else:
- self.c_attn = TFConv1D(n_state * 3, nx, initializer_range=config.initializer_range, name="c_attn")
-
- self.c_proj = TFConv1D(n_state, nx, initializer_range=config.initializer_range, name="c_proj")
- self.attn_dropout = keras.layers.Dropout(config.attn_pdrop)
- self.resid_dropout = keras.layers.Dropout(config.resid_pdrop)
- self.pruned_heads = set()
- self.embed_dim = n_state
-
- def prune_heads(self, heads):
- pass
-
- @staticmethod
- def causal_attention_mask(nd, ns, dtype):
- """
- 1's in the lower triangle, counting from the lower right corner. Same as tf.matrix_band_part(tf.ones([nd, ns]),
- -1, ns-nd), but doesn't produce garbage on TPUs.
- """
- i = tf.range(nd)[:, None]
- j = tf.range(ns)
- m = i >= j - ns + nd
- return tf.cast(m, dtype)
-
- def _attn(self, q, k, v, attention_mask, head_mask, output_attentions, training=False):
- # q, k, v have shape [batch, heads, sequence, features]
- w = tf.matmul(q, k, transpose_b=True)
- if self.scale:
- dk = tf.cast(shape_list(k)[-1], dtype=w.dtype) # scale attention_scores
- w = w / tf.math.sqrt(dk)
-
- if not self.is_cross_attention:
- # if only "normal" attention layer implements causal mask
-
- # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
- _, _, nd, ns = shape_list(w)
- b = self.causal_attention_mask(nd, ns, dtype=w.dtype)
- b = tf.reshape(b, [1, 1, nd, ns])
- w = w * b - 1e4 * (1 - b)
-
- if attention_mask is not None:
- # Apply the attention mask
- attention_mask = tf.cast(attention_mask, dtype=w.dtype)
- w = w + attention_mask
-
- w = stable_softmax(w, axis=-1)
- w = self.attn_dropout(w, training=training)
-
- # Mask heads if we want to
- if head_mask is not None:
- w = w * head_mask
-
- outputs = [tf.matmul(w, v)]
- if output_attentions:
- outputs.append(w)
- return outputs
-
- def merge_heads(self, x):
- x = tf.transpose(x, [0, 2, 1, 3])
- x_shape = shape_list(x)
- new_x_shape = x_shape[:-2] + [x_shape[-2] * x_shape[-1]]
- return tf.reshape(x, new_x_shape)
-
- def split_heads(self, x):
- x_shape = shape_list(x)
- new_x_shape = x_shape[:-1] + [self.n_head, x_shape[-1] // self.n_head]
- x = tf.reshape(x, new_x_shape)
- return tf.transpose(x, (0, 2, 1, 3)) # (batch, head, seq_length, head_features)
-
- def call(
- self,
- x,
- layer_past,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- use_cache,
- output_attentions,
- training=False,
- ):
- if encoder_hidden_states is not None:
- if not hasattr(self, "q_attn"):
- raise ValueError(
- "If class is used as cross attention, the weights `q_attn` have to be defined. "
- "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
- )
-
- query = self.q_attn(x)
- kv_out = self.c_attn(encoder_hidden_states)
- key, value = tf.split(kv_out, 2, axis=2)
- attention_mask = encoder_attention_mask
- else:
- x = self.c_attn(x)
- query, key, value = tf.split(x, 3, axis=2)
-
- query = self.split_heads(query)
- key = self.split_heads(key)
- value = self.split_heads(value)
- if layer_past is not None:
- past_key, past_value = tf.unstack(layer_past, axis=0, num=2)
- key = tf.concat([past_key, key], axis=-2)
- value = tf.concat([past_value, value], axis=-2)
-
- # to cope with keras serialization
- if use_cache:
- present = tf.stack([key, value], axis=0)
- else:
- present = (None,)
-
- attn_outputs = self._attn(query, key, value, attention_mask, head_mask, output_attentions, training=training)
- a = attn_outputs[0]
-
- a = self.merge_heads(a)
- a = self.c_proj(a)
- a = self.resid_dropout(a, training=training)
-
- outputs = [a, present] + attn_outputs[1:]
- return outputs # a, present, (attentions)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if self.is_cross_attention:
- c_attn_shape = 2 * self.embed_dim
- else:
- c_attn_shape = 3 * self.embed_dim
- if getattr(self, "c_proj", None) is not None:
- with tf.name_scope(self.c_proj.name):
- self.c_proj.build([None, None, self.embed_dim])
- if getattr(self, "c_attn", None) is not None:
- with tf.name_scope(self.c_attn.name):
- self.c_attn.build([None, None, c_attn_shape])
- if getattr(self, "q_attn", None) is not None:
- with tf.name_scope(self.q_attn.name):
- self.q_attn.build([None, None, self.embed_dim])
-
-
-class TFMLP(keras.layers.Layer):
- def __init__(self, n_state, config, **kwargs):
- super().__init__(**kwargs)
- nx = config.n_embd
- self.c_fc = TFConv1D(n_state, nx, initializer_range=config.initializer_range, name="c_fc")
- self.c_proj = TFConv1D(nx, n_state, initializer_range=config.initializer_range, name="c_proj")
- self.act = get_tf_activation(config.activation_function)
- self.dropout = keras.layers.Dropout(config.resid_pdrop)
- self.intermediate_size = n_state
- self.embed_dim = nx
-
- def call(self, x, training=False):
- h = self.act(self.c_fc(x))
- h2 = self.c_proj(h)
- h2 = self.dropout(h2, training=training)
- return h2
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "c_fc", None) is not None:
- with tf.name_scope(self.c_fc.name):
- self.c_fc.build([None, None, self.intermediate_size])
- if getattr(self, "c_proj", None) is not None:
- with tf.name_scope(self.c_proj.name):
- self.c_proj.build([None, None, self.embed_dim])
-
-
-class TFBlock(keras.layers.Layer):
- def __init__(self, config, scale=False, **kwargs):
- super().__init__(**kwargs)
- nx = config.n_embd
- inner_dim = config.n_inner if config.n_inner is not None else 4 * nx
- self.ln_1 = keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="ln_1")
- self.attn = TFAttention(nx, config, scale, name="attn")
- self.ln_2 = keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="ln_2")
-
- if config.add_cross_attention:
- self.crossattention = TFAttention(nx, config, scale, name="crossattention", is_cross_attention=True)
- self.ln_cross_attn = keras.layers.LayerNormalization(
- epsilon=config.layer_norm_epsilon, name="ln_cross_attn"
- )
-
- self.mlp = TFMLP(inner_dim, config, name="mlp")
- self.hidden_size = config.hidden_size
-
- def call(
- self,
- x,
- layer_past,
- attention_mask,
- head_mask,
- encoder_hidden_states,
- encoder_attention_mask,
- use_cache,
- output_attentions,
- training=False,
- ):
- a = self.ln_1(x)
- output_attn = self.attn(
- a,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- use_cache=use_cache,
- output_attentions=output_attentions,
- training=training,
- )
- a = output_attn[0] # output_attn: a, present, (attentions)
- outputs = output_attn[1:]
- x = x + a
-
- # Cross-Attention Block
- if encoder_hidden_states is not None:
- # add one self-attention block for cross-attention
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
- "cross-attention layers by setting `config.add_cross_attention=True`"
- )
-
- ca = self.ln_cross_attn(x)
- output_cross_attn = self.crossattention(
- ca,
- layer_past=None,
- attention_mask=attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- use_cache=False,
- output_attentions=output_attentions,
- training=training,
- )
- ca = output_cross_attn[0] # output_attn: a, present, (cross_attentions)
- x = x + ca
- outputs = outputs + output_cross_attn[2:] # add cross attentions if we output attention weights
-
- m = self.ln_2(x)
- m = self.mlp(m, training=training)
- x = x + m
-
- outputs = [x] + outputs
- return outputs # x, present, (attentions, cross_attentions)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "ln_1", None) is not None:
- with tf.name_scope(self.ln_1.name):
- self.ln_1.build([None, None, self.hidden_size])
- if getattr(self, "attn", None) is not None:
- with tf.name_scope(self.attn.name):
- self.attn.build(None)
- if getattr(self, "ln_2", None) is not None:
- with tf.name_scope(self.ln_2.name):
- self.ln_2.build([None, None, self.hidden_size])
- if getattr(self, "mlp", None) is not None:
- with tf.name_scope(self.mlp.name):
- self.mlp.build(None)
- if getattr(self, "crossattention", None) is not None:
- with tf.name_scope(self.crossattention.name):
- self.crossattention.build(None)
- if getattr(self, "ln_cross_attn", None) is not None:
- with tf.name_scope(self.ln_cross_attn.name):
- self.ln_cross_attn.build([None, None, self.hidden_size])
-
-
-@keras_serializable
-class TFGPT2MainLayer(keras.layers.Layer):
- config_class = GPT2Config
-
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- self.config = config
- self.output_attentions = config.output_attentions
- self.output_hidden_states = config.output_hidden_states
- self.use_cache = config.use_cache
- self.return_dict = config.use_return_dict
-
- self.num_hidden_layers = config.n_layer
- self.n_embd = config.n_embd
- self.n_positions = config.n_positions
- self.initializer_range = config.initializer_range
-
- self.wte = keras.layers.Embedding(
- input_dim=config.vocab_size,
- output_dim=config.hidden_size,
- embeddings_initializer=get_initializer(config.initializer_range),
- name="wte",
- )
- self.wpe = keras.layers.Embedding(
- input_dim=config.n_positions,
- output_dim=config.n_embd,
- embeddings_initializer=get_initializer(config.initializer_range),
- name="wpe",
- )
- self.drop = keras.layers.Dropout(config.embd_pdrop)
- self.h = [TFBlock(config, scale=True, name=f"h_._{i}") for i in range(config.n_layer)]
- self.ln_f = keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="ln_f")
- self.embed_dim = config.hidden_size
-
- def get_input_embeddings(self):
- return self.wte
-
- def set_input_embeddings(self, new_embeddings):
- self.wte = new_embeddings
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
- """
- raise NotImplementedError
-
- @unpack_inputs
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- encoder_hidden_states: np.ndarray | tf.Tensor | None = None,
- encoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithPastAndCrossAttentions, Tuple[tf.Tensor]]:
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- input_ids = tf.reshape(input_ids, [-1, input_shape[-1]])
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if past_key_values is None:
- past_length = 0
- past_key_values = [None] * len(self.h)
- else:
- past_length = shape_list(past_key_values[0][0])[-2]
-
- if position_ids is None:
- position_ids = tf.expand_dims(tf.range(past_length, input_shape[-1] + past_length), axis=0)
-
- if attention_mask is not None:
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask_shape = shape_list(attention_mask)
- attention_mask = tf.reshape(attention_mask, (attention_mask_shape[0], 1, 1, attention_mask_shape[1]))
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and -10000.0 for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- one_cst = tf.constant(1.0)
- attention_mask = tf.cast(attention_mask, dtype=one_cst.dtype)
- attention_mask = tf.multiply(tf.subtract(one_cst, attention_mask), tf.constant(-10000.0))
-
- # Copied from `modeling_tf_t5.py` with -1e9 -> -10000
- if self.config.add_cross_attention and encoder_attention_mask is not None:
- # If a 2D ou 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, mask_seq_length, mask_seq_length]
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- encoder_attention_mask = tf.cast(encoder_attention_mask, dtype=encoder_hidden_states.dtype)
- num_dims_encoder_attention_mask = len(shape_list(encoder_attention_mask))
- if num_dims_encoder_attention_mask == 3:
- encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
- if num_dims_encoder_attention_mask == 2:
- encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
-
- # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
- # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow/transformer/transformer_layers.py#L270
- # encoder_extended_attention_mask = tf.math.equal(encoder_extended_attention_mask,
- # tf.transpose(encoder_extended_attention_mask, perm=(-1, -2)))
-
- encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -10000.0
- else:
- encoder_extended_attention_mask = None
-
- encoder_attention_mask = encoder_extended_attention_mask
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.num_hidden_layers
- # head_mask = tf.constant([0] * self.num_hidden_layers)
-
- position_ids = tf.reshape(position_ids, [-1, shape_list(position_ids)[-1]])
-
- if inputs_embeds is None:
- check_embeddings_within_bounds(input_ids, self.config.vocab_size)
- inputs_embeds = self.wte(input_ids)
-
- position_embeds = self.wpe(position_ids)
-
- if token_type_ids is not None:
- token_type_ids = tf.reshape(token_type_ids, [-1, shape_list(token_type_ids)[-1]])
- token_type_embeds = self.wte(token_type_ids)
- else:
- token_type_embeds = tf.constant(0.0)
-
- position_embeds = tf.cast(position_embeds, dtype=inputs_embeds.dtype)
- token_type_embeds = tf.cast(token_type_embeds, dtype=inputs_embeds.dtype)
- hidden_states = inputs_embeds + position_embeds + token_type_embeds
- hidden_states = self.drop(hidden_states, training=training)
-
- output_shape = input_shape + [shape_list(hidden_states)[-1]]
-
- presents = () if use_cache else None
- all_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
- all_hidden_states = () if output_hidden_states else None
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (tf.reshape(hidden_states, output_shape),)
-
- outputs = block(
- hidden_states,
- layer_past,
- attention_mask,
- head_mask[i],
- encoder_hidden_states,
- encoder_attention_mask,
- use_cache,
- output_attentions,
- training=training,
- )
-
- hidden_states, present = outputs[:2]
- if use_cache:
- presents = presents + (present,)
-
- if output_attentions:
- all_attentions = all_attentions + (outputs[2],)
- if self.config.add_cross_attention and encoder_hidden_states is not None:
- all_cross_attentions = all_cross_attentions + (outputs[3],)
-
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = tf.reshape(hidden_states, output_shape)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if output_attentions:
- # let the number of heads free (-1) so we can extract attention even after head pruning
- attention_output_shape = input_shape[:-1] + [-1] + shape_list(all_attentions[0])[-2:]
- all_attentions = tuple(tf.reshape(t, attention_output_shape) for t in all_attentions)
-
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, presents, all_hidden_states, all_attentions, all_cross_attentions]
- if v is not None
- )
-
- return TFBaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- cross_attentions=all_cross_attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "wte", None) is not None:
- with tf.name_scope(self.wte.name):
- self.wte.build(None)
- if getattr(self, "wpe", None) is not None:
- with tf.name_scope(self.wpe.name):
- self.wpe.build(None)
- if getattr(self, "ln_f", None) is not None:
- with tf.name_scope(self.ln_f.name):
- self.ln_f.build([None, None, self.embed_dim])
- if getattr(self, "h", None) is not None:
- for layer in self.h:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-class TFGPT2PreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPT2Config
- base_model_prefix = "transformer"
- # names with a '.' represents the authorized unexpected/missing layers when a TF model is loaded from a PT model
- _keys_to_ignore_on_load_unexpected = [r"h.\d+.attn.bias", r"h.\d+.crossattention.bias"]
-
- @property
- def input_signature(self):
- # Although GPT-2 supports token_type_ids in theory, in practice they are rarely used, and the implementation
- # means that passing token_type_ids=0 yields different outputs from token_type_ids=None.
- # Therefore, we remove the token_type_ids argument by default, even though it would usually be included.
- return {
- "input_ids": tf.TensorSpec((None, None), tf.int32, name="input_ids"),
- "attention_mask": tf.TensorSpec((None, None), tf.int32, name="attention_mask"),
- }
-
-
-@dataclass
-class TFGPT2DoubleHeadsModelOutput(ModelOutput):
- """
- Base class for outputs of models predicting if two sentences are consecutive or not.
-
- Args:
- logits (`tf.Tensor` of shape `(batch_size, num_choices, sequence_length, config.vocab_size)`):
- Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
- mc_logits (`tf.Tensor` of shape `(batch_size, num_choices)`):
- Prediction scores of the multiple choice classification head (scores for each choice before SoftMax).
- past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads,
- sequence_length, embed_size_per_head)`).
-
- Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
- `past_key_values` input) to speed up sequential decoding.
- hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
- Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
- `(batch_size, sequence_length, hidden_size)`.
-
- Hidden-states of the model at the output of each layer plus the initial embedding outputs.
- attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
- Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
- sequence_length)`.
-
- Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
- heads.
- """
-
- logits: tf.Tensor = None
- mc_logits: tf.Tensor = None
- past_key_values: List[tf.Tensor] | None = None
- hidden_states: Tuple[tf.Tensor] | None = None
- attentions: Tuple[tf.Tensor] | None = None
-
-
-GPT2_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`GPT2Config`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPT2_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0].shape[-2]`
- (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
-
- If `past_key_values` is used, only input IDs that do not have their past calculated should be passed as
- `input_ids`.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`List[tf.Tensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past_key_values` output below). Can be used to speed up sequential decoding. The token ids which have
- their past given to this model should not be passed as input ids as they have already been computed.
- attention_mask (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- If `past_key_values` is used, `attention_mask` needs to contain the masking strategy that was used for
- `past_key_values`. In other words, the `attention_mask` always has to have the length:
- `len(past_key_values) + len(input_ids)`
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`Numpy array` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. This argument can be used in
- eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare GPT2 Model transformer outputting raw hidden-states without any specific head on top.",
- GPT2_START_DOCSTRING,
-)
-class TFGPT2Model(TFGPT2PreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFGPT2MainLayer(config, name="transformer")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutputWithPastAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- encoder_hidden_states: np.ndarray | tf.Tensor | None = None,
- encoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithPastAndCrossAttentions, Tuple[tf.Tensor]]:
- r"""
- encoder_hidden_states (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- past_key_values (`Tuple[Tuple[tf.Tensor]]` of length `config.n_layers`)
- contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
- If `past` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have
- their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*, defaults to `True`):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past`). Set to `False` during training, `True` during generation
- """
-
- outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
-
-
-@add_start_docstrings(
- """
- The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- GPT2_START_DOCSTRING,
-)
-class TFGPT2LMHeadModel(TFGPT2PreTrainedModel, TFCausalLanguageModelingLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFGPT2MainLayer(config, name="transformer")
-
- def get_output_embeddings(self):
- return self.get_input_embeddings()
-
- def set_output_embeddings(self, value):
- self.set_input_embeddings(value)
-
- def prepare_inputs_for_generation(self, inputs, past_key_values=None, use_cache=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # only last token for inputs_ids if past is defined in kwargs
- if past_key_values:
- inputs = tf.expand_dims(inputs[:, -1], -1)
- if token_type_ids is not None:
- token_type_ids = tf.expand_dims(token_type_ids[:, -1], -1)
-
- position_ids = kwargs.get("position_ids", None)
- attention_mask = kwargs.get("attention_mask", None)
-
- if attention_mask is not None and position_ids is None:
- position_ids = tf.math.cumsum(attention_mask, axis=-1, exclusive=True)
- if past_key_values:
- position_ids = tf.expand_dims(position_ids[:, -1], -1)
-
- return {
- "input_ids": inputs,
- "attention_mask": attention_mask,
- "position_ids": position_ids,
- "past_key_values": past_key_values,
- "use_cache": use_cache,
- "token_type_ids": token_type_ids,
- }
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFCausalLMOutputWithCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- encoder_hidden_states: np.ndarray | tf.Tensor | None = None,
- encoder_attention_mask: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFCausalLMOutputWithCrossAttentions, Tuple[tf.Tensor]]:
- r"""
- encoder_hidden_states (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
- the model is configured as a decoder.
- encoder_attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
- the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- past_key_values (`Tuple[Tuple[tf.Tensor]]` of length `config.n_layers`)
- contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
- If `past` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have
- their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*, defaults to `True`):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past`). Set to `False` during training, `True` during generation
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the cross entropy classification loss. Indices should be in `[0, ...,
- config.vocab_size - 1]`.
- """
-
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_states = transformer_outputs[0]
- logits = tf.matmul(hidden_states, self.transformer.wte.weights, transpose_b=True)
-
- loss = None
- if labels is not None:
- # shift labels to the left and cut last logit token
- shifted_logits = logits[:, :-1]
- labels = labels[:, 1:]
- loss = self.hf_compute_loss(labels, shifted_logits)
-
- if not return_dict:
- output = (logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFCausalLMOutputWithCrossAttentions(
- loss=loss,
- logits=logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- cross_attentions=transformer_outputs.cross_attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
-
-
-@add_start_docstrings(
- """
- The GPT2 Model transformer with a language modeling and a multiple-choice classification head on top e.g. for
- RocStories/SWAG tasks. The two heads are two linear layers. The language modeling head has its weights tied to the
- input embeddings, the classification head takes as input the input of a specified classification token index in the
- input sequence).
- """,
- GPT2_START_DOCSTRING,
-)
-class TFGPT2DoubleHeadsModel(TFGPT2PreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- config.num_labels = 1
- self.transformer = TFGPT2MainLayer(config, name="transformer")
- self.multiple_choice_head = TFSequenceSummary(
- config, initializer_range=config.initializer_range, name="multiple_choice_head"
- )
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- @replace_return_docstrings(output_type=TFGPT2DoubleHeadsModelOutput, config_class=_CONFIG_FOR_DOC)
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- mc_token_ids: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFGPT2DoubleHeadsModelOutput, Tuple[tf.Tensor]]:
- r"""
- mc_token_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, num_choices)`, *optional*, default to index of the last token of the input):
- Index of the classification token in each input sequence. Selected in the range `[0, input_ids.size(-1) -
- 1]`.
-
- Return:
-
- Examples:
-
- ```python
- >>> import tensorflow as tf
- >>> from transformers import AutoTokenizer, TFGPT2DoubleHeadsModel
-
- >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
- >>> model = TFGPT2DoubleHeadsModel.from_pretrained("openai-community/gpt2")
-
- >>> # Add a [CLS] to the vocabulary (we should train it also!)
- >>> num_added_tokens = tokenizer.add_special_tokens({"cls_token": "[CLS]"})
-
- >>> embedding_layer = model.resize_token_embeddings(
- ... len(tokenizer)
- ... ) # Update the model embeddings with the new vocabulary size
-
- >>> choices = ["Hello, my dog is cute [CLS]", "Hello, my cat is cute [CLS]"]
- >>> encoded_choices = [tokenizer.encode(s) for s in choices]
- >>> cls_token_location = [tokens.index(tokenizer.cls_token_id) for tokens in encoded_choices]
-
- >>> input_ids = tf.constant(encoded_choices)[None, :] # Batch size: 1, number of choices: 2
- >>> mc_token_ids = tf.constant([cls_token_location]) # Batch size: 1
-
- >>> outputs = model(input_ids, mc_token_ids=mc_token_ids)
- >>> lm_prediction_scores, mc_prediction_scores = outputs[:2]
- ```"""
-
- if input_ids is not None:
- input_shapes = shape_list(input_ids)
- else:
- input_shapes = shape_list(inputs_embeds)[:-1]
-
- seq_length = input_shapes[-1]
- flat_input_ids = tf.reshape(input_ids, (-1, seq_length)) if input_ids is not None else None
- flat_attention_mask = tf.reshape(attention_mask, (-1, seq_length)) if attention_mask is not None else None
- flat_token_type_ids = tf.reshape(token_type_ids, (-1, seq_length)) if token_type_ids is not None else None
- flat_position_ids = tf.reshape(position_ids, (-1, seq_length)) if position_ids is not None else None
- transformer_outputs = self.transformer(
- input_ids=flat_input_ids,
- past_key_values=past_key_values,
- attention_mask=flat_attention_mask,
- token_type_ids=flat_token_type_ids,
- position_ids=flat_position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=None,
- encoder_attention_mask=None,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_states = transformer_outputs[0]
- hidden_states = tf.reshape(hidden_states, input_shapes + shape_list(hidden_states)[-1:])
- if return_dict and output_hidden_states:
- # We do this to match the slightly odd PT behaviour - the final hidden state is reshaped to rank 4 when the
- # input is rank 3, but all other hidden states remain at rank-3 (with the first 2 dims merged)
- all_hidden_states = transformer_outputs.hidden_states[:-1] + (hidden_states,)
- else:
- all_hidden_states = None
- lm_logits = tf.matmul(hidden_states, self.transformer.wte.weights, transpose_b=True)
- mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids, training=training)
- mc_logits = tf.squeeze(mc_logits, axis=-1)
-
- if not return_dict:
- return (lm_logits, mc_logits) + transformer_outputs[1:]
-
- return TFGPT2DoubleHeadsModelOutput(
- logits=lm_logits,
- mc_logits=mc_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=all_hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- @property
- def input_signature(self):
- return {
- "input_ids": tf.TensorSpec((None, None, None), tf.int32, name="input_ids"),
- "attention_mask": tf.TensorSpec((None, None, None), tf.int32, name="attention_mask"),
- "mc_token_ids": tf.TensorSpec((None, None), tf.int32, name="mc_token_ids"),
- }
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "multiple_choice_head", None) is not None:
- with tf.name_scope(self.multiple_choice_head.name):
- self.multiple_choice_head.build(None)
-
-
-@add_start_docstrings(
- """
- The GPT2 Model transformer with a sequence classification head on top (linear layer).
-
- [`TFGPT2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT-1) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- GPT2_START_DOCSTRING,
-)
-class TFGPT2ForSequenceClassification(TFGPT2PreTrainedModel, TFSequenceClassificationLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
- self.score = keras.layers.Dense(
- config.num_labels,
- kernel_initializer=get_initializer(config.initializer_range),
- name="score",
- use_bias=False,
- )
- self.transformer = TFGPT2MainLayer(config, name="transformer")
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint="microsoft/DialogRPT-updown",
- output_type=TFSequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- labels: np.ndarray | tf.Tensor | None = None,
- training: Optional[bool] = False,
- ) -> Union[TFSequenceClassifierOutputWithPast, Tuple[tf.Tensor]]:
- r"""
- labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the cross entropy classification loss. Indices should be in `[0, ...,
- config.vocab_size - 1]`.
- """
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
- logits_shape = shape_list(logits)
- in_logits = None
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- sequence_lengths = (
- tf.argmax(tf.cast(tf.math.equal(input_ids, self.config.pad_token_id), input_ids.dtype), axis=-1)
- - 1
- )
- sequence_lengths = tf.where(sequence_lengths >= 0, sequence_lengths, input_ids.shape[-1] - 1)
- in_logits = tf.gather(logits, sequence_lengths, batch_dims=1, axis=1)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
- loss = None
-
- if labels is not None:
- assert (
- self.config.pad_token_id is not None or logits_shape[0] == 1
- ), "Cannot handle batch sizes > 1 if no padding token is defined."
-
- if not tf.is_tensor(sequence_lengths):
- in_logits = logits[0 : logits_shape[0], sequence_lengths]
-
- loss = self.hf_compute_loss(tf.reshape(labels, [-1]), tf.reshape(in_logits, [-1, self.num_labels]))
- pooled_logits = in_logits if in_logits is not None else logits
-
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "score", None) is not None:
- with tf.name_scope(self.score.name):
- self.score.build([None, None, self.config.n_embd])
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
diff --git a/transformers/models/gpt2/tokenization_gpt2.py b/transformers/models/gpt2/tokenization_gpt2.py
deleted file mode 100644
index 36f3ca8fadb527f76359ac6bd8f65989ad631992..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/tokenization_gpt2.py
+++ /dev/null
@@ -1,345 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for OpenAI GPT."""
-
-
-import json
-import os
-from functools import lru_cache
-from typing import List, Optional, Tuple
-
-import regex as re
-
-from ...tokenization_utils import AddedToken, PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {
- "vocab_file": "vocab.json",
- "merges_file": "merges.txt",
-}
-
-
-@lru_cache()
-def bytes_to_unicode():
- """
- Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
- characters the bpe code barfs on.
-
- The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
- if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
- decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
- tables between utf-8 bytes and unicode strings.
- """
- bs = (
- list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
- )
- cs = bs[:]
- n = 0
- for b in range(2**8):
- if b not in bs:
- bs.append(b)
- cs.append(2**8 + n)
- n += 1
- cs = [chr(n) for n in cs]
- return dict(zip(bs, cs))
-
-
-def get_pairs(word):
- """
- Return set of symbol pairs in a word.
-
- Word is represented as tuple of symbols (symbols being variable-length strings).
- """
- pairs = set()
- prev_char = word[0]
- for char in word[1:]:
- pairs.add((prev_char, char))
- prev_char = char
- return pairs
-
-
-class GPT2Tokenizer(PreTrainedTokenizer):
- """
- Construct a GPT-2 tokenizer. Based on byte-level Byte-Pair-Encoding.
-
- This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
- be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
- ```python
- >>> from transformers import GPT2Tokenizer
-
- >>> tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
- >>> tokenizer("Hello world")["input_ids"]
- [15496, 995]
-
- >>> tokenizer(" Hello world")["input_ids"]
- [18435, 995]
- ```
-
- You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer or when you
- call it on some text, but since the model was not pretrained this way, it might yield a decrease in performance.
-
-
-
- When used with `is_split_into_words=True`, this tokenizer will add a space before each word (even the first one).
-
-
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- merges_file (`str`):
- Path to the merges file.
- errors (`str`, *optional*, defaults to `"replace"`):
- Paradigm to follow when decoding bytes to UTF-8. See
- [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
- unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The end of sequence token.
- pad_token (`str`, *optional*):
- The token used for padding, for example when batching sequences of different lengths.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial space to the input. This allows to treat the leading word just as any
- other word. (GPT2 tokenizer detect beginning of words by the preceding space).
- add_bos_token (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial beginning of sentence token to the input. This allows to treat the leading
- word just as any other word.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- merges_file,
- errors="replace",
- unk_token="<|endoftext|>",
- bos_token="<|endoftext|>",
- eos_token="<|endoftext|>",
- pad_token=None,
- add_prefix_space=False,
- add_bos_token=False,
- **kwargs,
- ):
- bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
- eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
- unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
- pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
-
- self.add_bos_token = add_bos_token
-
- with open(vocab_file, encoding="utf-8") as vocab_handle:
- self.encoder = json.load(vocab_handle)
- self.decoder = {v: k for k, v in self.encoder.items()}
- self.errors = errors # how to handle errors in decoding
- self.byte_encoder = bytes_to_unicode()
- self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
- with open(merges_file, encoding="utf-8") as merges_handle:
- bpe_merges = merges_handle.read().split("\n")[1:-1]
- bpe_merges = [tuple(merge.split()) for merge in bpe_merges]
- self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
- self.cache = {}
- self.add_prefix_space = add_prefix_space
-
- # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
- self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
-
- super().__init__(
- errors=errors,
- unk_token=unk_token,
- bos_token=bos_token,
- eos_token=eos_token,
- pad_token=pad_token,
- add_prefix_space=add_prefix_space,
- add_bos_token=add_bos_token,
- **kwargs,
- )
-
- @property
- def vocab_size(self):
- return len(self.encoder)
-
- def get_vocab(self):
- return dict(self.encoder, **self.added_tokens_encoder)
-
- def bpe(self, token):
- if token in self.cache:
- return self.cache[token]
- word = tuple(token)
- pairs = get_pairs(word)
-
- if not pairs:
- return token
-
- while True:
- bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
- if bigram not in self.bpe_ranks:
- break
- first, second = bigram
- new_word = []
- i = 0
- while i < len(word):
- try:
- j = word.index(first, i)
- except ValueError:
- new_word.extend(word[i:])
- break
- else:
- new_word.extend(word[i:j])
- i = j
-
- if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
- new_word.append(first + second)
- i += 2
- else:
- new_word.append(word[i])
- i += 1
- new_word = tuple(new_word)
- word = new_word
- if len(word) == 1:
- break
- else:
- pairs = get_pairs(word)
- word = " ".join(word)
- self.cache[token] = word
- return word
-
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- if self.add_bos_token:
- bos_token_ids = [self.bos_token_id]
- else:
- bos_token_ids = []
-
- output = bos_token_ids + token_ids_0
-
- if token_ids_1 is None:
- return output
-
- return output + bos_token_ids + token_ids_1
-
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- if not self.add_bos_token:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=False
- )
-
- if token_ids_1 is None:
- return [1] + ([0] * len(token_ids_0))
- return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1))
-
- def _tokenize(self, text):
- """Tokenize a string."""
- bpe_tokens = []
- for token in re.findall(self.pat, text):
- token = "".join(
- self.byte_encoder[b] for b in token.encode("utf-8")
- ) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
- bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
- return bpe_tokens
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.encoder.get(token, self.encoder.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.decoder.get(index)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- text = "".join(tokens)
- text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
- return text
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- merge_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
- )
-
- with open(vocab_file, "w", encoding="utf-8") as f:
- f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
- index = 0
- with open(merge_file, "w", encoding="utf-8") as writer:
- writer.write("#version: 0.2\n")
- for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
- " Please check that the tokenizer is not corrupted!"
- )
- index = token_index
- writer.write(" ".join(bpe_tokens) + "\n")
- index += 1
-
- return vocab_file, merge_file
-
- def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
- add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space)
- if is_split_into_words or add_prefix_space:
- text = " " + text
- return (text, kwargs)
-
- @property
- def default_chat_template(self):
- """
- A simple chat template that ignores role information and just concatenates messages with EOS tokens.
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- return "{% for message in messages %}" "{{ message.content }}{{ eos_token }}" "{% endfor %}"
diff --git a/transformers/models/gpt2/tokenization_gpt2_fast.py b/transformers/models/gpt2/tokenization_gpt2_fast.py
deleted file mode 100644
index fb3a5d4a0ce3f21dd7c7cb31ff4a740d8cdd8037..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/tokenization_gpt2_fast.py
+++ /dev/null
@@ -1,156 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for OpenAI GPT."""
-
-
-import json
-from typing import Optional, Tuple
-
-from tokenizers import pre_tokenizers
-
-from ...tokenization_utils_base import BatchEncoding
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import logging
-from .tokenization_gpt2 import GPT2Tokenizer
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}
-
-
-class GPT2TokenizerFast(PreTrainedTokenizerFast):
- """
- Construct a "fast" GPT-2 tokenizer (backed by HuggingFace's *tokenizers* library). Based on byte-level
- Byte-Pair-Encoding.
-
- This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
- be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
- ```python
- >>> from transformers import GPT2TokenizerFast
-
- >>> tokenizer = GPT2TokenizerFast.from_pretrained("openai-community/gpt2")
- >>> tokenizer("Hello world")["input_ids"]
- [15496, 995]
-
- >>> tokenizer(" Hello world")["input_ids"]
- [18435, 995]
- ```
-
- You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer, but since
- the model was not pretrained this way, it might yield a decrease in performance.
-
-
-
- When used with `is_split_into_words=True`, this tokenizer needs to be instantiated with `add_prefix_space=True`.
-
-
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`, *optional*):
- Path to the vocabulary file.
- merges_file (`str`, *optional*):
- Path to the merges file.
- tokenizer_file (`str`, *optional*):
- Path to [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that
- contains everything needed to load the tokenizer.
- unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The end of sequence token.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial space to the input. This allows to treat the leading word just as any
- other word. (GPT2 tokenizer detect beginning of words by the preceding space).
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
- slow_tokenizer_class = GPT2Tokenizer
-
- def __init__(
- self,
- vocab_file=None,
- merges_file=None,
- tokenizer_file=None,
- unk_token="<|endoftext|>",
- bos_token="<|endoftext|>",
- eos_token="<|endoftext|>",
- add_prefix_space=False,
- **kwargs,
- ):
- super().__init__(
- vocab_file,
- merges_file,
- tokenizer_file=tokenizer_file,
- unk_token=unk_token,
- bos_token=bos_token,
- eos_token=eos_token,
- add_prefix_space=add_prefix_space,
- **kwargs,
- )
-
- self.add_bos_token = kwargs.pop("add_bos_token", False)
-
- pre_tok_state = json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__())
- if pre_tok_state.get("add_prefix_space", add_prefix_space) != add_prefix_space:
- pre_tok_class = getattr(pre_tokenizers, pre_tok_state.pop("type"))
- pre_tok_state["add_prefix_space"] = add_prefix_space
- self.backend_tokenizer.pre_tokenizer = pre_tok_class(**pre_tok_state)
-
- self.add_prefix_space = add_prefix_space
-
- def _batch_encode_plus(self, *args, **kwargs) -> BatchEncoding:
- is_split_into_words = kwargs.get("is_split_into_words", False)
- assert self.add_prefix_space or not is_split_into_words, (
- f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
- "to use it with pretokenized inputs."
- )
-
- return super()._batch_encode_plus(*args, **kwargs)
-
- def _encode_plus(self, *args, **kwargs) -> BatchEncoding:
- is_split_into_words = kwargs.get("is_split_into_words", False)
-
- assert self.add_prefix_space or not is_split_into_words, (
- f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
- "to use it with pretokenized inputs."
- )
-
- return super()._encode_plus(*args, **kwargs)
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
-
- @property
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.default_chat_template
- def default_chat_template(self):
- """
- A simple chat template that ignores role information and just concatenates messages with EOS tokens.
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- return "{% for message in messages %}" "{{ message.content }}{{ eos_token }}" "{% endfor %}"
diff --git a/transformers/models/gpt2/tokenization_gpt2_tf.py b/transformers/models/gpt2/tokenization_gpt2_tf.py
deleted file mode 100644
index d763eb848550157528f2fab408c54139c83b9e30..0000000000000000000000000000000000000000
--- a/transformers/models/gpt2/tokenization_gpt2_tf.py
+++ /dev/null
@@ -1,104 +0,0 @@
-import os
-from typing import Dict, List, Union
-
-import tensorflow as tf
-from keras_nlp.tokenizers import BytePairTokenizer
-from tensorflow_text import pad_model_inputs
-
-from ...modeling_tf_utils import keras
-from .tokenization_gpt2 import GPT2Tokenizer
-
-
-class TFGPT2Tokenizer(keras.layers.Layer):
- """
- This is an in-graph tokenizer for GPT2. It should be initialized similarly to other tokenizers, using the
- `from_pretrained()` method. It can also be initialized with the `from_tokenizer()` method, which imports settings
- from an existing standard tokenizer object.
-
- In-graph tokenizers, unlike other Hugging Face tokenizers, are actually Keras layers and are designed to be run
- when the model is called, rather than during preprocessing. As a result, they have somewhat more limited options
- than standard tokenizer classes. They are most useful when you want to create an end-to-end model that goes
- straight from `tf.string` inputs to outputs.
-
- Args:
- vocab (Dict[str, int]): Vocabulary dict for Byte Pair Tokenizer
- merges (List[str]): Merges list for Byte Pair Tokenizer
- """
-
- def __init__(self, vocab: Dict[str, int], merges: List[str], max_length: int = None, pad_token_id: int = None):
- super().__init__()
- self.pad_token_id = pad_token_id
- self.max_length = max_length
- self.vocab = vocab
- self.merges = merges
- self.tf_tokenizer = BytePairTokenizer(vocab, merges, sequence_length=max_length)
-
- @classmethod
- def from_tokenizer(cls, tokenizer: GPT2Tokenizer, *args, **kwargs):
- """Creates TFGPT2Tokenizer from GPT2Tokenizer
-
- Args:
- tokenizer (GPT2Tokenizer)
-
- Examples:
-
- ```python
- from transformers import AutoTokenizer, TFGPT2Tokenizer
-
- tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
- tf_tokenizer = TFGPT2Tokenizer.from_tokenizer(tokenizer)
- ```
- """
- merges = [" ".join(m) for m in tokenizer.bpe_ranks.keys()]
- vocab = tokenizer.get_vocab()
- return cls(vocab, merges, *args, **kwargs)
-
- @classmethod
- def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], *init_inputs, **kwargs):
- """Creates TFGPT2Tokenizer from pretrained GPT2Tokenizer
-
- Args:
- pretrained_model_name_or_path (Union[str, os.PathLike]): Path to pretrained model
-
- Examples:
-
- ```python
- from transformers import TFGPT2Tokenizer
-
- tf_tokenizer = TFGPT2Tokenizer.from_pretrained("openai-community/gpt2")
- ```
- """
- tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)
- return cls.from_tokenizer(tokenizer, *init_inputs, **kwargs)
-
- @classmethod
- def from_config(cls, config):
- """Creates TFGPT2Tokenizer from configurations
-
- Args:
- config (Dict): Dictionary with keys such as stated in `get_config`.
- """
- return cls(**config)
-
- def get_config(self):
- return {
- "vocab": self.vocab,
- "merges": self.merges,
- "max_length": self.max_length,
- "pad_token_id": self.pad_token_id,
- }
-
- def call(self, x, max_length: int = None):
- input_ids = self.tf_tokenizer(x)
- attention_mask = tf.ones_like(input_ids)
-
- if self.pad_token_id is not None:
- # pad the tokens up to max length
- max_length = max_length if max_length is not None else self.max_length
-
- if max_length is not None:
- input_ids, attention_mask = pad_model_inputs(
- input_ids, max_seq_length=max_length, pad_value=self.pad_token_id
- )
-
- return {"attention_mask": attention_mask, "input_ids": input_ids}
diff --git a/transformers/models/gpt_bigcode/__init__.py b/transformers/models/gpt_bigcode/__init__.py
deleted file mode 100644
index 33660eb81e4faebb7938bbba7ba165a2d7079d81..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_bigcode/__init__.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_gpt_bigcode": ["GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTBigCodeConfig"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_gpt_bigcode"] = [
- "GPT_BIGCODE_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GPTBigCodeForSequenceClassification",
- "GPTBigCodeForTokenClassification",
- "GPTBigCodeForCausalLM",
- "GPTBigCodeModel",
- "GPTBigCodePreTrainedModel",
- ]
-
-if TYPE_CHECKING:
- from .configuration_gpt_bigcode import GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTBigCodeConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_gpt_bigcode import (
- GPT_BIGCODE_PRETRAINED_MODEL_ARCHIVE_LIST,
- GPTBigCodeForCausalLM,
- GPTBigCodeForSequenceClassification,
- GPTBigCodeForTokenClassification,
- GPTBigCodeModel,
- GPTBigCodePreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gpt_bigcode/__pycache__/__init__.cpython-310.pyc b/transformers/models/gpt_bigcode/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index fb064049f84cba850ae3d64a7c2c996f72e0e67e..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_bigcode/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_bigcode/__pycache__/configuration_gpt_bigcode.cpython-310.pyc b/transformers/models/gpt_bigcode/__pycache__/configuration_gpt_bigcode.cpython-310.pyc
deleted file mode 100644
index e2abd70b663affc510df36180cbde27a0791e6ff..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_bigcode/__pycache__/configuration_gpt_bigcode.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_bigcode/__pycache__/modeling_gpt_bigcode.cpython-310.pyc b/transformers/models/gpt_bigcode/__pycache__/modeling_gpt_bigcode.cpython-310.pyc
deleted file mode 100644
index 66014702324ea43a0eaf7321f917f8e1c5b192d5..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_bigcode/__pycache__/modeling_gpt_bigcode.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py b/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py
deleted file mode 100644
index ef5e02ffdc43af1a72f7daaff59f18a3f95b5000..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_bigcode/configuration_gpt_bigcode.py
+++ /dev/null
@@ -1,144 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The BigCode team and HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPTBigCode configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GPTBigCodeConfig(PretrainedConfig):
- """
- This is the configuration class to store the configuration of a [`GPTBigCodeModel`]. It is used to instantiate a
- GPTBigCode model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the GPTBigCode
- [gpt_bigcode](https://huggingface.co/gpt_bigcode) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 50257):
- Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`GPTBigCodeModel`].
- n_positions (`int`, *optional*, defaults to 1024):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- n_embd (`int`, *optional*, defaults to 768):
- Dimensionality of the embeddings and hidden states.
- n_layer (`int`, *optional*, defaults to 12):
- Number of hidden layers in the Transformer encoder.
- n_head (`int`, *optional*, defaults to 12):
- Number of attention heads for each attention layer in the Transformer encoder.
- n_inner (`int`, *optional*, defaults to None):
- Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
- activation_function (`str`, *optional*, defaults to `"gelu_pytorch_tanh"`):
- Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new",
- "gelu_pytorch_tanh"]`.
- resid_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- embd_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the embeddings.
- attn_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention.
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
- The epsilon to use in the layer normalization layers.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- scale_attn_weights (`bool`, *optional*, defaults to `True`):
- Scale attention weights by dividing by sqrt(hidden_size)..
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
- attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`):
- Whether to call the fused softmax in float32.
- scale_attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`):
- Whether to scale the attention softmax in float32.
- attention_type (`bool`, *optional*, defaults to `True`):
- Whether to use Multi-Query Attion (`True`) or Multi-Head Attention (`False`).
- Example:
-
- ```python
- >>> from transformers import GPTBigCodeConfig, GPTBigCodeModel
-
- >>> # Initializing a GPTBigCode configuration
- >>> configuration = GPTBigCodeConfig()
-
- >>> # Initializing a model (with random weights) from the configuration
- >>> model = GPTBigCodeModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "gpt_bigcode"
- keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {
- "hidden_size": "n_embd",
- "max_position_embeddings": "n_positions",
- "num_attention_heads": "n_head",
- "num_hidden_layers": "n_layer",
- }
-
- def __init__(
- self,
- vocab_size=50257,
- n_positions=1024,
- n_embd=768,
- n_layer=12,
- n_head=12,
- n_inner=None,
- activation_function="gelu_pytorch_tanh",
- resid_pdrop=0.1,
- embd_pdrop=0.1,
- attn_pdrop=0.1,
- layer_norm_epsilon=1e-5,
- initializer_range=0.02,
- scale_attn_weights=True,
- use_cache=True,
- bos_token_id=50256,
- eos_token_id=50256,
- attention_softmax_in_fp32=True,
- scale_attention_softmax_in_fp32=True,
- multi_query=True,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.n_positions = n_positions
- self.n_embd = n_embd
- self.n_layer = n_layer
- self.n_head = n_head
- self.n_inner = n_inner
- self.activation_function = activation_function
- self.resid_pdrop = resid_pdrop
- self.embd_pdrop = embd_pdrop
- self.attn_pdrop = attn_pdrop
- self.layer_norm_epsilon = layer_norm_epsilon
- self.initializer_range = initializer_range
- self.scale_attn_weights = scale_attn_weights
- self.use_cache = use_cache
- self.attention_softmax_in_fp32 = attention_softmax_in_fp32
- self.scale_attention_softmax_in_fp32 = scale_attention_softmax_in_fp32
- self.multi_query = multi_query
-
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
-
- super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
diff --git a/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py b/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py
deleted file mode 100644
index 4e3b8498480c9e76a087b5fb4339a83d1b809710..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py
+++ /dev/null
@@ -1,1504 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The Bigcode team and HuggingFace Inc. team.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch GPTBigCode model."""
-import math
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_attn_mask_utils import AttentionMaskConverter
-from ...modeling_outputs import (
- BaseModelOutputWithPastAndCrossAttentions,
- CausalLMOutputWithCrossAttentions,
- SequenceClassifierOutputWithPast,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import is_torch_greater_or_equal_than_2_2
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- logging,
-)
-from .configuration_gpt_bigcode import GPTBigCodeConfig
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "bigcode/gpt_bigcode-santacoder"
-_CONFIG_FOR_DOC = "GPTBigCodeConfig"
-
-
-from ..deprecated._archive_maps import GPT_BIGCODE_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Fused kernels
-# Use separate functions for each case because conditionals prevent kernel fusion.
-# TODO: Could have better fused kernels depending on scaling, dropout and head mask.
-# Is it doable without writing 32 functions?
-@torch.jit.script
-def upcast_masked_softmax(
- x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor, scale: float, softmax_dtype: torch.dtype
-):
- input_dtype = x.dtype
- x = x.to(softmax_dtype) * scale
- x = torch.where(mask, x, mask_value)
- x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype)
- return x
-
-
-@torch.jit.script
-def upcast_softmax(x: torch.Tensor, scale: float, softmax_dtype: torch.dtype):
- input_dtype = x.dtype
- x = x.to(softmax_dtype) * scale
- x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype)
- return x
-
-
-@torch.jit.script
-def masked_softmax(x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor):
- x = torch.where(mask, x, mask_value)
- x = torch.nn.functional.softmax(x, dim=-1)
- return x
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-class GPTBigCodeAttention(nn.Module):
- def __init__(self, config, is_cross_attention=False, layer_idx=None):
- super().__init__()
- self.config = config
-
- self.mask_value = None
- self.multi_query = config.multi_query
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_heads
- self.kv_heads = 1 if self.multi_query else self.num_heads
- self.kv_dim = self.kv_heads * self.head_dim
- self.split_size = self.embed_dim
- self.is_causal = True
-
- if self.head_dim * self.num_heads != self.embed_dim:
- raise ValueError(
- f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {self.num_heads})."
- )
-
- self.scale_attn_weights = config.scale_attn_weights
- self.is_cross_attention = is_cross_attention
-
- self.layer_idx = layer_idx
- self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
- self.scale_attention_softmax_in_fp32 = (
- config.scale_attention_softmax_in_fp32 and config.attention_softmax_in_fp32
- )
- self.attn_pdrop = config.attn_pdrop
-
- if self.is_cross_attention:
- if self.multi_query:
- raise NotImplementedError("Multi-Query Attention not supported for cross_attention")
-
- self.c_attn = nn.Linear(self.embed_dim, 2 * self.embed_dim)
- self.q_attn = nn.Linear(self.embed_dim, self.embed_dim)
- else:
- self.c_attn = nn.Linear(self.embed_dim, self.embed_dim + 2 * self.kv_dim)
-
- self.c_proj = nn.Linear(self.embed_dim, self.embed_dim)
-
- self.attn_dropout = nn.Dropout(config.attn_pdrop)
- self.resid_dropout = nn.Dropout(config.resid_pdrop)
-
- def _get_mask_value(self, device, dtype):
- # torch.where expects a tensor. We use a cache to avoid recreating it every time.
- if self.mask_value is None or self.mask_value.dtype != dtype or self.mask_value.device != device:
- self.mask_value = torch.full([], torch.finfo(dtype).min, dtype=dtype, device=device)
- return self.mask_value
-
- def _attn(self, query, key, value, attention_mask=None, head_mask=None):
- dtype = query.dtype
- softmax_dtype = torch.float32 if self.attention_softmax_in_fp32 else dtype
- upcast = dtype != softmax_dtype
-
- unscale = self.layer_idx + 1 if self.scale_attention_softmax_in_fp32 and upcast else 1
- scale_factor = unscale**-1
- if self.scale_attn_weights:
- scale_factor /= self.head_dim**0.5
-
- # MQA models: (batch_size, query_length, num_heads * head_dim)
- # MHA models: (batch_size, num_heads, query_length, head_dim)
- query_shape = query.shape
- batch_size = query_shape[0]
- key_length = key.size(-1)
- if self.multi_query:
- # (batch_size, query_length, num_heads, head_dim) x (batch_size, head_dim, key_length)
- # -> (batch_size, query_length, num_heads, key_length)
- query_length = query_shape[1]
- attn_shape = (batch_size, query_length, self.num_heads, key_length)
- attn_view = (batch_size, query_length * self.num_heads, key_length)
- # No copy needed for MQA 2, or when layer_past is provided.
- query = query.reshape(batch_size, query_length * self.num_heads, self.head_dim)
- else:
- # (batch_size, num_heads, query_length, head_dim) x (batch_size, num_heads, head_dim, key_length)
- # -> (batch_size, num_heads, query_length, key_length)
- query_length = query_shape[2]
- attn_shape = (batch_size, self.num_heads, query_length, key_length)
- attn_view = (batch_size * self.num_heads, query_length, key_length)
- # Always copies
- query = query.reshape(batch_size * self.num_heads, query_length, self.head_dim)
- # No copy when layer_past is provided.
- key = key.reshape(batch_size * self.num_heads, self.head_dim, key_length)
-
- attn_weights = torch.empty(attn_view, device=query.device, dtype=query.dtype)
- if query.device.type == "cpu":
- # This is needed because of a bug in pytorch https://github.com/pytorch/pytorch/issues/80588.
- # The bug was fixed in https://github.com/pytorch/pytorch/pull/96086,
- # but the fix has not been released as of pytorch version 2.0.0.
- attn_weights = torch.zeros_like(attn_weights)
- beta = 1
- else:
- beta = 0
- attn_weights = torch.baddbmm(attn_weights, query, key, beta=beta, alpha=scale_factor).view(attn_shape)
-
- if upcast:
- # Use a fused kernel to prevent a large overhead from casting and scaling.
- # Sub-optimal when the key length is not a multiple of 8.
- if attention_mask is None:
- attn_weights = upcast_softmax(attn_weights, unscale, softmax_dtype)
- else:
- mask_value = self._get_mask_value(attn_weights.device, softmax_dtype)
- attn_weights = upcast_masked_softmax(attn_weights, attention_mask, mask_value, unscale, softmax_dtype)
- else:
- if attention_mask is not None:
- mask_value = self._get_mask_value(attn_weights.device, softmax_dtype)
-
- # The fused kernel is very slow when the key length is not a multiple of 8, so we skip fusion.
- attn_weights = torch.where(attention_mask, attn_weights, mask_value)
-
- attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
-
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- if self.multi_query:
- head_mask = head_mask.transpose(1, 2)
- attn_weights = attn_weights * head_mask
-
- if self.multi_query:
- attn_output = torch.bmm(attn_weights.view(attn_view), value).view(query_shape)
- else:
- attn_output = torch.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- layer_past: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[
- Tuple[torch.Tensor, Optional[torch.Tensor]],
- Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[torch.Tensor, ...]],
- ]:
- if encoder_hidden_states is not None:
- if not hasattr(self, "q_attn") or not self.is_cross_attention:
- raise ValueError(
- "If class is used as cross attention, the weights `q_attn` have to be defined. "
- "Please make sure to instantiate class with `GPTBigCodeAttention(..., is_cross_attention=True)`."
- )
-
- query = self.q_attn(hidden_states)
- key_value = self.c_attn(encoder_hidden_states)
- attention_mask = encoder_attention_mask
- elif self.multi_query:
- query, key_value = self.c_attn(hidden_states).split((self.embed_dim, 2 * self.kv_dim), dim=2)
- else:
- # Note: We split as (self.num_heads, 3, self.head_dim) instead of (3, self.num_heads, self.head_dim),
- # i.e., the memory layout is not the same as GPT2.
- # This makes the concatenation with past_key_value more efficient.
- query, key_value = (
- self.c_attn(hidden_states)
- .view(*hidden_states.shape[:2], self.num_heads, 3 * self.head_dim)
- .transpose(1, 2)
- .split((self.head_dim, 2 * self.head_dim), dim=3)
- )
-
- if layer_past is not None:
- key_value = torch.cat((layer_past, key_value), dim=-2)
- present = key_value if use_cache else None
-
- key, value = key_value.split((self.head_dim, self.head_dim), dim=-1)
-
- attn_output, attn_weights = self._attn(query, key.transpose(-1, -2), value, attention_mask, head_mask)
-
- if not self.multi_query:
- attn_output = attn_output.transpose(1, 2).reshape(hidden_states.shape)
- attn_output = self.c_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- if self.multi_query:
- # Transpose to return weights in the usual format (batch_size, num_heads, query_length, key_length)
- attn_weights = attn_weights.transpose(1, 2)
- outputs += (attn_weights,)
-
- return outputs # a, present, (attentions)
-
-
-class GPTBigCodeFlashAttention2(GPTBigCodeAttention):
- """
- GPTBigCode flash attention module. This module inherits from `GPTBigCodeAttention` as the weights of the module
- stays untouched. The only required change would be on the forward pass where it needs to correctly call the public
- API of flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- layer_past: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[
- Tuple[torch.Tensor, Optional[torch.Tensor]],
- Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[torch.Tensor, ...]],
- ]:
- if encoder_hidden_states is not None:
- if not hasattr(self, "q_attn") or not self.is_cross_attention:
- raise ValueError(
- "If class is used as cross attention, the weights `q_attn` have to be defined. "
- "Please make sure to instantiate class with `GPTBigCodeAttention(..., is_cross_attention=True)`."
- )
-
- query = self.q_attn(hidden_states)
- key_value = self.c_attn(encoder_hidden_states)
- attention_mask = encoder_attention_mask
- elif self.multi_query:
- query, key_value = self.c_attn(hidden_states).split((self.embed_dim, 2 * self.kv_dim), dim=2)
- else:
- # Note: We split as (self.num_heads, 3, self.head_dim) instead of (3, self.num_heads, self.head_dim),
- # i.e., the memory layout is not the same as GPT2.
- # This makes the concatenation with past_key_value more efficient.
- query, key_value = (
- self.c_attn(hidden_states)
- .view(*hidden_states.shape[:2], self.num_heads, 3 * self.head_dim)
- .transpose(1, 2)
- .split((self.head_dim, 2 * self.head_dim), dim=3)
- )
-
- if layer_past is not None:
- key_value = torch.cat((layer_past, key_value), dim=-2)
- present = key_value if use_cache else None
-
- key, value = key_value.split((self.head_dim, self.head_dim), dim=-1)
-
- # Flash attention requires the input to have the shape
- # batch_size x seq_length x head_dim x hidden_dim
- if self.multi_query:
- batch_size, query_length, _ = query.shape
- query = query.reshape(batch_size, query_length, self.num_heads, self.head_dim)
- key = key.unsqueeze(2)
- value = value.unsqueeze(2)
- else:
- query_length = query.shape[2]
- batch_size, _, tgt, _ = key.shape
- query = query.transpose(1, 2).reshape(batch_size, query_length, self.num_heads, self.head_dim)
- key = key.transpose(1, 2).reshape(batch_size, tgt, self.num_heads, self.head_dim)
- value = value.transpose(1, 2).reshape(batch_size, tgt, self.num_heads, self.head_dim)
-
- attn_dropout = self.attn_pdrop if self.training else 0.0
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in float16 just to be sure everything works as expected.
- input_dtype = query.dtype
- if input_dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.c_attn.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
- query = query.to(target_dtype)
- key = key.to(target_dtype)
- value = value.to(target_dtype)
-
- attn_output = self._flash_attention_forward(
- query, key, value, attention_mask, query_length, dropout=attn_dropout
- )
-
- attn_weights_reshaped = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
- attn_output = self.c_proj(attn_weights_reshaped)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
-
- if output_attentions:
- if self.multi_query:
- # Transpose to return weights in the usual format (batch_size, num_heads, query_length, key_length)
- attn_weights_reshaped = attn_weights_reshaped.transpose(1, 2)
- else:
- attn_weights_reshaped = None
-
- outputs += (attn_weights_reshaped,)
-
- return outputs # a, present, (attentions)
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-class GPTBigCodeSdpaAttention(GPTBigCodeAttention):
- def _attn(self, query, key, value, attention_mask=None, head_mask=None):
- if head_mask is not None:
- # The super dispatch is done in the forward.
- raise ValueError(
- "PyTorch SDPA does not support head_mask. Please open an issue in Transformers repository."
- )
-
- scale = None
- if not self.scale_attn_weights:
- scale = 1
-
- # MQA models: (batch_size, query_length, num_heads * head_dim)
- # MHA models: (batch_size, num_heads, query_length, head_dim)
- query_shape = query.shape
- batch_size = query_shape[0]
- key.shape[-2]
-
- if self.multi_query:
- query_length = query_shape[1]
-
- # SDPA requires the dimension [..., sequence_length, head_dim].
- query = query.view(batch_size, query_length, self.num_heads, self.head_dim).transpose(1, 2)
-
- # Without these unsqueeze, SDPA complains as the query and key/value have a different number of dimensions.
- key = key.unsqueeze(1)
- value = value.unsqueeze(1)
-
- # Although these expand are not numerically useful, PyTorch can not dispatch to memory-efficient backend
- # and flash attention backend (No available kernel. Aborting execution.) from the shapes
- # query = [batch_size, num_heads, query_length, head_dim]
- # key = [batch_size, 1, past_length, head_dim]
- # value = [batch_size, 1, past_length, head_dim]
- #
- # torch==2.1.2 is bugged with non-contiguous inputs with custom attn_mask (https://github.com/pytorch/pytorch/issues/112577), hence the check.
- if is_torch_greater_or_equal_than_2_2:
- key = key.expand(-1, self.num_heads, -1, -1)
- value = value.expand(-1, self.num_heads, -1, -1)
- else:
- query_length = query_shape[-1]
-
- # See the comment above.
- if query.device.type == "cuda" and attention_mask is not None:
- query = query.contiguous()
- key = key.contiguous()
- value = value.contiguous()
-
- sdpa_result = torch.nn.functional.scaled_dot_product_attention(
- query,
- key,
- value,
- attn_mask=attention_mask,
- dropout_p=self.attn_pdrop if self.training else 0.0,
- # The query_length > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case query_length == 1.
- is_causal=self.is_causal and attention_mask is None and query_length > 1,
- scale=scale,
- )
-
- if self.multi_query:
- # (batch_size, num_heads, seq_len, head_dim) --> (batch_size, seq_len, num_heads, head_dim)
- sdpa_result = sdpa_result.transpose(1, 2)
-
- # Reshape is kind of expensive here, as it does a memory copy,
- # but I did not manage to make away without it (logits do not match when using view)
- # (batch_size, seq_len, num_heads, head_dim) --> (batch_size, seq_len, num_heads * head_dim)
- sdpa_result = sdpa_result.reshape(query_shape)
-
- return sdpa_result, None
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- layer_past: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[
- Tuple[torch.Tensor, Optional[torch.Tensor]],
- Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[torch.Tensor, ...]],
- ]:
- if encoder_hidden_states is not None:
- if not hasattr(self, "q_attn") or not self.is_cross_attention:
- raise ValueError(
- "If class is used as cross attention, the weights `q_attn` have to be defined. "
- "Please make sure to instantiate class with `GPTBigCodeAttention(..., is_cross_attention=True)`."
- )
-
- query = self.q_attn(hidden_states)
- key_value = self.c_attn(encoder_hidden_states)
- attention_mask = encoder_attention_mask
- elif self.multi_query:
- query, key_value = self.c_attn(hidden_states).split((self.embed_dim, 2 * self.kv_dim), dim=2)
- else:
- # Note: We split as (self.num_heads, 3, self.head_dim) instead of (3, self.num_heads, self.head_dim),
- # i.e., the memory layout is not the same as GPT2.
- # This makes the concatenation with past_key_value more efficient.
- query, key_value = (
- self.c_attn(hidden_states)
- .view(*hidden_states.shape[:2], self.num_heads, 3 * self.head_dim)
- .transpose(1, 2)
- .split((self.head_dim, 2 * self.head_dim), dim=3)
- )
-
- if layer_past is not None:
- key_value = torch.cat((layer_past, key_value), dim=-2)
- present = key_value if use_cache else None
-
- key, value = key_value.split((self.head_dim, self.head_dim), dim=-1)
-
- if not output_attentions and head_mask is None:
- # Difference with the original implementation: there is no need to transpose the key here,
- # as SDPA expects seq_length to be at index -2 for the key as well
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
- else:
- # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once this is implemented.
- logger.warning_once(
- "GPTBigCodeModel is using GPTBigCodeSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` and `head_mask` not None."
- ' Falling back to the manual attention implementation, but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
- )
- attn_output, attn_weights = super()._attn(query, key.transpose(-1, -2), value, attention_mask, head_mask)
-
- if not self.multi_query:
- attn_output = attn_output.transpose(1, 2).reshape(hidden_states.shape)
- attn_output = self.c_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- if self.multi_query:
- # Transpose to return weights in the usual format (batch_size, num_heads, query_length, key_length)
- attn_weights = attn_weights.transpose(1, 2)
- outputs += (attn_weights,)
-
- return outputs
-
-
-class GPTBigCodeMLP(nn.Module):
- def __init__(self, intermediate_size, config):
- super().__init__()
- embed_dim = config.hidden_size
- self.c_fc = nn.Linear(embed_dim, intermediate_size)
- self.c_proj = nn.Linear(intermediate_size, embed_dim)
- self.act = ACT2FN[config.activation_function]
- self.dropout = nn.Dropout(config.resid_pdrop)
-
- # Copied from transformers.models.gpt2.modeling_gpt2.GPT2MLP.forward
- def forward(self, hidden_states: Optional[Tuple[torch.FloatTensor]]) -> torch.FloatTensor:
- hidden_states = self.c_fc(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.c_proj(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-GPTBIGCODE_ATTENTION_CLASSES = {
- "eager": GPTBigCodeAttention,
- "flash_attention_2": GPTBigCodeFlashAttention2,
- "sdpa": GPTBigCodeSdpaAttention,
-}
-
-
-class GPTBigCodeBlock(nn.Module):
- def __init__(self, config, layer_idx=None):
- super().__init__()
- hidden_size = config.hidden_size
- self.inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
-
- self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
- self.attn = GPTBIGCODE_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
-
- self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
- if config.add_cross_attention:
- if config.multi_query:
- raise NotImplementedError("Cross-attention not implemented for MQA")
-
- self.crossattention = GPTBIGCODE_ATTENTION_CLASSES[config._attn_implementation](
- config, is_cross_attention=True, layer_idx=layer_idx
- )
-
- self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
- self.mlp = GPTBigCodeMLP(self.inner_dim, config)
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.Tensor]],
- layer_past: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[
- Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
- ]:
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
- outputs = attn_outputs[1:]
- # residual connection
- hidden_states = attn_output + residual
-
- if encoder_hidden_states is not None:
- # add one self-attention block for cross-attention
- if not hasattr(self, "crossattention"):
- raise ValueError(
- f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
- "cross-attention layers by setting `config.add_cross_attention=True`"
- )
- residual = hidden_states
- hidden_states = self.ln_cross_attn(hidden_states)
- cross_attn_outputs = self.crossattention(
- hidden_states,
- attention_mask=attention_mask,
- head_mask=head_mask,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- output_attentions=output_attentions,
- )
- attn_output = cross_attn_outputs[0]
- # residual connection
- hidden_states = residual + attn_output
- outputs = outputs + cross_attn_outputs[2:] # add cross attentions if we output attention weights
-
- residual = hidden_states
- hidden_states = self.ln_2(hidden_states)
- feed_forward_hidden_states = self.mlp(hidden_states)
- # residual connection
- hidden_states = residual + feed_forward_hidden_states
-
- if use_cache:
- outputs = (hidden_states,) + outputs
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- return outputs # hidden_states, present, (attentions, cross_attentions)
-
-
-class GPTBigCodePreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTBigCodeConfig
- base_model_prefix = "transformer"
- supports_gradient_checkpointing = True
- _no_split_modules = ["GPTBigCodeBlock"]
- _skip_keys_device_placement = "past_key_values"
- _supports_flash_attn_2 = True
- _supports_sdpa = True
-
- def __init__(self, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, (GPTBigCodeMLP, GPTBigCodeAttention)):
- # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
- # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
- # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
- # > -- GPT-2 :: https://openai.com/blog/better-language-models/
- #
- # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
- module.c_proj.weight.data.normal_(
- mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.n_layer))
- )
- module.c_proj._is_hf_initialized = True
- elif isinstance(module, nn.Linear):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-GPT_BIGCODE_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`GPTBigCodeConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPT_BIGCODE_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
- `past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input
- sequence tokens in the vocabulary.
-
- If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
- `input_ids`.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`Tuple[torch.Tensor]` of length `config.n_layers`):
- Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
- their past given to this model should not be passed as `input_ids` as they have already been computed.
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- If `past_key_values` is used, `attention_mask` needs to contain the masking strategy that was used for
- `past_key_values`. In other words, the `attention_mask` always has to have the length:
- `len(past_key_values) + len(input_ids)`
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
-
- If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see
- `past_key_values`).
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare GPT_BIGCODE Model transformer outputting raw hidden-states without any specific head on top.",
- GPT_BIGCODE_START_DOCSTRING,
-)
-class GPTBigCodeModel(GPTBigCodePreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.multi_query = config.multi_query
- self.embed_dim = config.hidden_size
-
- self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
- self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
-
- self.drop = nn.Dropout(config.embd_pdrop)
- self.h = nn.ModuleList([GPTBigCodeBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)])
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
-
- max_positions = config.max_position_embeddings
- self.register_buffer(
- "bias", torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)), persistent=False
- )
-
- self.gradient_checkpointing = False
-
- self._use_sdpa = config._attn_implementation == "sdpa"
- self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.wte
-
- def set_input_embeddings(self, new_embeddings):
- self.wte = new_embeddings
-
- @add_start_docstrings_to_model_forward(GPT_BIGCODE_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPastAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- past_key_values: Optional[List[torch.Tensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- input_ids = input_ids.view(-1, input_shape[-1])
- batch_size = input_ids.shape[0]
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- batch_size = inputs_embeds.shape[0]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if batch_size <= 0:
- raise ValueError("batch_size has to be defined and > 0")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if token_type_ids is not None:
- token_type_ids = token_type_ids.view(-1, input_shape[-1])
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * len(self.h))
- else:
- past_length = past_key_values[0].size(-2)
-
- if attention_mask is not None and len(attention_mask.shape) == 2 and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_length > 0:
- position_ids = position_ids[:, past_length : input_shape[-1] + past_length :]
- elif position_ids is None:
- position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0)
-
- # Self-attention mask.
- query_length = input_shape[-1]
- key_length = past_length + query_length
- self_attention_mask = self.bias[None, key_length - query_length : key_length, :key_length]
-
- if self._use_flash_attention_2:
- # 2d mask is passed through the layers
- attention_mask = attention_mask.bool() if (attention_mask is not None and 0 in attention_mask) else None
- encoder_attention_mask = (
- encoder_attention_mask.bool()
- if (encoder_attention_mask is not None and 0 in encoder_attention_mask)
- else None
- )
- else:
- # 4d mask is passed through the layers
- if attention_mask is not None:
- self_attention_mask = self_attention_mask * attention_mask.view(batch_size, 1, -1).to(
- dtype=torch.bool, device=self_attention_mask.device
- )
-
- # MQA models: (batch_size, query_length, n_heads, key_length)
- # MHA models: (batch_size, n_heads, query_length, key_length)
- self_attention_mask = self_attention_mask.unsqueeze(2 if self.multi_query else 1)
-
- if self._use_sdpa and head_mask is None and not output_attentions:
- # SDPA with a custom mask is much faster in fp16/fp32 dtype rather than bool. Cast here to floating point instead of at every layer.
- dtype = self.wte.weight.dtype
- min_dtype = torch.finfo(dtype).min
- self_attention_mask = torch.where(
- self_attention_mask,
- torch.full([], 0.0, dtype=dtype, device=self_attention_mask.device),
- torch.full([], min_dtype, dtype=dtype, device=self_attention_mask.device),
- )
-
- # output_attentions=True can not be supported when using SDPA, and we fall back on
- # the manual implementation that requires a 4D causal mask in all cases.
- if self.multi_query:
- # gpt_bigcode using MQA has the bad taste to use a causal mask with shape
- # [batch_size, target_length, 1, source_length], not compatible with SDPA, hence this transpose.
- self_attention_mask = self_attention_mask.transpose(1, 2)
-
- if query_length > 1 and attention_mask is not None and attention_mask.device.type == "cuda":
- # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
- # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
- self_attention_mask = AttentionMaskConverter._unmask_unattended(
- self_attention_mask, min_dtype=min_dtype
- )
-
- attention_mask = self_attention_mask
-
- # If a 2D or 3D attention mask is provided for the cross-attention
- # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
- if (
- self.config.add_cross_attention
- and encoder_hidden_states is not None
- and encoder_attention_mask is not None
- ):
- if encoder_attention_mask.dim() == 2:
- encoder_attention_mask.unsqueeze(1)
- assert encoder_attention_mask.dim() == 3
- encoder_attention_mask = encoder_attention_mask.bool().unsqueeze(2 if self.multi_query else 1)
- else:
- encoder_attention_mask = None
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # head_mask has shape n_layer x batch x n_heads x N x N
- head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-
- if inputs_embeds is None:
- inputs_embeds = self.wte(input_ids)
- position_embeds = self.wpe(position_ids)
- hidden_states = inputs_embeds + position_embeds
-
- if token_type_ids is not None:
- token_type_embeds = self.wte(token_type_ids)
- hidden_states = hidden_states + token_type_embeds
-
- hidden_states = self.drop(hidden_states)
-
- output_shape = input_shape + (hidden_states.size(-1),)
-
- presents = [] if use_cache else None
- all_self_attentions = () if output_attentions else None
- all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
- all_hidden_states = () if output_hidden_states else None
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- None,
- attention_mask,
- head_mask[i],
- encoder_hidden_states,
- encoder_attention_mask,
- use_cache,
- output_attentions,
- )
- else:
- outputs = block(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask[i],
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
-
- hidden_states = outputs[0]
- if use_cache:
- presents.append(outputs[1])
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
- if self.config.add_cross_attention:
- all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],)
-
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = hidden_states.view(output_shape)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions]
- if v is not None
- )
-
- return BaseModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- cross_attentions=all_cross_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT_BIGCODE Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- GPT_BIGCODE_START_DOCSTRING,
-)
-class GPTBigCodeForCausalLM(GPTBigCodePreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.transformer = GPTBigCodeModel(config)
- self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # Omit tokens covered by past_key_values
- if past_key_values:
- if self.config.multi_query:
- past_length = past_key_values[0].shape[1]
- else:
- past_length = past_key_values[0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
- if token_type_ids is not None:
- token_type_ids = token_type_ids[:, -input_ids.shape[1] :]
-
- attention_mask = kwargs.get("attention_mask", None)
- position_ids = kwargs.get("position_ids", None)
-
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
- else:
- position_ids = None
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- model_inputs = {"input_ids": input_ids}
-
- model_inputs.update(
- {
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "position_ids": position_ids,
- "attention_mask": attention_mask,
- "token_type_ids": token_type_ids,
- }
- )
- return model_inputs
-
- @add_start_docstrings_to_model_forward(GPT_BIGCODE_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutputWithCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- encoder_hidden_states: Optional[torch.Tensor] = None,
- encoder_attention_mask: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
- r"""
- labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- encoder_hidden_states=encoder_hidden_states,
- encoder_attention_mask=encoder_attention_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
-
- lm_logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous().to(shift_logits.device)
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
- if not return_dict:
- output = (lm_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithCrossAttentions(
- loss=loss,
- logits=lm_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- cross_attentions=transformer_outputs.cross_attentions,
- )
-
- @staticmethod
- def _reorder_cache(
- past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
- ) -> Tuple[Tuple[torch.Tensor]]:
- """
- This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
- [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
- beam_idx at every generation step.
- """
- return tuple(layer_past.index_select(0, beam_idx.to(layer_past.device)) for layer_past in past_key_values)
-
-
-@add_start_docstrings(
- """
- The GPTBigCode Model transformer with a sequence classification head on top (linear layer).
-
- [`GPTBigCodeForSequenceClassification`] uses the last token in order to do the classification, as other causal
- models (e.g. GPT-1) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- GPT_BIGCODE_START_DOCSTRING,
-)
-class GPTBigCodeForSequenceClassification(GPTBigCodePreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = GPTBigCodeModel(config)
- self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT_BIGCODE_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size, sequence_length = input_ids.shape[:2]
- else:
- batch_size, sequence_length = inputs_embeds.shape[:2]
-
- assert (
- self.config.pad_token_id is not None or batch_size == 1
- ), "Cannot handle batch sizes > 1 if no padding token is defined."
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
-
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
-
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- GPT_BIGCODE Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g.
- for Named-Entity-Recognition (NER) tasks.
- """,
- GPT_BIGCODE_START_DOCSTRING,
-)
-class GPTBigCodeForTokenClassification(GPTBigCodePreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.transformer = GPTBigCodeModel(config)
- if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
- classifier_dropout = config.classifier_dropout
- elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
- classifier_dropout = config.hidden_dropout
- else:
- classifier_dropout = 0.1
- self.dropout = nn.Dropout(classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT_BIGCODE_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = transformer_outputs[0]
- hidden_states = self.dropout(hidden_states)
- logits = self.classifier(hidden_states)
-
- loss = None
- if labels is not None:
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1).to(logits.device))
-
- if not return_dict:
- output = (logits,) + transformer_outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
diff --git a/transformers/models/gpt_neo/__init__.py b/transformers/models/gpt_neo/__init__.py
deleted file mode 100644
index 02ca0a11949b73ecef0329412d869ce1996d1bc6..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neo/__init__.py
+++ /dev/null
@@ -1,85 +0,0 @@
-# Copyright 2021 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available
-
-
-_import_structure = {
- "configuration_gpt_neo": ["GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTNeoConfig", "GPTNeoOnnxConfig"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_gpt_neo"] = [
- "GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GPTNeoForCausalLM",
- "GPTNeoForQuestionAnswering",
- "GPTNeoForSequenceClassification",
- "GPTNeoForTokenClassification",
- "GPTNeoModel",
- "GPTNeoPreTrainedModel",
- "load_tf_weights_in_gpt_neo",
- ]
-
-try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flax_gpt_neo"] = [
- "FlaxGPTNeoForCausalLM",
- "FlaxGPTNeoModel",
- "FlaxGPTNeoPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_gpt_neo import GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoConfig, GPTNeoOnnxConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_gpt_neo import (
- GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST,
- GPTNeoForCausalLM,
- GPTNeoForQuestionAnswering,
- GPTNeoForSequenceClassification,
- GPTNeoForTokenClassification,
- GPTNeoModel,
- GPTNeoPreTrainedModel,
- load_tf_weights_in_gpt_neo,
- )
-
- try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel, FlaxGPTNeoPreTrainedModel
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gpt_neo/__pycache__/__init__.cpython-310.pyc b/transformers/models/gpt_neo/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index f8c1d6f82a3c3f99028502d898cb44d98ae4d676..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neo/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neo/__pycache__/configuration_gpt_neo.cpython-310.pyc b/transformers/models/gpt_neo/__pycache__/configuration_gpt_neo.cpython-310.pyc
deleted file mode 100644
index 73f5517e45213b9d6855ed56ed72364cfdac0407..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neo/__pycache__/configuration_gpt_neo.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neo/__pycache__/convert_gpt_neo_mesh_tf_to_pytorch.cpython-310.pyc b/transformers/models/gpt_neo/__pycache__/convert_gpt_neo_mesh_tf_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 04b1ef6b43a2e3247d0cd86aaf502a35494768b7..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neo/__pycache__/convert_gpt_neo_mesh_tf_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neo/__pycache__/modeling_flax_gpt_neo.cpython-310.pyc b/transformers/models/gpt_neo/__pycache__/modeling_flax_gpt_neo.cpython-310.pyc
deleted file mode 100644
index 55a64f3258553fb47480a63587c06fd45a318fae..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neo/__pycache__/modeling_flax_gpt_neo.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neo/__pycache__/modeling_gpt_neo.cpython-310.pyc b/transformers/models/gpt_neo/__pycache__/modeling_gpt_neo.cpython-310.pyc
deleted file mode 100644
index 4d24e1f95d7709383271f78fb2c568735a9dfa9e..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neo/__pycache__/modeling_gpt_neo.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neo/configuration_gpt_neo.py b/transformers/models/gpt_neo/configuration_gpt_neo.py
deleted file mode 100644
index 411b392180b018ca027d08138f363a7335e9d384..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neo/configuration_gpt_neo.py
+++ /dev/null
@@ -1,272 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPT Neo model configuration"""
-
-from collections import OrderedDict
-from typing import Any, Mapping, Optional
-
-from ... import PreTrainedTokenizer, TensorType, is_torch_available
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfigWithPast
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GPTNeoConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GPTNeoModel`]. It is used to instantiate a GPT
- Neo model according to the specified arguments, defining the model architecture. Instantiating a configuration with
- the defaults will yield a similar configuration to that of the GPTNeo
- [EleutherAI/gpt-neo-1.3B](https://huggingface.co/EleutherAI/gpt-neo-1.3B) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 50257):
- Vocabulary size of the GPT Neo model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`GPTNeoModel`]. Vocabulary size of the model. Defines the different
- tokens that can be represented by the *inputs_ids* passed to the forward method of [`GPTNeoModel`].
- max_position_embeddings (`int`, *optional*, defaults to 2048):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- hidden_size (`int`, *optional*, defaults to 2048):
- Dimensionality of the encoder layers and the pooler layer.
- num_layers (`int`, *optional*, defaults to 24):
- Number of hidden layers in the Transformer encoder.
- attention_types (`List`, *optional*, defaults to `[[['global', 'local'], 12]]`):
- The type of attention for each layer in a `List` of the following format `[[["attention_type"],
- num_layerss]]` e.g. for a 24 layer model `[[["global"], 24]]` or `[[["global", "local"], 12]]` Choose the
- value of `attention_type` from `["global", "local"]`
- num_heads (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 8192):
- Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- window_size (`int`, *optional*, defaults to 256):
- The size of the sliding window for local attention.
- activation_function (`str` or `function`, *optional*, defaults to `"gelu_new"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- resid_dropout (`float`, *optional*, defaults to 0.0):
- Residual dropout used in the attention pattern.
- embed_dropout (`float`, *optional*, defaults to 0.0):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the attention probabilities.
- classifier_dropout (`float`, *optional*, defaults to 0.1):
- Argument used when doing token classification, used in the model [`GPTNeoForTokenClassification`]. The
- dropout ratio for the hidden layer.
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
- The epsilon used by the layer normalization layers.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- bos_token_id (`int`, *optional*, defaults to 50256):
- The id of the beginning of sentence token in the vocabulary.
- eos_token_id (`int`, *optional*, defaults to 50256):
- The id of the end of sentence token in the vocabulary.
-
- Example:
-
- ```python
- >>> from transformers import GPTNeoConfig, GPTNeoModel
-
- >>> # Initializing a GPTNeo EleutherAI/gpt-neo-1.3B style configuration
- >>> configuration = GPTNeoConfig()
-
- >>> # Initializing a model (with random weights) from the EleutherAI/gpt-neo-1.3B style configuration
- >>> model = GPTNeoModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "gpt_neo"
- keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {"num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}
-
- def __init__(
- self,
- vocab_size=50257,
- max_position_embeddings=2048,
- hidden_size=2048,
- num_layers=24,
- attention_types=[[["global", "local"], 12]],
- num_heads=16,
- intermediate_size=None,
- window_size=256,
- activation_function="gelu_new",
- resid_dropout=0.0,
- embed_dropout=0.0,
- attention_dropout=0.0,
- classifier_dropout=0.1,
- layer_norm_epsilon=1e-5,
- initializer_range=0.02,
- use_cache=True,
- bos_token_id=50256,
- eos_token_id=50256,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.hidden_size = hidden_size
- self.num_layers = num_layers
- self.num_heads = num_heads
- self.intermediate_size = intermediate_size
- self.window_size = window_size
- self.activation_function = activation_function
- self.resid_dropout = resid_dropout
- self.embed_dropout = embed_dropout
- self.attention_dropout = attention_dropout
- self.classifier_dropout = classifier_dropout
- self.layer_norm_epsilon = layer_norm_epsilon
- self.initializer_range = initializer_range
- self.use_cache = use_cache
-
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
-
- self.attention_types = attention_types
- self.attention_layers = self.expand_attention_types_params(attention_types)
-
- if len(self.attention_layers) != self.num_layers:
- raise ValueError(
- "Configuration for convolutional module is incorrect. "
- "It is required that `len(config.attention_layers)` == `config.num_layers` "
- f"but is `len(config.attention_layers) = {len(self.attention_layers)}`, "
- f"`config.num_layers = {self.num_layers}`. "
- "`config.attention_layers` is prepared using `config.attention_types`. "
- "Please verify the value of `config.attention_types` argument."
- )
-
- super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
-
- @staticmethod
- def expand_attention_types_params(attention_types):
- attentions = []
- for item in attention_types:
- for _ in range(item[1]):
- attentions.extend(item[0])
- return attentions
-
-
-def custom_unfold(input, dimension, size, step):
- """Custom torch.Tensor.unfold implementation to enable the export to ONNX."""
- import torch
-
- shape = input.size()
- rank = len(shape)
- sizedim = shape[dimension]
-
- low_indices = torch.arange(0, sizedim, step)
- min_length = torch.div(sizedim - size, step, rounding_mode="floor") + 1
- indices = torch.arange(size) + low_indices[:min_length][:, None]
-
- s = [slice(None)] * rank
- s[dimension] = indices
- sliced = input[s]
-
- perm = list(range(0, rank + 1))
- perm.append(perm.pop(dimension + 1))
-
- return sliced.permute(perm)
-
-
-def custom_get_block_length_and_num_blocks(seq_length, window_size):
- """
- Custom implementation for GPTNeoAttentionMixin._get_block_length_and_num_blocks to enable the export to ONNX as
- original implementation uses Python variables and control flow.
- """
- import torch
-
- candidates = torch.arange(1, window_size)
- remainders = torch.remainder(seq_length, candidates)
- divisor_indices = remainders == 0
- divisors = candidates[divisor_indices]
- largest_divisor = torch.max(divisors)
- return largest_divisor, torch.div(seq_length, largest_divisor, rounding_mode="floor")
-
-
-class GPTNeoOnnxConfig(OnnxConfigWithPast):
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
- if self.use_past:
- self.fill_with_past_key_values_(common_inputs, direction="inputs")
- common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"}
- else:
- common_inputs["attention_mask"] = {0: "batch", 1: "sequence"}
-
- return common_inputs
-
- @property
- def num_attention_heads(self) -> int:
- return self._config.num_heads
-
- def generate_dummy_inputs(
- self,
- tokenizer: PreTrainedTokenizer,
- batch_size: int = -1,
- seq_length: int = -1,
- is_pair: bool = False,
- framework: Optional[TensorType] = None,
- ) -> Mapping[str, Any]:
- common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
- tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
- )
-
- # We need to order the input in the way they appears in the forward()
- ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]})
-
- # Need to add the past_keys
- if self.use_past:
- if not is_torch_available():
- raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
- else:
- import torch
-
- batch, seqlen = common_inputs["input_ids"].shape
- # Not using the same length for past_key_values
- past_key_values_length = seqlen + 2
- past_shape = (
- batch,
- self.num_attention_heads,
- past_key_values_length,
- self._config.hidden_size // self.num_attention_heads,
- )
- ordered_inputs["past_key_values"] = [
- (torch.zeros(past_shape), torch.zeros(past_shape)) for _ in range(self.num_layers)
- ]
-
- ordered_inputs["attention_mask"] = common_inputs["attention_mask"]
- if self.use_past:
- mask_dtype = ordered_inputs["attention_mask"].dtype
- ordered_inputs["attention_mask"] = torch.cat(
- [ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1
- )
-
- return ordered_inputs
-
- @property
- def default_onnx_opset(self) -> int:
- return 13
diff --git a/transformers/models/gpt_neo/convert_gpt_neo_mesh_tf_to_pytorch.py b/transformers/models/gpt_neo/convert_gpt_neo_mesh_tf_to_pytorch.py
deleted file mode 100644
index 4a5fddd0a9d0f95b83777ebc9207a40811940535..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neo/convert_gpt_neo_mesh_tf_to_pytorch.py
+++ /dev/null
@@ -1,72 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The Eleuther AI and HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Convert GPT Neo checkpoint."""
-
-
-import argparse
-import json
-
-from transformers import GPTNeoConfig, GPTNeoForCausalLM, load_tf_weights_in_gpt_neo
-from transformers.utils import logging
-
-
-logging.set_verbosity_info()
-
-
-def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path):
- # Initialise PyTorch model
- config_json = json.load(open(config_file, "r"))
- config = GPTNeoConfig(
- hidden_size=config_json["n_embd"],
- num_layers=config_json["n_layer"],
- num_heads=config_json["n_head"],
- attention_types=config_json["attention_types"],
- max_position_embeddings=config_json["n_positions"],
- resid_dropout=config_json["res_dropout"],
- embed_dropout=config_json["embed_dropout"],
- attention_dropout=config_json["attn_dropout"],
- )
- print(f"Building PyTorch model from configuration: {config}")
- model = GPTNeoForCausalLM(config)
-
- # Load weights from tf checkpoint
- load_tf_weights_in_gpt_neo(model, config, tf_checkpoint_path)
-
- # Save pytorch-model
- print(f"Save PyTorch model to {pytorch_dump_path}")
- model.save_pretrained(pytorch_dump_path)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- # Required parameters
- parser.add_argument(
- "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
- )
- parser.add_argument(
- "--config_file",
- default=None,
- type=str,
- required=True,
- help=(
- "The config json file corresponding to the pre-trained mesh-tf model. \n"
- "This specifies the model architecture."
- ),
- )
- parser.add_argument(
- "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
- )
- args = parser.parse_args()
- convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)
diff --git a/transformers/models/gpt_neo/modeling_flax_gpt_neo.py b/transformers/models/gpt_neo/modeling_flax_gpt_neo.py
deleted file mode 100644
index 5639ca50f166a272968b497df696d15410f180ea..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neo/modeling_flax_gpt_neo.py
+++ /dev/null
@@ -1,684 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The Eleuther AI and The Google Flax Team Authors and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from functools import partial
-from typing import Optional, Tuple
-
-import flax.linen as nn
-import jax
-import jax.numpy as jnp
-from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
-from flax.linen import combine_masks, make_causal_mask
-from flax.linen.attention import dot_product_attention_weights
-from flax.traverse_util import flatten_dict, unflatten_dict
-from jax import lax
-
-from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput
-from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_gpt_neo import GPTNeoConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "GPTNeoConfig"
-_CHECKPOINT_FOR_DOC = "EleutherAI/gpt-neo-1.3B"
-
-
-GPT_NEO_START_DOCSTRING = r"""
-
- This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a Flax Linen
- [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a
- regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.
-
- Finally, this model supports inherent JAX features such as:
-
- - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
-
- Parameters:
- config ([`GPTNeoConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.
- dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
- The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
- `jax.numpy.bfloat16` (on TPUs).
-
- This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
- specified all the computation will be performed with the given `dtype`.
-
- **Note that this only specifies the dtype of the computation and does not influence the dtype of model
- parameters.**
-
- If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and
- [`~FlaxPreTrainedModel.to_bf16`].
-"""
-
-GPT_NEO_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`numpy.ndarray` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length`. Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
- past_key_values (`Dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):
- Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
- auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-class FlaxGPTNeoSelfAttention(nn.Module):
- config: GPTNeoConfig
- attention_type: str
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- config = self.config
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_heads
- if self.head_dim * self.num_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and "
- f"`num_heads`: {self.num_heads})."
- )
-
- self.attn_dropout = nn.Dropout(config.attention_dropout)
- self.resid_dropout = nn.Dropout(config.resid_dropout)
-
- dense = partial(
- nn.Dense,
- self.embed_dim,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
- )
-
- self.q_proj, self.k_proj, self.v_proj = dense(use_bias=False), dense(use_bias=False), dense(use_bias=False)
- self.out_proj = dense()
-
- self.causal_mask = make_causal_mask(jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool")
- if self.attention_type == "local":
- self.causal_mask = self.causal_mask ^ jnp.tril(self.causal_mask, -config.window_size)
-
- def _split_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))
-
- def _merge_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))
-
- @nn.compact
- def _concatenate_to_cache(self, key, value, query, attention_mask):
- """
- This function takes projected key, value states from a single input token and concatenates the states to cached
- states from previous steps. This function is slighly adapted from the official Flax repository:
- https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252
- """
- # detect if we're initializing by absence of existing cache data.
- is_initialized = self.has_variable("cache", "cached_key")
- cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype)
- cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype)
- cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
-
- if is_initialized:
- *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape
- # update key, value caches with our new 1d spatial slices
- cur_index = cache_index.value
- indices = (0,) * len(batch_dims) + (cur_index, 0, 0)
- key = lax.dynamic_update_slice(cached_key.value, key, indices)
- value = lax.dynamic_update_slice(cached_value.value, value, indices)
- cached_key.value = key
- cached_value.value = value
- num_updated_cache_vectors = query.shape[1]
- cache_index.value = cache_index.value + num_updated_cache_vectors
- # causal mask for cached decoder self-attention: our single query position should only attend to those key positions that have already been generated and cached, not the remaining zero elements.
- pad_mask = jnp.broadcast_to(
- jnp.arange(max_length) < cur_index + num_updated_cache_vectors,
- tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),
- )
- attention_mask = combine_masks(pad_mask, attention_mask)
- return key, value, attention_mask
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- query = self.q_proj(hidden_states) * jnp.sqrt(self.head_dim).astype(self.dtype)
- key = self.k_proj(hidden_states)
- value = self.v_proj(hidden_states)
-
- query = self._split_heads(query)
- key = self._split_heads(key)
- value = self._split_heads(value)
-
- query_length, key_length = query.shape[1], key.shape[1]
-
- if self.has_variable("cache", "cached_key"):
- mask_shift = self.variables["cache"]["cache_index"]
- max_decoder_length = self.variables["cache"]["cached_key"].shape[1]
- causal_mask = lax.dynamic_slice(
- self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)
- )
- else:
- causal_mask = self.causal_mask[:, :, :query_length, :key_length]
-
- batch_size = hidden_states.shape[0]
- causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])
-
- attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
- attention_mask = combine_masks(attention_mask, causal_mask)
-
- dropout_rng = None
- if not deterministic and self.config.attention_dropout > 0.0:
- dropout_rng = self.make_rng("dropout")
-
- # During fast autoregressive decoding, we feed one position at a time,
- # and cache the keys and values step by step.
- if self.has_variable("cache", "cached_key") or init_cache:
- key, value, attention_mask = self._concatenate_to_cache(key, value, query, attention_mask)
-
- # transform boolean mask into float mask
- attention_bias = lax.select(
- attention_mask > 0,
- jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
- jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
- )
-
- # usual dot product attention
- attn_weights = dot_product_attention_weights(
- query,
- key,
- bias=attention_bias,
- dropout_rng=dropout_rng,
- dropout_rate=self.config.attention_dropout,
- deterministic=deterministic,
- dtype=self.dtype,
- precision=None,
- )
-
- attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value)
- attn_output = self._merge_heads(attn_output)
- attn_output = self.out_proj(attn_output)
- attn_output = self.resid_dropout(attn_output, deterministic=deterministic)
-
- outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)
- return outputs
-
-
-class FlaxGPTNeoAttention(nn.Module):
- config: GPTNeoConfig
- layer_id: int = 0
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- attention_type = self.config.attention_layers[self.layer_id]
- self.attention = FlaxGPTNeoSelfAttention(self.config, attention_type, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- return self.attention(
- hidden_states,
- attention_mask=attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
-
-
-class FlaxGPTNeoMLP(nn.Module):
- config: GPTNeoConfig
- intermediate_size: int
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- embed_dim = self.config.hidden_size
- kernel_init = jax.nn.initializers.normal(self.config.initializer_range)
- self.c_fc = nn.Dense(self.intermediate_size, dtype=self.dtype, kernel_init=kernel_init)
- self.c_proj = nn.Dense(embed_dim, dtype=self.dtype, kernel_init=kernel_init)
- self.act = ACT2FN[self.config.activation_function]
- self.dropout = nn.Dropout(rate=self.config.resid_dropout)
-
- def __call__(self, hidden_states, deterministic: bool = True):
- hidden_states = self.c_fc(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.c_proj(hidden_states)
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- return hidden_states
-
-
-class FlaxGPTNeoBlock(nn.Module):
- config: GPTNeoConfig
- layer_id: int = 0
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- hidden_size = self.config.hidden_size
- inner_dim = self.config.intermediate_size if self.config.intermediate_size is not None else 4 * hidden_size
-
- self.ln_1 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
- self.attn = FlaxGPTNeoAttention(self.config, layer_id=self.layer_id, dtype=self.dtype)
- self.ln_2 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
- self.mlp = FlaxGPTNeoMLP(self.config, inner_dim, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- outputs = self.attn(
- hidden_states,
- attention_mask=attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
- # residual connection
- attn_output = outputs[0]
- hidden_states = attn_output + residual
-
- residual = hidden_states
- hidden_states = self.ln_2(hidden_states)
- feed_forward_hidden_states = self.mlp(hidden_states, deterministic=deterministic)
- # residual connection
- hidden_states = residual + feed_forward_hidden_states
-
- return (hidden_states,) + outputs[1:]
-
-
-class FlaxGPTNeoPreTrainedModel(FlaxPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTNeoConfig
- base_model_prefix = "transformer"
- module_class: nn.Module = None
-
- def __init__(
- self,
- config: GPTNeoConfig,
- input_shape: Tuple = (1, 1),
- seed: int = 0,
- dtype: jnp.dtype = jnp.float32,
- _do_init: bool = True,
- **kwargs,
- ):
- module = self.module_class(config=config, dtype=dtype, **kwargs)
- super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
-
- def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
- # init input tensors
- input_ids = jnp.zeros(input_shape, dtype="i4")
- attention_mask = jnp.ones_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape)
- params_rng, dropout_rng = jax.random.split(rng)
- rngs = {"params": params_rng, "dropout": dropout_rng}
-
- random_params = self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False)["params"]
-
- if params is not None:
- random_params = flatten_dict(unfreeze(random_params))
- params = flatten_dict(unfreeze(params))
- for missing_key in self._missing_keys:
- params[missing_key] = random_params[missing_key]
- self._missing_keys = set()
- return freeze(unflatten_dict(params))
- else:
- return random_params
-
- def init_cache(self, batch_size, max_length):
- r"""
- Args:
- batch_size (`int`):
- batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
- max_length (`int`):
- maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
- cache.
- """
- # init input variables to retrieve cache
- input_ids = jnp.ones((batch_size, max_length))
- attention_mask = jnp.ones_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
-
- init_variables = self.module.init(
- jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True
- )
- return unfreeze(init_variables["cache"])
-
- @add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- position_ids=None,
- params: dict = None,
- past_key_values: dict = None,
- dropout_rng: jax.random.PRNGKey = None,
- train: bool = False,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ):
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- batch_size, sequence_length = input_ids.shape
-
- if position_ids is None:
- if past_key_values is not None:
- raise ValueError("Make sure to provide `position_ids` when passing `past_key_values`.")
-
- position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
-
- if attention_mask is None:
- attention_mask = jnp.ones((batch_size, sequence_length))
-
- # Handle any PRNG if needed
- rngs = {}
- if dropout_rng is not None:
- rngs["dropout"] = dropout_rng
-
- inputs = {"params": params or self.params}
-
- # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be changed by FlaxGPTNeoAttention module
- if past_key_values:
- inputs["cache"] = past_key_values
- mutable = ["cache"]
- else:
- mutable = False
-
- outputs = self.module.apply(
- inputs,
- jnp.array(input_ids, dtype="i4"),
- jnp.array(attention_mask, dtype="i4"),
- jnp.array(position_ids, dtype="i4"),
- not train,
- False,
- output_attentions,
- output_hidden_states,
- return_dict,
- rngs=rngs,
- mutable=mutable,
- )
-
- # add updated cache to model output
- if past_key_values is not None and return_dict:
- outputs, past_key_values = outputs
- outputs["past_key_values"] = unfreeze(past_key_values["cache"])
- return outputs
- elif past_key_values is not None and not return_dict:
- outputs, past_key_values = outputs
- outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:]
-
- return outputs
-
-
-class FlaxGPTNeoBlockCollection(nn.Module):
- config: GPTNeoConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.blocks = [
- FlaxGPTNeoBlock(self.config, layer_id=i, name=str(i), dtype=self.dtype)
- for i in range(self.config.num_hidden_layers)
- ]
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- all_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
-
- for block in self.blocks:
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- layer_outputs = block(
- hidden_states,
- attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions += (layer_outputs[1],)
-
- # this contains possible `None` values - `FlaxGPTNeoModule` will filter them out
- outputs = (hidden_states, all_hidden_states, all_attentions)
-
- return outputs
-
-
-class FlaxGPTNeoModule(nn.Module):
- config: GPTNeoConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.embed_dim = self.config.hidden_size
- embedding_init = jax.nn.initializers.normal(stddev=self.config.initializer_range)
- self.wte = nn.Embed(
- self.config.vocab_size,
- self.embed_dim,
- embedding_init=embedding_init,
- )
- self.wpe = nn.Embed(
- self.config.max_position_embeddings,
- self.embed_dim,
- embedding_init=embedding_init,
- )
- self.dropout = nn.Dropout(rate=self.config.embed_dropout)
- self.h = FlaxGPTNeoBlockCollection(self.config, dtype=self.dtype)
- self.ln_f = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- position_ids,
- deterministic=True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- input_embeds = self.wte(input_ids.astype("i4"))
- position_embeds = self.wpe(position_ids.astype("i4"))
-
- hidden_states = input_embeds + position_embeds
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
-
- outputs = self.h(
- hidden_states,
- attention_mask,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = outputs[0]
- hidden_states = self.ln_f(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = outputs[1] + (hidden_states,)
- outputs = (hidden_states, all_hidden_states) + outputs[2:]
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- if not return_dict:
- return tuple(v for v in outputs if v is not None)
-
- return FlaxBaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=outputs[1],
- attentions=outputs[-1],
- )
-
-
-@add_start_docstrings(
- "The bare GPTNeo Model transformer outputting raw hidden-states without any specific head on top.",
- GPT_NEO_START_DOCSTRING,
-)
-class FlaxGPTNeoModel(FlaxGPTNeoPreTrainedModel):
- module_class = FlaxGPTNeoModule
-
-
-append_call_sample_docstring(FlaxGPTNeoModel, _CHECKPOINT_FOR_DOC, FlaxBaseModelOutput, _CONFIG_FOR_DOC)
-
-
-class FlaxGPTNeoForCausalLMModule(nn.Module):
- config: GPTNeoConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.transformer = FlaxGPTNeoModule(self.config, dtype=self.dtype)
- self.lm_head = nn.Dense(
- self.config.vocab_size,
- use_bias=False,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- position_ids,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- outputs = self.transformer(
- input_ids,
- attention_mask,
- position_ids,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
-
- if self.config.tie_word_embeddings:
- shared_kernel = self.transformer.variables["params"]["wte"]["embedding"].T
- lm_logits = self.lm_head.apply({"params": {"kernel": shared_kernel}}, hidden_states)
- else:
- lm_logits = self.lm_head(hidden_states)
-
- if not return_dict:
- return (lm_logits,) + outputs[1:]
-
- return FlaxCausalLMOutput(logits=lm_logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
-
-
-@add_start_docstrings(
- """
- The GPTNeo Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- GPT_NEO_START_DOCSTRING,
-)
-class FlaxGPTNeoForCausalLM(FlaxGPTNeoPreTrainedModel):
- module_class = FlaxGPTNeoForCausalLMModule
-
- def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jax.Array] = None):
- # initializing the cache
- batch_size, seq_length = input_ids.shape
-
- past_key_values = self.init_cache(batch_size, max_length)
- # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
- # But since GPTNeo uses a causal mask, those positions are masked anyways.
- # Thus we can create a single static attention_mask here, which is more efficient for compilation
- extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
- if attention_mask is not None:
- position_ids = attention_mask.cumsum(axis=-1) - 1
- extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, attention_mask, (0, 0))
- else:
- position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))
-
- return {
- "past_key_values": past_key_values,
- "attention_mask": extended_attention_mask,
- "position_ids": position_ids,
- }
-
- def update_inputs_for_generation(self, model_outputs, model_kwargs):
- model_kwargs["past_key_values"] = model_outputs.past_key_values
- model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1
- return model_kwargs
-
-
-append_call_sample_docstring(FlaxGPTNeoForCausalLM, _CHECKPOINT_FOR_DOC, FlaxCausalLMOutput, _CONFIG_FOR_DOC)
diff --git a/transformers/models/gpt_neo/modeling_gpt_neo.py b/transformers/models/gpt_neo/modeling_gpt_neo.py
deleted file mode 100644
index 2fbf4677ca6f442a1ba8d0a1cb515269f08ce2f5..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neo/modeling_gpt_neo.py
+++ /dev/null
@@ -1,1346 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The Eleuther AI and HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch GPT Neo model."""
-
-
-import os
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
-from ...modeling_outputs import (
- BaseModelOutputWithPast,
- BaseModelOutputWithPastAndCrossAttentions,
- CausalLMOutputWithCrossAttentions,
- CausalLMOutputWithPast,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutputWithPast,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...pytorch_utils import is_torch_greater_or_equal_than_1_13
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- is_torch_fx_available,
- logging,
-)
-from .configuration_gpt_neo import GPTNeoConfig
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-
-# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
-# It means that the function will not be traced through and simply appear as a node in the graph.
-if is_torch_fx_available():
- if not is_torch_greater_or_equal_than_1_13:
- import torch.fx
-
- _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "GPTNeoConfig"
-
-
-from ..deprecated._archive_maps import GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-_CHECKPOINT_FOR_DOC = "EleutherAI/gpt-neo-1.3B"
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-def load_tf_weights_in_gpt_neo(model, config, gpt_neo_checkpoint_path):
- """Load tf checkpoints in a pytorch model"""
- try:
- import re
-
- import tensorflow as tf
- except ImportError:
- logger.error(
- "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
- "https://www.tensorflow.org/install/ for installation instructions."
- )
- raise
- tf_path = os.path.abspath(gpt_neo_checkpoint_path)
- logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
- # Load weights from TF model
- init_vars = tf.train.list_variables(tf_path)
- names = []
- arrays = []
- for name, shape in init_vars:
- if "global_step" not in name and "adam" not in name:
- array = tf.train.load_variable(tf_path, name)
- array = tf.dtypes.cast(array.squeeze(), tf.float32).numpy()
- name = name.replace("attn/q", "attn/attention/q_proj/w")
- name = name.replace("attn/k", "attn/attention/k_proj/w")
- name = name.replace("attn/v", "attn/attention/v_proj/w")
- name = name.replace("attn/o", "attn/attention/out_proj/w")
- name = name.replace("norm_1", "ln_1")
- name = name.replace("norm_2", "ln_2")
- name = name.replace("attn/compute_output_bias/o_b", "attn/attention/out_proj/b")
- name = name.replace("conv1d_main/c_fc/kernel", "c_fc/w")
- name = name.replace("conv1d_main/c_fc/bias", "c_fc/b")
- name = name.replace("conv1d_main/c_proj/kernel", "c_proj/w")
- name = name.replace("conv1d_main/c_proj/bias", "c_proj/b")
-
- names.append(name)
- arrays.append(array)
-
- for name, array in zip(names, arrays):
- name = name[5:] # skip "gpt2/"
- name = name.split("/")
- pointer = model.transformer
- for m_name in name:
- if re.fullmatch(r"[A-Za-z]+\d+", m_name):
- scope_names = re.split(r"(\d+)", m_name)
- else:
- scope_names = [m_name]
- if scope_names[0] == "w" or scope_names[0] == "g":
- pointer = getattr(pointer, "weight")
- elif scope_names[0] == "b":
- pointer = getattr(pointer, "bias")
- elif scope_names[0] == "wpe" or scope_names[0] == "wte":
- pointer = getattr(pointer, scope_names[0])
- pointer = getattr(pointer, "weight")
- else:
- pointer = getattr(pointer, scope_names[0])
- if len(scope_names) >= 2:
- num = int(scope_names[1])
- pointer = pointer[num]
-
- if name[-1] == "w" and name[-2] in ["out_proj", "k_proj", "q_proj", "v_proj", "c_proj", "c_fc"]:
- array = array.transpose()
-
- if name == ["wte"]:
- # if vocab is padded, then trim off the padding embeddings
- array = array[: config.vocab_size]
-
- if pointer.shape != array.shape:
- raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched {name}")
-
- print(f"Initialize PyTorch weight {name}")
- pointer.data = torch.from_numpy(array)
-
- # init the final linear layer using word embeddings
- embs = model.transformer.wte.weight
- lin = nn.Linear(embs.size()[1], embs.size()[0], bias=False)
- lin.weight = embs
- model.set_output_embeddings(lin)
- return model
-
-
-class GPTNeoSelfAttention(nn.Module):
- def __init__(self, config, attention_type):
- super().__init__()
- self.config = config
-
- max_positions = config.max_position_embeddings
- bias = torch.tril(torch.ones((max_positions, max_positions), dtype=bool)).view(
- 1, 1, max_positions, max_positions
- )
-
- # local causal self attention is a sliding window where each token can only attend to the previous
- # window_size tokens. This is implemented by updating the causal mask such that for each token
- # all other tokens are masked except the previous window_size tokens.
- if attention_type == "local":
- bias = torch.bitwise_xor(bias, torch.tril(bias, -config.window_size))
-
- self.register_buffer("bias", bias, persistent=False)
- self.register_buffer("masked_bias", torch.tensor(-1e9), persistent=False)
-
- self.attn_dropout = nn.Dropout(float(config.attention_dropout))
- self.resid_dropout = nn.Dropout(float(config.resid_dropout))
- self.is_causal = True
-
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_heads
- self.head_dim = self.embed_dim // self.num_heads
- if self.head_dim * self.num_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
- f" {self.num_heads})."
- )
-
- self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
- self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
- self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
- self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
-
- def _split_heads(self, tensor, num_heads, attn_head_size):
- """
- Splits hidden_size dim into attn_head_size and num_heads
- """
- new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
- tensor = tensor.view(new_shape)
- return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)
-
- def _merge_heads(self, tensor, num_heads, attn_head_size):
- """
- Merges attn_head_size dim and num_attn_heads dim into hidden_size
- """
- tensor = tensor.permute(0, 2, 1, 3).contiguous()
- new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
- return tensor.view(new_shape)
-
- def _attn(self, query, key, value, attention_mask=None, head_mask=None):
- # Keep the attention weights computation in fp32 to avoid overflow issues
- query = query.to(torch.float32)
- key = key.to(torch.float32)
-
- attn_weights = torch.matmul(query, key.transpose(-1, -2))
-
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
- mask_value = torch.finfo(attn_weights.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
- attn_weights = torch.where(causal_mask, attn_weights, mask_value)
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
- attn_weights = attn_weights.to(value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- layer_past=None,
- head_mask=None,
- use_cache=False,
- output_attentions=False,
- ):
- query = self.q_proj(hidden_states)
- key = self.k_proj(hidden_states)
- value = self.v_proj(hidden_states)
-
- query = self._split_heads(query, self.num_heads, self.head_dim)
- key = self._split_heads(key, self.num_heads, self.head_dim)
- value = self._split_heads(value, self.num_heads, self.head_dim)
-
- if layer_past is not None:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- if use_cache is True:
- present = (key, value)
- else:
- present = None
-
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
- attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
- attn_output = self.out_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs # a, present, (attentions)
-
-
-class GPTNeoFlashAttention2(GPTNeoSelfAttention):
- """
- GPTNeo flash attention module. This module inherits from `GPTNeoSelfAttention` as the weights of the module stays
- untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
- flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- layer_past=None,
- head_mask=None,
- use_cache=False,
- output_attentions=False,
- ):
- bsz, _, _ = hidden_states.size()
-
- query = self.q_proj(hidden_states)
- key = self.k_proj(hidden_states)
- value = self.v_proj(hidden_states)
-
- query = self._split_heads(query, self.num_heads, self.head_dim)
- key = self._split_heads(key, self.num_heads, self.head_dim)
- value = self._split_heads(value, self.num_heads, self.head_dim)
-
- if layer_past is not None:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- if use_cache is True:
- present = (key, value)
- else:
- present = None
-
- query_length = query.shape[2]
- tgt_len = key.shape[2]
-
- # Flash attention requires the input to have the shape
- # batch_size x seq_length x head_dim x hidden_dim
- query = query.transpose(1, 2).view(bsz, query_length, self.num_heads, self.head_dim)
- key = key.transpose(1, 2).view(bsz, tgt_len, self.num_heads, self.head_dim)
- value = value.transpose(1, 2).view(bsz, tgt_len, self.num_heads, self.head_dim)
-
- attn_dropout = self.config.attention_dropout if self.training else 0.0
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in the correct dtype just to be sure everything works as expected.
- # This might slowdown training & inference so it is recommended to not cast the LayerNorms
- # in fp32. (LlamaRMSNorm handles it correctly)
-
- if query.dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.q_proj.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
-
- query = query.to(target_dtype)
- key = key.to(target_dtype)
- value = value.to(target_dtype)
-
- attn_output = self._flash_attention_forward(
- query, key, value, attention_mask, query_length, dropout=attn_dropout, softmax_scale=1.0
- )
-
- attn_weights_reshaped = attn_output.reshape(bsz, query_length, self.num_heads * self.head_dim)
- attn_output = self.out_proj(attn_weights_reshaped)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights_reshaped,)
-
- return outputs
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-GPT_NEO_ATTENTION_CLASSES = {
- "eager": GPTNeoSelfAttention,
- "flash_attention_2": GPTNeoFlashAttention2,
-}
-
-
-class GPTNeoAttention(nn.Module):
- def __init__(self, config, layer_id=0):
- super().__init__()
- self.layer_id = layer_id
- self.attention_layers = config.attention_layers
- self.attention_type = self.attention_layers[layer_id]
-
- if self.attention_type in ["global", "local"]:
- self.attention = GPT_NEO_ATTENTION_CLASSES[config._attn_implementation](config, self.attention_type)
- else:
- raise NotImplementedError(
- "Only attn layer types 'global' and 'local' exist, but got `config.attention_layers`: "
- f"{config.attention_layers}. Select attn layer types from ['global', 'local'] only."
- )
-
- def forward(
- self,
- hidden_states,
- layer_past=None,
- attention_mask=None,
- head_mask=None,
- use_cache=False,
- output_attentions=False,
- ):
- return self.attention(
- hidden_states,
- attention_mask=attention_mask,
- layer_past=layer_past,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
-
-
-class GPTNeoMLP(nn.Module):
- def __init__(self, intermediate_size, config): # in MLP: intermediate_size= 4 * hidden_size
- super().__init__()
- embed_dim = config.hidden_size
- self.c_fc = nn.Linear(embed_dim, intermediate_size)
- self.c_proj = nn.Linear(intermediate_size, embed_dim)
- self.act = ACT2FN[config.activation_function]
- self.dropout = nn.Dropout(float(config.resid_dropout))
-
- def forward(self, hidden_states):
- hidden_states = self.c_fc(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.c_proj(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-class GPTNeoBlock(nn.Module):
- def __init__(self, config, layer_id):
- super().__init__()
- hidden_size = config.hidden_size
- inner_dim = config.intermediate_size if config.intermediate_size is not None else 4 * hidden_size
- self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
- self.attn = GPTNeoAttention(config, layer_id)
- self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
- self.mlp = GPTNeoMLP(inner_dim, config)
-
- def forward(
- self,
- hidden_states,
- layer_past=None,
- attention_mask=None,
- head_mask=None,
- use_cache=False,
- output_attentions=False,
- ):
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
- outputs = attn_outputs[1:]
- # residual connection
- hidden_states = attn_output + residual
-
- residual = hidden_states
- hidden_states = self.ln_2(hidden_states)
- feed_forward_hidden_states = self.mlp(hidden_states)
- # residual connection
- hidden_states = residual + feed_forward_hidden_states
-
- if use_cache:
- outputs = (hidden_states,) + outputs
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- return outputs # hidden_states, present, (attentions, cross_attentions)
-
-
-class GPTNeoPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTNeoConfig
- load_tf_weights = load_tf_weights_in_gpt_neo
- base_model_prefix = "transformer"
- supports_gradient_checkpointing = True
- _no_split_modules = ["GPTNeoBlock"]
- _skip_keys_device_placement = "past_key_values"
- _supports_flash_attn_2 = True
-
- def __init__(self, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, (nn.Linear,)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-GPT_NEO_START_DOCSTRING = r"""
-
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`GPTNeoConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPT_NEO_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
- `past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input
- sequence tokens in the vocabulary.
-
- If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
- `input_ids`.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.num_layers`):
- Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
- their past given to this model should not be passed as `input_ids` as they have already been computed.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
-
- If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see
- `past_key_values`).
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare GPT Neo Model transformer outputting raw hidden-states without any specific head on top.",
- GPT_NEO_START_DOCSTRING,
-)
-class GPTNeoModel(GPTNeoPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.embed_dim = config.hidden_size
- self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
- self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
- self.drop = nn.Dropout(float(config.embed_dropout))
- self.h = nn.ModuleList([GPTNeoBlock(config, layer_id=i) for i in range(config.num_layers)])
- self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
-
- self.gradient_checkpointing = False
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.wte
-
- def set_input_embeddings(self, new_embeddings):
- self.wte = new_embeddings
-
- @add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPastAndCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- input_ids = input_ids.view(-1, input_shape[-1])
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if token_type_ids is not None:
- token_type_ids = token_type_ids.view(-1, input_shape[-1])
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * len(self.h))
- else:
- past_length = past_key_values[0][0].size(-2)
-
- if position_ids is None:
- position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0)
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x num_heads x N x N
- # head_mask has shape n_layer x batch x num_heads x N x N
- head_mask = self.get_head_mask(head_mask, self.config.num_layers)
-
- if inputs_embeds is None:
- inputs_embeds = self.wte(input_ids)
- position_embeds = self.wpe(position_ids)
- hidden_states = inputs_embeds + position_embeds
-
- # Attention mask.
- if self._use_flash_attention_2:
- # 2d mask is passed through the layers
- attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
- else:
- # 4d mask is passed through the layers
- attention_mask = _prepare_4d_causal_attention_mask(attention_mask, input_shape, inputs_embeds, past_length)
-
- if token_type_ids is not None:
- token_type_embeds = self.wte(token_type_ids)
- hidden_states = hidden_states + token_type_embeds
-
- hidden_states = self.drop(hidden_states)
-
- output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- presents = () if use_cache else None
- all_self_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- None,
- attention_mask,
- head_mask[i],
- use_cache,
- output_attentions,
- )
- else:
- outputs = block(
- hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- head_mask=head_mask[i],
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
-
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
-
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = hidden_states.view(output_shape)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
-
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT Neo Model transformer with a language modeling head on top (linear layer with weights tied to the input
- embeddings).
- """,
- GPT_NEO_START_DOCSTRING,
-)
-class GPTNeoForCausalLM(GPTNeoPreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.transformer = GPTNeoModel(config)
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # Omit tokens covered by past_key_values
- if past_key_values:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
- if token_type_ids is not None:
- token_type_ids = token_type_ids[:, -input_ids.shape[1] :]
-
- attention_mask = kwargs.get("attention_mask", None)
- position_ids = kwargs.get("position_ids", None)
-
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- model_inputs = {"input_ids": input_ids}
-
- model_inputs.update(
- {
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "position_ids": position_ids,
- "attention_mask": attention_mask,
- "token_type_ids": token_type_ids,
- }
- )
-
- return model_inputs
-
- @add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutputWithCrossAttentions,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
-
- lm_logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # Compute loss in fp32 to match with mesh-tf version
- # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179
- lm_logits = lm_logits.to(torch.float32)
-
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
- lm_logits = lm_logits.to(hidden_states.dtype)
- loss = loss.to(hidden_states.dtype)
-
- if not return_dict:
- output = (lm_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=lm_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- @staticmethod
- def _reorder_cache(
- past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
- ) -> Tuple[Tuple[torch.Tensor]]:
- """
- This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
- [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
- beam_idx at every generation step.
- """
- return tuple(
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
- for layer_past in past_key_values
- )
-
-
-@add_start_docstrings(
- """
- The GPTNeo Model transformer with a sequence classification head on top (linear layer).
-
- [`GPTNeoForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT-1) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- GPT_NEO_START_DOCSTRING,
-)
-class GPTNeoForSequenceClassification(GPTNeoPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = GPTNeoModel(config)
- self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.Tensor] = None,
- past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- token_type_ids: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- head_mask: Optional[torch.Tensor] = None,
- inputs_embeds: Optional[torch.Tensor] = None,
- labels: Optional[torch.Tensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size, sequence_length = input_ids.shape[:2]
- else:
- batch_size, sequence_length = inputs_embeds.shape[:2]
-
- if self.config.pad_token_id is None and batch_size != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
-
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
- loss = None
- if labels is not None:
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- GPT Neo model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
- Named-Entity-Recognition (NER) tasks.
- """,
- GPT_NEO_START_DOCSTRING,
-)
-class GPTNeoForTokenClassification(GPTNeoPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.transformer = GPTNeoModel(config)
- self.dropout = nn.Dropout(config.classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint="EleutherAI/gpt-neo-125m",
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_loss=0.25,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = transformer_outputs[0]
- hidden_states = self.dropout(hidden_states)
- logits = self.classifier(hidden_states)
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + transformer_outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT-Neo Model transformer with a span classification head on top for extractive question-answering tasks like
- SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- GPT_NEO_START_DOCSTRING,
-)
-class GPTNeoForQuestionAnswering(GPTNeoPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = GPTNeoModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, 2)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- real_checkpoint=_CHECKPOINT_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- start_positions: Optional[torch.LongTensor] = None,
- end_positions: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/gpt_neox/__init__.py b/transformers/models/gpt_neox/__init__.py
deleted file mode 100644
index 46f06b1991afe78c5fc58c14ef3c68a75c49e0f4..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neox/__init__.py
+++ /dev/null
@@ -1,80 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...file_utils import _LazyModule, is_tokenizers_available, is_torch_available
-from ...utils import OptionalDependencyNotAvailable
-
-
-_import_structure = {"configuration_gpt_neox": ["GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTNeoXConfig"]}
-
-try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_gpt_neox_fast"] = ["GPTNeoXTokenizerFast"]
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_gpt_neox"] = [
- "GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GPTNeoXForCausalLM",
- "GPTNeoXForQuestionAnswering",
- "GPTNeoXForSequenceClassification",
- "GPTNeoXForTokenClassification",
- "GPTNeoXLayer",
- "GPTNeoXModel",
- "GPTNeoXPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_gpt_neox import GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXConfig
-
- try:
- if not is_tokenizers_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_gpt_neox_fast import GPTNeoXTokenizerFast
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_gpt_neox import (
- GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST,
- GPTNeoXForCausalLM,
- GPTNeoXForQuestionAnswering,
- GPTNeoXForSequenceClassification,
- GPTNeoXForTokenClassification,
- GPTNeoXLayer,
- GPTNeoXModel,
- GPTNeoXPreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gpt_neox/__pycache__/__init__.cpython-310.pyc b/transformers/models/gpt_neox/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 48537c5382c99cbd95d908fbf438827def57bc39..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neox/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neox/__pycache__/configuration_gpt_neox.cpython-310.pyc b/transformers/models/gpt_neox/__pycache__/configuration_gpt_neox.cpython-310.pyc
deleted file mode 100644
index c6d299ad812ffec77b492f044cb39005a3df0523..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neox/__pycache__/configuration_gpt_neox.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neox/__pycache__/modeling_gpt_neox.cpython-310.pyc b/transformers/models/gpt_neox/__pycache__/modeling_gpt_neox.cpython-310.pyc
deleted file mode 100644
index e4b482257740a38c2bcfdae11f4dc41c825ee010..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neox/__pycache__/modeling_gpt_neox.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neox/__pycache__/tokenization_gpt_neox_fast.cpython-310.pyc b/transformers/models/gpt_neox/__pycache__/tokenization_gpt_neox_fast.cpython-310.pyc
deleted file mode 100644
index 7a3a980368f883e128157050f4386f83bfdd6d98..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neox/__pycache__/tokenization_gpt_neox_fast.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neox/configuration_gpt_neox.py b/transformers/models/gpt_neox/configuration_gpt_neox.py
deleted file mode 100644
index 7f583f139448f9a0b147f7df306be5a293358803..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neox/configuration_gpt_neox.py
+++ /dev/null
@@ -1,179 +0,0 @@
-# coding=utf-8
-# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPTNeoX model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GPTNeoXConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GPTNeoXModel`]. It is used to instantiate an
- GPTNeoX model according to the specified arguments, defining the model architecture. Instantiating a configuration
- with the defaults will yield a similar configuration to that of the GPTNeoX
- [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
-
- Args:
- vocab_size (`int`, *optional*, defaults to 50432):
- Vocabulary size of the GPTNeoX model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`GPTNeoXModel`].
- hidden_size (`int`, *optional*, defaults to 6144):
- Dimension of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 44):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 64):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, *optional*, defaults to 24576):
- Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
- `"relu"`, `"selu"` and `"gelu_new"` are supported.
- rotary_pct (`float`, *optional*, defaults to 0.25):
- percentage of hidden dimensions to allocate to rotary embeddings
- rotary_emb_base (`int`, *optional*, defaults to 10000)
- base for computing rotary embeddings frequency
- attention_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio probability of the attention score.
- hidden_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio of (1) the word embeddings, (2) the post-attention hidden states, and (3) the post-mlp
- hidden states.
- classifier_dropout (`float`, *optional*, defaults to 0.1):
- Argument used when doing token classification, used in the model [`GPTNeoXForTokenClassification`].
-
- The dropout ratio for the hidden layer.
- max_position_embeddings (`int`, *optional*, defaults to 2048):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- initializer_range (`float`, *optional*, defaults to 1e-5):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-12):
- The epsilon used by the layer normalization layers.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- use_parallel_residual (`bool`, *optional*, defaults to `True`):
- Whether to use a "parallel" formulation in each Transformer layer, which can provide a slight training
- speedup at large scales (e.g. 20B).
- rope_scaling (`Dict`, *optional*):
- Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
- strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
- `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
- `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
- these scaling strategies behave:
- https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
- experimental feature, subject to breaking API changes in future versions.
- attention_bias (`bool`, *optional*, defaults to `True`):
- Whether to use a bias in the query, key, value and output projection layers during self-attention.
-
- Example:
-
- ```python
- >>> from transformers import GPTNeoXConfig, GPTNeoXModel
-
- >>> # Initializing a GPTNeoX gpt-neox-20b style configuration
- >>> configuration = GPTNeoXConfig()
-
- >>> # Initializing a model (with random weights) from the gpt-neox-20b style configuration
- >>> model = GPTNeoXModel(configuration) # doctest: +SKIP
-
- >>> # Accessing the model configuration
- >>> configuration = model.config # doctest: +SKIP
- ```"""
-
- model_type = "gpt_neox"
- keys_to_ignore_at_inference = ["past_key_values"]
-
- def __init__(
- self,
- vocab_size=50432,
- hidden_size=6144,
- num_hidden_layers=44,
- num_attention_heads=64,
- intermediate_size=24576,
- hidden_act="gelu",
- rotary_pct=0.25,
- rotary_emb_base=10000,
- attention_dropout=0.0,
- hidden_dropout=0.0,
- classifier_dropout=0.1,
- max_position_embeddings=2048,
- initializer_range=0.02,
- layer_norm_eps=1e-5,
- use_cache=True,
- bos_token_id=0,
- eos_token_id=2,
- tie_word_embeddings=False,
- use_parallel_residual=True,
- rope_scaling=None,
- attention_bias=True,
- **kwargs,
- ):
- super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_size = intermediate_size
- self.hidden_act = hidden_act
- self.rotary_pct = rotary_pct
- self.rotary_emb_base = rotary_emb_base
- self.attention_dropout = attention_dropout
- self.hidden_dropout = hidden_dropout
- self.classifier_dropout = classifier_dropout
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.use_cache = use_cache
- self.tie_word_embeddings = tie_word_embeddings
- self.use_parallel_residual = use_parallel_residual
- self.rope_scaling = rope_scaling
- self.attention_bias = attention_bias
- self._rope_scaling_validation()
-
- if self.hidden_size % self.num_attention_heads != 0:
- raise ValueError(
- "The hidden size is not divisble by the number of attention heads! Make sure to update them!"
- )
-
- # Copied from transformers.models.llama.configuration_llama.LlamaConfig._rope_scaling_validation
- def _rope_scaling_validation(self):
- """
- Validate the `rope_scaling` configuration.
- """
- if self.rope_scaling is None:
- return
-
- if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
- raise ValueError(
- "`rope_scaling` must be a dictionary with two fields, `type` and `factor`, " f"got {self.rope_scaling}"
- )
- rope_scaling_type = self.rope_scaling.get("type", None)
- rope_scaling_factor = self.rope_scaling.get("factor", None)
- if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
- raise ValueError(
- f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
- )
- if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
- raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
diff --git a/transformers/models/gpt_neox/modeling_gpt_neox.py b/transformers/models/gpt_neox/modeling_gpt_neox.py
deleted file mode 100644
index 83c99202ac9379ed023c2ac36d87a6d195433328..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neox/modeling_gpt_neox.py
+++ /dev/null
@@ -1,1426 +0,0 @@
-# coding=utf-8
-# Copyright 2022 EleutherAI The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch GPTNeoX model."""
-
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-from torch.nn import functional as F
-
-from ...activations import ACT2FN
-from ...file_utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- replace_return_docstrings,
-)
-from ...modeling_outputs import (
- BaseModelOutputWithPast,
- CausalLMOutputWithPast,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutputWithPast,
- TokenClassifierOutput,
-)
-from ...modeling_utils import PreTrainedModel
-from ...utils import is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10, logging
-from .configuration_gpt_neox import GPTNeoXConfig
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "trl-internal-testing/tiny-random-GPTNeoXForCausalLM"
-_REAL_CHECKPOINT_FOR_DOC = "EleutherAI/gpt-neox-20b"
-_CONFIG_FOR_DOC = "GPTNeoXConfig"
-
-
-from ..deprecated._archive_maps import GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-class GPTNeoXPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTNeoXConfig
- base_model_prefix = "gpt_neox"
- supports_gradient_checkpointing = True
- _no_split_modules = ["GPTNeoXLayer"]
- _skip_keys_device_placement = "past_key_values"
- _supports_flash_attn_2 = True
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-class GPTNeoXAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- self.num_attention_heads = config.num_attention_heads
- self.hidden_size = config.hidden_size
- if self.hidden_size % self.num_attention_heads != 0:
- raise ValueError(
- "The hidden size is not divisble by the number of attention heads! Make sure to update them"
- )
- self.head_size = self.hidden_size // self.num_attention_heads
- self.rotary_ndims = int(self.head_size * config.rotary_pct)
- self._init_bias(config.max_position_embeddings)
-
- self.register_buffer("masked_bias", torch.tensor(-1e9), persistent=False)
- self._init_rope()
-
- self.norm_factor = self.head_size**-0.5
- self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=config.attention_bias)
- self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
- self.attention_dropout = nn.Dropout(config.attention_dropout)
- self.is_causal = True
-
- def _init_bias(self, max_positions, device=None):
- self.register_buffer(
- "bias",
- torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
- 1, 1, max_positions, max_positions
- ),
- persistent=False,
- )
- if device is not None:
- self.bias = self.bias.to(device)
-
- def _init_rope(self):
- if self.config.rope_scaling is None:
- self.rotary_emb = GPTNeoXRotaryEmbedding(
- self.rotary_ndims, self.config.max_position_embeddings, base=self.config.rotary_emb_base
- )
- else:
- scaling_type = self.config.rope_scaling["type"]
- scaling_factor = self.config.rope_scaling["factor"]
- if scaling_type == "linear":
- self.rotary_emb = GPTNeoXLinearScalingRotaryEmbedding(
- self.rotary_ndims,
- self.config.max_position_embeddings,
- base=self.config.rotary_emb_base,
- scaling_factor=scaling_factor,
- )
- elif scaling_type == "dynamic":
- self.rotary_emb = GPTNeoXDynamicNTKScalingRotaryEmbedding(
- self.rotary_ndims,
- self.config.max_position_embeddings,
- base=self.config.rotary_emb_base,
- scaling_factor=scaling_factor,
- )
- else:
- raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- attention_mask: torch.FloatTensor,
- position_ids: torch.LongTensor,
- head_mask: Optional[torch.FloatTensor] = None,
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- padding_mask: Optional[torch.Tensor] = None,
- ):
- has_layer_past = layer_past is not None
-
- # Compute QKV
- # Attention heads [batch, seq_len, hidden_size]
- # --> [batch, seq_len, (np * 3 * head_size)]
- qkv = self.query_key_value(hidden_states)
-
- # [batch, seq_len, (num_heads * 3 * head_size)]
- # --> [batch, seq_len, num_heads, 3 * head_size]
- new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size)
- qkv = qkv.view(*new_qkv_shape)
-
- # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
- query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
- key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3)
- value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3)
-
- # Compute rotary embeddings on rotary_ndims
- query_rot = query[..., : self.rotary_ndims]
- query_pass = query[..., self.rotary_ndims :]
- key_rot = key[..., : self.rotary_ndims]
- key_pass = key[..., self.rotary_ndims :]
-
- # Compute token offset for rotary embeddings (when decoding)
- seq_len = key.shape[-2]
- if has_layer_past:
- seq_len += layer_past[0].shape[-2]
- cos, sin = self.rotary_emb(value, seq_len=seq_len)
- query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
- query = torch.cat((query, query_pass), dim=-1)
- key = torch.cat((key, key_pass), dim=-1)
-
- # Cache QKV values
- if has_layer_past:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
- present = (key, value) if use_cache else None
-
- # Compute attention
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
- # Reshape outputs
- attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size)
- attn_output = self.dense(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
- @classmethod
- def _split_heads(cls, tensor, num_attention_heads, attn_head_size):
- """
- Splits hidden dim into attn_head_size and num_attention_heads
- """
- # tensor: [bs, seq_len, hidden_size]
- new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
- # -> [bs, seq_len, num_attention_heads, attn_head_size]
- tensor = tensor.view(new_shape)
- # -> [bs, num_attention_heads, seq_len, attn_head_size]
- tensor = tensor.permute(0, 2, 1, 3)
- return tensor
-
- @classmethod
- def _merge_heads(cls, tensor, num_attention_heads, attn_head_size):
- """
- Merges attn_head_size dim and num_attn_heads dim into hidden dim
- """
- # tensor [bs, num_attention_heads, seq_len, attn_head_size]
- tensor = tensor.permute(0, 2, 1, 3).contiguous()
- # -> [bs, seq_len, num_attention_heads, attn_head_size]
- tensor = tensor.view(tensor.size(0), tensor.size(1), num_attention_heads * attn_head_size)
- # -> [bs, seq_len, hidden_size]
- return tensor
-
- def _attn(self, query, key, value, attention_mask=None, head_mask=None):
- # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size]
- # compute causal mask from causal mask buffer
- batch_size, num_attention_heads, query_length, attn_head_size = query.size()
- key_length = key.size(-2)
-
- # dynamically increase the causal mask with the key length, if needed.
- if key_length > self.bias.shape[-1]:
- self._init_bias(key_length, device=key.device)
- causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
-
- query = query.view(batch_size * num_attention_heads, query_length, attn_head_size)
- key = key.view(batch_size * num_attention_heads, key_length, attn_head_size)
- attn_scores = torch.zeros(
- batch_size * num_attention_heads,
- query_length,
- key_length,
- dtype=query.dtype,
- device=key.device,
- )
- attn_scores = torch.baddbmm(
- attn_scores,
- query,
- key.transpose(1, 2),
- beta=1.0,
- alpha=self.norm_factor,
- )
- attn_scores = attn_scores.view(batch_size, num_attention_heads, query_length, key_length)
-
- mask_value = torch.finfo(attn_scores.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.tensor(mask_value, dtype=attn_scores.dtype).to(attn_scores.device)
- attn_scores = torch.where(causal_mask, attn_scores, mask_value)
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_scores = attn_scores + attention_mask
-
- attn_weights = nn.functional.softmax(attn_scores, dim=-1)
- attn_weights = attn_weights.to(value.dtype)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_weights = self.attention_dropout(attn_weights)
-
- attn_output = torch.matmul(attn_weights, value)
- return attn_output, attn_weights
-
-
-class GPTNeoXFlashAttention2(GPTNeoXAttention):
- """
- GPTNeoX flash attention module. This module inherits from `GPTNeoXAttention` as the weights of the module stays
- untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
- flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- attention_mask: torch.FloatTensor,
- position_ids: torch.LongTensor,
- head_mask: Optional[torch.FloatTensor] = None,
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ):
- has_layer_past = layer_past is not None
-
- # Compute QKV
- # Attention heads [batch, seq_len, hidden_size]
- # --> [batch, seq_len, (np * 3 * head_size)]
- qkv = self.query_key_value(hidden_states)
-
- # [batch, seq_len, (num_heads * 3 * head_size)]
- # --> [batch, seq_len, num_heads, 3 * head_size]
- new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size)
- qkv = qkv.view(*new_qkv_shape)
-
- # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
- query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
- key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3)
- value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3)
-
- query_length = query.shape[-2]
-
- # Compute rotary embeddings on rotary_ndims
- query_rot = query[..., : self.rotary_ndims]
- query_pass = query[..., self.rotary_ndims :]
- key_rot = key[..., : self.rotary_ndims]
- key_pass = key[..., self.rotary_ndims :]
-
- # Compute token offset for rotary embeddings (when decoding)
- seq_len = key.shape[-2]
- if has_layer_past:
- seq_len += layer_past[0].shape[-2]
- cos, sin = self.rotary_emb(value, seq_len=seq_len)
- query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
- query = torch.cat((query, query_pass), dim=-1)
- key = torch.cat((key, key_pass), dim=-1)
-
- # Cache QKV values
- if has_layer_past:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
- present = (key, value) if use_cache else None
-
- # GPT-neo-X casts query and key in fp32 to apply rotary embedding in full precision
- target_dtype = value.dtype
- if query.dtype != target_dtype:
- query = query.to(target_dtype)
- if key.dtype != target_dtype:
- key = key.to(target_dtype)
-
- # Permute to get the expected shape for Flash Attention
- query = query.permute(0, 2, 1, 3)
- key = key.permute(0, 2, 1, 3)
- value = value.permute(0, 2, 1, 3)
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in float16 / bfloat16 just to be sure everything works as expected.
- # This might slowdown training & inference so it is recommended to not cast the LayerNorms
- input_dtype = query.dtype
- if input_dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.query_key_value.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
-
- query = query.to(target_dtype)
- key = key.to(target_dtype)
- value = value.to(target_dtype)
-
- attention_dropout = self.config.attention_dropout if self.training else 0.0
-
- # Compute attention
- attn_weights = self._flash_attention_forward(
- query, key, value, attention_mask, query_length, dropout=attention_dropout, softmax_scale=self.norm_factor
- )
-
- # Reshape outputs
- attn_output = attn_weights.reshape(
- attn_weights.shape[0], attn_weights.shape[1], self.num_attention_heads * self.head_size
- )
- attn_output = self.dense(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input with num_heads->num_attention_heads
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_attention_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-def attention_mask_func(attention_scores, ltor_mask):
- attention_scores.masked_fill_(~ltor_mask, torch.finfo(attention_scores.dtype).min)
- return attention_scores
-
-
-class GPTNeoXRotaryEmbedding(nn.Module):
- # Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding.__init__
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
- super().__init__()
-
- self.dim = dim
- self.max_position_embeddings = max_position_embeddings
- self.base = base
- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
-
- # Build here to make `torch.jit.trace` work.
- self._set_cos_sin_cache(
- seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
- )
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos(), persistent=False)
- self.register_buffer("sin_cached", emb.sin(), persistent=False)
-
- def forward(self, x, seq_len=None):
- # x: [bs, num_attention_heads, seq_len, head_size]
- if seq_len > self.max_seq_len_cached:
- self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-
- return (
- self.cos_cached[:seq_len],
- self.sin_cached[:seq_len],
- )
-
-
-# copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding.__init__
-# TODO @gante bring compatibility back
-class GPTNeoXLinearScalingRotaryEmbedding(GPTNeoXRotaryEmbedding):
- """GPTNeoXRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
-
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
- self.scaling_factor = scaling_factor
- super().__init__(dim, max_position_embeddings, base, device)
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
- t = t / self.scaling_factor
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos(), persistent=False)
- self.register_buffer("sin_cached", emb.sin(), persistent=False)
-
-
-class GPTNeoXDynamicNTKScalingRotaryEmbedding(GPTNeoXRotaryEmbedding):
- """GPTNeoXRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
-
- # copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding.__init__
- # TODO @gante no longer copied from
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
- self.scaling_factor = scaling_factor
- super().__init__(dim, max_position_embeddings, base, device)
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
-
- if seq_len > self.max_position_embeddings:
- base = self.base * (
- (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
- ) ** (self.dim / (self.dim - 2))
- inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
-
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos(), persistent=False)
- self.register_buffer("sin_cached", emb.sin(), persistent=False)
-
-
-def rotate_half(x):
- """Rotates half the hidden dims of the input."""
- x1 = x[..., : x.shape[-1] // 2]
- x2 = x[..., x.shape[-1] // 2 :]
- return torch.cat((-x2, x1), dim=-1)
-
-
-# Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
- """Applies Rotary Position Embedding to the query and key tensors.
-
- Args:
- q (`torch.Tensor`): The query tensor.
- k (`torch.Tensor`): The key tensor.
- cos (`torch.Tensor`): The cosine part of the rotary embedding.
- sin (`torch.Tensor`): The sine part of the rotary embedding.
- position_ids (`torch.Tensor`):
- The position indices of the tokens corresponding to the query and key tensors. For example, this can be
- used to pass offsetted position ids when working with a KV-cache.
- unsqueeze_dim (`int`, *optional*, defaults to 1):
- The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
- sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
- that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
- k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
- cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
- the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
- Returns:
- `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
- """
- cos = cos[position_ids].unsqueeze(unsqueeze_dim)
- sin = sin[position_ids].unsqueeze(unsqueeze_dim)
- q_embed = (q * cos) + (rotate_half(q) * sin)
- k_embed = (k * cos) + (rotate_half(k) * sin)
- return q_embed, k_embed
-
-
-class GPTNeoXMLP(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size)
- self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size)
- self.act = ACT2FN[config.hidden_act]
-
- def forward(self, hidden_states):
- hidden_states = self.dense_h_to_4h(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.dense_4h_to_h(hidden_states)
- return hidden_states
-
-
-GPT_NEOX_ATTENTION_CLASSES = {
- "eager": GPTNeoXAttention,
- "flash_attention_2": GPTNeoXFlashAttention2,
-}
-
-
-class GPTNeoXLayer(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.use_parallel_residual = config.use_parallel_residual
- self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.post_attention_dropout = nn.Dropout(config.hidden_dropout)
- self.post_mlp_dropout = nn.Dropout(config.hidden_dropout)
- self.attention = GPT_NEOX_ATTENTION_CLASSES[config._attn_implementation](config)
- self.mlp = GPTNeoXMLP(config)
-
- def forward(
- self,
- hidden_states: Optional[torch.FloatTensor],
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: Optional[bool] = False,
- ):
- attention_layer_outputs = self.attention(
- self.input_layernorm(hidden_states),
- attention_mask=attention_mask,
- position_ids=position_ids,
- layer_past=layer_past,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attention_layer_outputs[0] # output_attn: attn_output, present, (attn_weights)
- attn_output = self.post_attention_dropout(attn_output)
- outputs = attention_layer_outputs[1:]
-
- if self.use_parallel_residual:
- # pseudocode:
- # x = x + attn(ln1(x)) + mlp(ln2(x))
- mlp_output = self.mlp(self.post_attention_layernorm(hidden_states))
- mlp_output = self.post_mlp_dropout(mlp_output)
- hidden_states = mlp_output + attn_output + hidden_states
- else:
- # pseudocode:
- # x = x + attn(ln1(x))
- # x = x + mlp(ln2(x))
- attn_output = attn_output + hidden_states
- mlp_output = self.mlp(self.post_attention_layernorm(attn_output))
- mlp_output = self.post_mlp_dropout(mlp_output)
- hidden_states = mlp_output + attn_output
-
- if use_cache:
- outputs = (hidden_states,) + outputs # hidden_states, present, (attn_weights)
- else:
- outputs = (hidden_states,) + outputs[1:] # hidden_states, (attn_weights)
-
- return outputs
-
-
-GPT_NEOX_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`~GPTNeoXConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPT_NEOX_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare GPTNeoX Model transformer outputting raw hidden-states without any specific head on top.",
- GPT_NEOX_START_DOCSTRING,
-)
-class GPTNeoXModel(GPTNeoXPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.config = config
-
- self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
- self.emb_dropout = nn.Dropout(config.hidden_dropout)
- self.layers = nn.ModuleList([GPTNeoXLayer(config) for _ in range(config.num_hidden_layers)])
- self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
-
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embed_in
-
- def set_input_embeddings(self, value):
- self.embed_in = value
-
- @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPast]:
- r"""
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- use_cache = use_cache if use_cache is not None else self.config.use_cache
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * self.config.num_hidden_layers)
- else:
- past_length = past_key_values[0][0].size(-2)
-
- if position_ids is None:
- device = input_ids.device if input_ids is not None else inputs_embeds.device
- position_ids = torch.arange(past_length, seq_length + past_length, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0)
-
- # Attention mask.
- if attention_mask is not None:
- assert batch_size > 0, "batch_size has to be defined and > 0"
- attention_mask = attention_mask.view(batch_size, -1)
- if self._use_flash_attention_2:
- attention_mask = attention_mask if 0 in attention_mask else None
- else:
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask = attention_mask[:, None, None, :]
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and the dtype's smallest value for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
- attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- if inputs_embeds is None:
- inputs_embeds = self.embed_in(input_ids)
-
- hidden_states = self.emb_dropout(inputs_embeds)
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- presents = () if use_cache else None
- all_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
- for i, (layer, layer_past) in enumerate(zip(self.layers, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- layer.__call__,
- hidden_states,
- attention_mask,
- position_ids,
- head_mask[i],
- use_cache,
- None,
- output_attentions,
- )
- else:
- outputs = layer(
- hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask[i],
- layer_past=layer_past,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
- if output_attentions:
- all_attentions = all_attentions + (outputs[2 if use_cache else 1],)
-
- hidden_states = self.final_layer_norm(hidden_states)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)
-
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- )
-
-
-@add_start_docstrings(
- """GPTNeoX Model with a `language modeling` head on top for CLM fine-tuning.""", GPT_NEOX_START_DOCSTRING
-)
-class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel):
- _tied_weights_keys = ["embed_out.weight"]
-
- def __init__(self, config):
- super().__init__(config)
-
- self.gpt_neox = GPTNeoXModel(config)
- self.embed_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.embed_out
-
- def set_output_embeddings(self, new_embeddings):
- self.embed_out = new_embeddings
-
- @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
- `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
- `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. The two additional tensors are
- only required when the model is used as a decoder in a Sequence to Sequence model.
-
- Contains pre-computed hidden-states (key and values in the self-attention blocks that can be used (see
- `past_key_values` input) to speed up sequential decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
- `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
- ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, GPTNeoXForCausalLM, GPTNeoXConfig
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
- >>> config = GPTNeoXConfig.from_pretrained("EleutherAI/gpt-neox-20b")
- >>> config.is_decoder = True
- >>> model = GPTNeoXForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b", config=config)
-
- >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> prediction_logits = outputs.logits
- ```"""
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.gpt_neox(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- lm_logits = self.embed_out(hidden_states)
-
- lm_loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # we are doing next-token prediction; shift prediction scores and input ids by one
- shift_logits = lm_logits[:, :-1, :].contiguous()
- labels = labels[:, 1:].contiguous()
- loss_fct = CrossEntropyLoss()
- lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
-
- if not return_dict:
- output = (lm_logits,) + outputs[1:]
- return ((lm_loss,) + output) if lm_loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=lm_loss,
- logits=lm_logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def prepare_inputs_for_generation(
- self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
- ):
- input_shape = input_ids.shape
- # cut decoder_input_ids if past is used
- if past_key_values is not None:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
-
- position_ids = kwargs.get("position_ids", None)
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
- if attention_mask is None:
- attention_mask = input_ids.new_ones(input_shape)
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- model_inputs = {"input_ids": input_ids}
- model_inputs.update(
- {
- "attention_mask": attention_mask,
- "past_key_values": past_key_values,
- "position_ids": position_ids,
- }
- )
-
- return model_inputs
-
- def _reorder_cache(self, past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past[:2])
- + layer_past[2:],
- )
- return reordered_past
-
-
-@add_start_docstrings(
- """
- The GPTNeoX Model transformer with a sequence classification head on top (linear layer).
-
- [`GPTNeoXForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT-1) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- GPT_NEOX_START_DOCSTRING,
-)
-class GPTNeoXForSequenceClassification(GPTNeoXPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.gpt_neox = GPTNeoXModel(config)
- self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=SequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.gpt_neox(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size, sequence_length = input_ids.shape[:2]
- else:
- batch_size, sequence_length = inputs_embeds.shape[:2]
-
- if self.config.pad_token_id is None and batch_size != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
-
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-class GPTNeoXForTokenClassification(GPTNeoXPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
-
- self.gpt_neox = GPTNeoXModel(config)
- self.dropout = nn.Dropout(config.classifier_dropout)
- self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint="LarsJonasson/pythia-410m-deduped-sft-swedish",
- output_type=TokenClassifierOutput,
- config_class=_CONFIG_FOR_DOC,
- expected_loss=0.25,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, TokenClassifierOutput]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.gpt_neox(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- hidden_states = self.dropout(hidden_states)
- logits = self.classifier(hidden_states)
-
- loss = None
- if labels is not None:
- labels = labels.to(logits.device)
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-
- if not return_dict:
- output = (logits,) + outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TokenClassifierOutput(
- loss=loss,
- logits=logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT-NeoX Model transformer with a span classification head on top for extractive question-answering tasks like
- SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- GPT_NEOX_START_DOCSTRING,
-)
-class GPTNeoXForQuestionAnswering(GPTNeoXPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.gpt_neox = GPTNeoXModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, 2)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- start_positions: Optional[torch.LongTensor] = None,
- end_positions: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.gpt_neox(
- input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1).to(start_logits.device)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1).to(end_logits.device)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/gpt_neox/tokenization_gpt_neox_fast.py b/transformers/models/gpt_neox/tokenization_gpt_neox_fast.py
deleted file mode 100644
index fd49572d7fe656d4640bb64ab55a1327a996f205..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neox/tokenization_gpt_neox_fast.py
+++ /dev/null
@@ -1,243 +0,0 @@
-# coding=utf-8
-# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for GPTNeoX."""
-import json
-from typing import List, Optional, Tuple
-
-from tokenizers import pre_tokenizers, processors
-
-from ...tokenization_utils_fast import PreTrainedTokenizerFast
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}
-
-
-class GPTNeoXTokenizerFast(PreTrainedTokenizerFast):
- """
- Construct a "fast" GPT-NeoX-20B tokenizer (backed by HuggingFace's *tokenizers* library). Based on byte-level
- Byte-Pair-Encoding.
-
- This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
- be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
- ```python
- >>> from transformers import GPTNeoXTokenizerFast
-
- >>> tokenizer = GPTNeoXTokenizerFast.from_pretrained("openai-community/gpt2")
- >>> tokenizer("Hello world")["input_ids"]
- [15496, 995]
-
- >>> tokenizer(" Hello world")["input_ids"]
- [18435, 995]
- ```
-
- You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer, but since
- the model was not pretrained this way, it might yield a decrease in performance.
-
-
-
- When used with `is_split_into_words=True`, this tokenizer needs to be instantiated with `add_prefix_space=True`.
-
-
-
- This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
- refer to this superclass for more information regarding those methods.
-
- Args:
- vocab_file (`str`):
- Path to the vocabulary file.
- merges_file (`str`):
- Path to the merges file.
- errors (`str`, *optional*, defaults to `"replace"`):
- Paradigm to follow when decoding bytes to UTF-8. See
- [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
- unk_token (`str`, *optional*, defaults to `<|endoftext|>`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- bos_token (`str`, *optional*, defaults to `<|endoftext|>`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `<|endoftext|>`):
- The end of sequence token.
- pad_token (`str`, *optional*):
- Token for padding a sequence.
- add_prefix_space (`bool`, *optional*, defaults to `False`):
- Whether or not to add an initial space to the input. This allows to treat the leading word just as any
- other word. (GPTNeoX tokenizer detect beginning of words by the preceding space).
- add_bos_token (`bool`, *optional*, defaults to `False`):
- Whether or not to add a `bos_token` at the start of sequences.
- add_eos_token (`bool`, *optional*, defaults to `False`):
- Whether or not to add an `eos_token` at the end of sequences.
- trim_offsets (`bool`, *optional*, defaults to `True`):
- Whether or not the post-processing step should trim offsets to avoid including whitespaces.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file=None,
- merges_file=None,
- tokenizer_file=None,
- unk_token="<|endoftext|>",
- bos_token="<|endoftext|>",
- eos_token="<|endoftext|>",
- pad_token=None,
- add_bos_token=False,
- add_eos_token=False,
- add_prefix_space=False,
- **kwargs,
- ):
- super().__init__(
- vocab_file,
- merges_file,
- tokenizer_file=tokenizer_file,
- unk_token=unk_token,
- bos_token=bos_token,
- eos_token=eos_token,
- pad_token=pad_token,
- add_bos_token=add_bos_token,
- add_eos_token=add_eos_token,
- add_prefix_space=add_prefix_space,
- **kwargs,
- )
-
- self._add_bos_token = add_bos_token
- self._add_eos_token = add_eos_token
- self.update_post_processor()
-
- pre_tok_state = json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__())
- if pre_tok_state.get("add_prefix_space", add_prefix_space) != add_prefix_space:
- pre_tok_class = getattr(pre_tokenizers, pre_tok_state.pop("type"))
- pre_tok_state["add_prefix_space"] = add_prefix_space
- self.backend_tokenizer.pre_tokenizer = pre_tok_class(**pre_tok_state)
-
- self.add_prefix_space = add_prefix_space
-
- @property
- def add_eos_token(self):
- return self._add_eos_token
-
- @property
- def add_bos_token(self):
- return self._add_bos_token
-
- @add_eos_token.setter
- def add_eos_token(self, value):
- self._add_eos_token = value
- self.update_post_processor()
-
- @add_bos_token.setter
- def add_bos_token(self, value):
- self._add_bos_token = value
- self.update_post_processor()
-
- # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.update_post_processor
- def update_post_processor(self):
- """
- Updates the underlying post processor with the current `bos_token` and `eos_token`.
- """
- bos = self.bos_token
- bos_token_id = self.bos_token_id
- if bos is None and self.add_bos_token:
- raise ValueError("add_bos_token = True but bos_token = None")
-
- eos = self.eos_token
- eos_token_id = self.eos_token_id
- if eos is None and self.add_eos_token:
- raise ValueError("add_eos_token = True but eos_token = None")
-
- single = f"{(bos+':0 ') if self.add_bos_token else ''}$A:0{(' '+eos+':0') if self.add_eos_token else ''}"
- pair = f"{single}{(' '+bos+':1') if self.add_bos_token else ''} $B:1{(' '+eos+':1') if self.add_eos_token else ''}"
-
- special_tokens = []
- if self.add_bos_token:
- special_tokens.append((bos, bos_token_id))
- if self.add_eos_token:
- special_tokens.append((eos, eos_token_id))
- self._tokenizer.post_processor = processors.TemplateProcessing(
- single=single, pair=pair, special_tokens=special_tokens
- )
-
- # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.get_special_tokens_mask
- def get_special_tokens_mask(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
- ) -> List[int]:
- """
- Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
- special tokens using the tokenizer `prepare_for_model` method.
-
- Args:
- token_ids_0 (`List[int]`):
- List of IDs.
- token_ids_1 (`List[int]`, *optional*):
- Optional second list of IDs for sequence pairs.
- already_has_special_tokens (`bool`, *optional*, defaults to `False`):
- Whether or not the token list is already formatted with special tokens for the model.
-
- Returns:
- `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
- """
- if already_has_special_tokens:
- return super().get_special_tokens_mask(
- token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
- )
-
- bos_token_id = [1] if self.add_bos_token else []
- eos_token_id = [1] if self.add_eos_token else []
-
- if token_ids_1 is None:
- return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
- return (
- bos_token_id
- + ([0] * len(token_ids_0))
- + eos_token_id
- + bos_token_id
- + ([0] * len(token_ids_1))
- + eos_token_id
- )
-
- # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.build_inputs_with_special_tokens
- def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
- bos_token_id = [self.bos_token_id] if self.add_bos_token else []
- eos_token_id = [self.eos_token_id] if self.add_eos_token else []
-
- output = bos_token_id + token_ids_0 + eos_token_id
-
- if token_ids_1 is not None:
- output = output + bos_token_id + token_ids_1 + eos_token_id
-
- return output
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- files = self._tokenizer.model.save(save_directory, name=filename_prefix)
- return tuple(files)
-
- @property
- # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.default_chat_template
- def default_chat_template(self):
- """
- A simple chat template that ignores role information and just concatenates messages with EOS tokens.
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- return "{% for message in messages %}" "{{ message.content }}{{ eos_token }}" "{% endfor %}"
diff --git a/transformers/models/gpt_neox_japanese/__init__.py b/transformers/models/gpt_neox_japanese/__init__.py
deleted file mode 100644
index bf04db7676c8b6d3871a8d56b330f19ddea7c6a7..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neox_japanese/__init__.py
+++ /dev/null
@@ -1,62 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...file_utils import _LazyModule, is_torch_available
-from ...utils import OptionalDependencyNotAvailable
-
-
-_import_structure = {
- "configuration_gpt_neox_japanese": ["GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTNeoXJapaneseConfig"],
- "tokenization_gpt_neox_japanese": ["GPTNeoXJapaneseTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_gpt_neox_japanese"] = [
- "GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GPTNeoXJapaneseForCausalLM",
- "GPTNeoXJapaneseLayer",
- "GPTNeoXJapaneseModel",
- "GPTNeoXJapanesePreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_gpt_neox_japanese import GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXJapaneseConfig
- from .tokenization_gpt_neox_japanese import GPTNeoXJapaneseTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_gpt_neox_japanese import (
- GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST,
- GPTNeoXJapaneseForCausalLM,
- GPTNeoXJapaneseLayer,
- GPTNeoXJapaneseModel,
- GPTNeoXJapanesePreTrainedModel,
- )
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gpt_neox_japanese/__pycache__/__init__.cpython-310.pyc b/transformers/models/gpt_neox_japanese/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 731b36a7fee283041584ad33a208f0714fe4b8d0..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neox_japanese/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neox_japanese/__pycache__/configuration_gpt_neox_japanese.cpython-310.pyc b/transformers/models/gpt_neox_japanese/__pycache__/configuration_gpt_neox_japanese.cpython-310.pyc
deleted file mode 100644
index 0500bf44dd1ae901c013be110e62b58826061db3..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neox_japanese/__pycache__/configuration_gpt_neox_japanese.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neox_japanese/__pycache__/modeling_gpt_neox_japanese.cpython-310.pyc b/transformers/models/gpt_neox_japanese/__pycache__/modeling_gpt_neox_japanese.cpython-310.pyc
deleted file mode 100644
index a23922e3fb39bea011a0ff4c4ea8ad95610db50b..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neox_japanese/__pycache__/modeling_gpt_neox_japanese.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neox_japanese/__pycache__/tokenization_gpt_neox_japanese.cpython-310.pyc b/transformers/models/gpt_neox_japanese/__pycache__/tokenization_gpt_neox_japanese.cpython-310.pyc
deleted file mode 100644
index 91653a7e0892b7178bc63d34cd9245797d31b159..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_neox_japanese/__pycache__/tokenization_gpt_neox_japanese.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_neox_japanese/configuration_gpt_neox_japanese.py b/transformers/models/gpt_neox_japanese/configuration_gpt_neox_japanese.py
deleted file mode 100644
index 8ee73257b64c7c1a93f975af7b61f707decb9cd9..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neox_japanese/configuration_gpt_neox_japanese.py
+++ /dev/null
@@ -1,120 +0,0 @@
-# coding=utf-8
-# Copyright 2022 ABEJA, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPTNeoX Japanese model configuration"""
-
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GPTNeoXJapaneseConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GPTNeoXModelJapanese`]. It is used to instantiate
- a GPTNeoX model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the GPTNeoXJapanese
- [abeja/gpt-neox-japanese-2.7b](https://huggingface.co/abeja/gpt-neox-japanese-2.7b) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information. Default configs is set as 2.7B model
-
- Args:
- vocab_size (`int`, *optional*, defaults to 32000):
- Vocabulary size of the GPTNeoXJapanese model. Defines the number of different tokens that can be
- represented by the `inputs_ids` passed when calling [`GPTNeoXJapanese`].
- hidden_size (`int`, *optional*, defaults to 2560):
- Dimension of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, *optional*, defaults to 32):
- Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, *optional*, defaults to 32):
- Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_multiple_size (`int`, *optional*, defaults to 4):
- Dimension of the "intermediate" layer in the Transformer encoder is calculated by hidden_size *
- intermediate_multiple_size.
- hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
- The non-linear activation function (function or string) in the encoder and pooler.
- rotary_pct (`float`, *optional*, defaults to 1.00):
- percentage of hidden dimensions to allocate to rotary embeddings
- rotary_emb_base (`int`, *optional*, defaults to 10000)
- base for computing rotary embeddings frequency
- max_position_embeddings (`int`, *optional*, defaults to 2048):
- The maximum sequence length that this model might ever be used with.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, *optional*, defaults to 1e-5):
- The epsilon used by the layer normalization layers.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models). Only
- relevant if `config.is_decoder=True`.
- attention_dropout (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention.
- hidden_dropout (`float`, *optional*, defaults to 0.0):
- The dropout ratio for the hidden layer.
- Example:
-
- ```python
- >>> from transformers import GPTNeoXJapaneseConfig, GPTNeoXJapaneseModel
-
- >>> # Initializing a GPTNeoXJapanese gpt-neox-japanese-2.7b style configuration
- >>> configuration = GPTNeoXJapaneseConfig()
-
- >>> # Initializing a model (with random weights) from the gpt-neox-japanese-2.7b style configuration
- >>> model = GPTNeoXJapaneseModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "gpt_neox_japanese"
-
- def __init__(
- self,
- vocab_size=32000,
- hidden_size=2560,
- num_hidden_layers=32,
- num_attention_heads=32,
- intermediate_multiple_size=4,
- hidden_act="gelu",
- rotary_pct=1.00,
- rotary_emb_base=10000,
- max_position_embeddings=2048,
- initializer_range=0.02,
- layer_norm_eps=1e-5,
- use_cache=True,
- bos_token_id=31996,
- eos_token_id=31999,
- attention_dropout=0.1,
- hidden_dropout=0.0,
- **kwargs,
- ):
- super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.hidden_size = hidden_size
- self.num_hidden_layers = num_hidden_layers
- self.num_attention_heads = num_attention_heads
- self.intermediate_multiple_size = intermediate_multiple_size
- self.hidden_act = hidden_act
- self.rotary_pct = rotary_pct
- self.rotary_emb_base = rotary_emb_base
- self.initializer_range = initializer_range
- self.layer_norm_eps = layer_norm_eps
- self.use_cache = use_cache
- self.attention_dropout = attention_dropout
- self.hidden_dropout = hidden_dropout
diff --git a/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py b/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py
deleted file mode 100644
index 9fdff2c838700639bb1055a584dbe0a06dadab30..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py
+++ /dev/null
@@ -1,729 +0,0 @@
-# coding=utf-8
-# Copyright 2022 ABEJA, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch GPTNeoX model."""
-
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import Tensor, nn
-from torch.nn import CrossEntropyLoss
-
-from ...activations import ACT2FN
-from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
-from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
-from ...modeling_utils import PreTrainedModel
-from ...utils import logging
-from .configuration_gpt_neox_japanese import GPTNeoXJapaneseConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "abeja/gpt-neox-japanese-2.7b"
-_CONFIG_FOR_DOC = "GPTNeoXJapaneseConfig"
-
-
-from ..deprecated._archive_maps import GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-class GPTNeoXJapanesePreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTNeoXJapaneseConfig
- base_model_prefix = "gpt_neox_japanese"
- _no_split_modules = ["GPTNeoXJapaneseLayer"]
- _skip_keys_device_placement = "past_key_values"
-
- def _init_weights(self, module):
- """Initialize the weights"""
- if isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-class GPTNeoXJapaneseAttention(nn.Module):
- def __init__(self, config, use_bias=False):
- super().__init__()
- self.num_attention_heads = config.num_attention_heads
- self.hidden_size = config.hidden_size
- self.head_size = self.hidden_size // self.num_attention_heads
-
- self.rotary_ndims = int(self.head_size * config.rotary_pct)
- self.rotary_emb = RotaryEmbedding(
- self.rotary_ndims, config.max_position_embeddings, base=config.rotary_emb_base
- )
- self.max_positions = config.max_position_embeddings
- self.attention_dropout = nn.Dropout(config.attention_dropout)
- self.norm_factor = torch.sqrt(torch.tensor(self.head_size, dtype=torch.float32)).to(torch.get_default_dtype())
-
- self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
- self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
- # Activate bias if the last layer
- self.use_bias = use_bias
- self.dense_bias = nn.Parameter(torch.zeros(config.hidden_size)) if use_bias else None
-
- def forward(
- self,
- hidden_states,
- attention_mask,
- head_mask=None,
- layer_past=None,
- use_cache=False,
- output_attentions=False,
- ):
- has_layer_past = layer_past is not None and layer_past[0].numel() > 0
-
- # Compute QKV
- # Attention heads [batch, seq_len, hidden_size]
- # --> [batch, seq_len, (np * 3 * head_size)]
- qkv = self.query_key_value(hidden_states)
-
- # [batch, seq_len, (num_heads * 3 * head_size)]
- # --> [batch, seq_len, num_heads, 3 * head_size]
- new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size)
- qkv = qkv.view(*new_qkv_shape)
-
- # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
- query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
- key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3)
- value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3)
-
- # Compute rotary embeddings on rotary_ndims
- query_rot = query[..., : self.rotary_ndims]
- query_pass = query[..., self.rotary_ndims :]
- key_rot = key[..., : self.rotary_ndims]
- key_pass = key[..., self.rotary_ndims :]
-
- # Compute token offset for rotary embeddings (when decoding)
- seq_len = key.shape[-2]
- offset = 0
- if has_layer_past:
- offset = layer_past[0].shape[-2]
- seq_len += offset
- cos, sin = self.rotary_emb(value, seq_len=seq_len)
- query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, offset=offset)
- query = torch.cat((query, query_pass), dim=-1)
- key = torch.cat((key, key_pass), dim=-1)
-
- # Cache QKV values
- if has_layer_past:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
- present = (key, value) if use_cache else None
-
- # Compute attention
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
- # Reshape outputs
- attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size)
- attn_output = self.dense(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs, self.dense_bias
-
- @classmethod
- def _split_heads(cls, tensor, num_attention_heads, attn_head_size):
- """
- Splits hidden dim into attn_head_size and num_attention_heads
- """
- # tensor: [bs, seq_len, hidden_size]
- new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
- # -> [bs, seq_len, num_attention_heads, attn_head_size]
- tensor = tensor.view(new_shape)
- # -> [bs, num_attention_heads, seq_len, attn_head_size]
- tensor = tensor.permute(0, 2, 1, 3)
- return tensor
-
- @classmethod
- def _merge_heads(cls, tensor, num_attention_heads, attn_head_size):
- """
- Merges attn_head_size dim and num_attn_heads dim into hidden dim
- """
- # tensor [bs, num_attention_heads, seq_len, attn_head_size]
- tensor = tensor.permute(0, 2, 1, 3).contiguous()
- # -> [bs, seq_len, num_attention_heads, attn_head_size]
- tensor = tensor.view(tensor.size(0), tensor.size(1), num_attention_heads * attn_head_size)
- # -> [bs, seq_len, hidden_size]
- return tensor
-
- def _create_causal_mask(self, key_length, query_length):
- causal_mask = torch.tril(
- torch.ones((self.max_positions, self.max_positions), dtype=torch.bool).view(
- 1, 1, self.max_positions, self.max_positions
- )
- )
- return causal_mask[:, :, key_length - query_length : key_length, :key_length]
-
- def _attn(self, query, key, value, attention_mask=None, head_mask=None):
- # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size]
- # compute causal mask from causal mask buffer
- batch_size, num_attention_heads, query_length, attn_head_size = query.size()
- key_length = key.size(-2)
-
- causal_mask = self._create_causal_mask(key_length, query_length)
-
- query = query.view(batch_size * num_attention_heads, query_length, attn_head_size)
- key = key.view(batch_size * num_attention_heads, key_length, attn_head_size)
- attn_scores = torch.zeros(
- batch_size * num_attention_heads,
- query_length,
- key_length,
- dtype=query.dtype,
- device=key.device,
- )
- attn_scores = torch.baddbmm(
- attn_scores,
- query,
- key.transpose(1, 2),
- beta=1.0,
- alpha=(torch.tensor(1.0, dtype=self.norm_factor.dtype, device=self.norm_factor.device) / self.norm_factor),
- )
- attn_scores = attn_scores.view(batch_size, num_attention_heads, query_length, key_length)
-
- mask_value = torch.finfo(attn_scores.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.tensor(mask_value, dtype=attn_scores.dtype).to(attn_scores.device)
- causal_mask = causal_mask.to(attn_scores.device)
- attn_scores = torch.where(causal_mask, attn_scores, mask_value)
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_scores = attn_scores + attention_mask
-
- attn_weights = nn.functional.softmax(attn_scores, dim=-1)
- attn_weights = self.attention_dropout(attn_weights)
- attn_weights = attn_weights.to(value.dtype)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
- return attn_output, attn_weights
-
-
-# Copied from transformers.models.gpt_neox.modeling_gpt_neox.GPTNeoXRotaryEmbedding with GPTNeoXRotaryEmbedding->RotaryEmbedding
-class RotaryEmbedding(nn.Module):
- # Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding.__init__
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
- super().__init__()
-
- self.dim = dim
- self.max_position_embeddings = max_position_embeddings
- self.base = base
- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
- self.register_buffer("inv_freq", inv_freq, persistent=False)
-
- # Build here to make `torch.jit.trace` work.
- self._set_cos_sin_cache(
- seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
- )
-
- def _set_cos_sin_cache(self, seq_len, device, dtype):
- self.max_seq_len_cached = seq_len
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
-
- freqs = torch.outer(t, self.inv_freq)
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
- emb = torch.cat((freqs, freqs), dim=-1)
- self.register_buffer("cos_cached", emb.cos(), persistent=False)
- self.register_buffer("sin_cached", emb.sin(), persistent=False)
-
- def forward(self, x, seq_len=None):
- # x: [bs, num_attention_heads, seq_len, head_size]
- if seq_len > self.max_seq_len_cached:
- self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-
- return (
- self.cos_cached[:seq_len],
- self.sin_cached[:seq_len],
- )
-
-
-def rotate_half(x):
- """Rotates half the hidden dims of the input."""
- x1 = x[..., : x.shape[-1] // 2]
- x2 = x[..., x.shape[-1] // 2 :]
- return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, offset: int = 0):
- cos = cos[..., offset : q.shape[-2] + offset, :]
- sin = sin[..., offset : q.shape[-2] + offset, :]
- q_embed = (q * cos) + (rotate_half(q) * sin)
- k_embed = (k * cos) + (rotate_half(k) * sin)
- return q_embed, k_embed
-
-
-def bias_dropout_add(x: Tensor, bias: Tensor, residual: Optional[Tensor], prob: float, training: bool) -> Tensor:
- """add bias to x, apply dropout and residual connection
-
- Args:
- x (Tensor): main path of output
- bias (Tensor): None or attn_bias of the last attention layer
- residual (Optional[Tensor]): residual value
- prob (float): dropout probability
- training (bool): whether in training mode or not
-
- Returns:
- Tensor: dropout(x + bias) + residual
- """
- if bias is not None:
- x = x + bias
- out = torch.nn.functional.dropout(x, p=prob, training=training)
- if residual is not None:
- out = residual + out
- return out
-
-
-class GPTNeoXJapaneseMLP(nn.Module):
- def __init__(self, config):
- super().__init__()
- intermediate_size = int(config.hidden_size * config.intermediate_multiple_size)
- self.dense_h_to_4h = nn.Linear(config.hidden_size, intermediate_size, bias=False)
- # Project back to h.
- self.dense_4h_to_h = nn.Linear(intermediate_size, config.hidden_size, bias=False)
- self.act = ACT2FN[config.hidden_act]
-
- def forward(self, hidden_states):
- intermediate = self.dense_h_to_4h(hidden_states)
- intermediate = self.act(intermediate)
- output = self.dense_4h_to_h(intermediate)
- return output
-
-
-class GPTNeoXJapaneseLayer(nn.Module):
- def __init__(self, config, layer_number):
- super().__init__()
- self.layer_number = layer_number
- self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
- # activate bias only last layer
- self.attention = GPTNeoXJapaneseAttention(config=config, use_bias=layer_number == config.num_hidden_layers - 1)
- self.mlp = GPTNeoXJapaneseMLP(config)
- self.hidden_dropout = config.hidden_dropout
-
- def forward(
- self,
- hidden_states,
- attention_mask=None,
- head_mask=None,
- use_cache=False,
- layer_past=None,
- output_attentions=False,
- ):
- residual = hidden_states
- ln_out = self.input_layernorm(hidden_states)
- attention_layer_outputs, attn_bias = self.attention(
- ln_out,
- attention_mask=attention_mask,
- layer_past=layer_past,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attention_layer_outputs[0] # output_attn: a, present, (attentions)
- outputs = attention_layer_outputs[1:]
-
- # attn_output = (atten_output + bias) + residual
- attn_output = bias_dropout_add(
- attn_output,
- bias=attn_bias.expand_as(residual) if attn_bias is not None else attn_bias,
- residual=residual,
- prob=self.hidden_dropout,
- training=self.training,
- )
- mlp_output = self.mlp(self.post_attention_layernorm(attn_output))
-
- # attn_output = (mlp_output + mlp_bias) + atten_output
- attn_output = bias_dropout_add(
- mlp_output, bias=None, residual=attn_output, prob=self.hidden_dropout, training=self.training
- )
-
- if use_cache:
- outputs = (attn_output,) + outputs
- else:
- outputs = (attn_output,) + outputs[1:]
-
- return outputs # hidden_states, present, (attentions)
-
-
-GPT_NEOX_JAPANESE_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`~GPTNeoXJapaneseConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPT_NEOX_JAPANESE_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`].
-
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-@add_start_docstrings(
- "The bare GPTNeoXJapanese Model transformer outputting raw hidden-states without any specific head on top.",
- GPT_NEOX_JAPANESE_START_DOCSTRING,
-)
-class GPTNeoXJapaneseModel(GPTNeoXJapanesePreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.config = config
-
- self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
- self.layers = nn.ModuleList(
- [GPTNeoXJapaneseLayer(config=config, layer_number=i) for i in range(config.num_hidden_layers)]
- )
- self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embed_in
-
- def set_input_embeddings(self, value):
- self.embed_in = value
-
- @add_start_docstrings_to_model_forward(GPT_NEOX_JAPANESE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=BaseModelOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPast]:
- r"""
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, GPTNeoXJapaneseModel
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("abeja/gpt-neox-japanese-2.7b")
- >>> model = GPTNeoXJapaneseModel.from_pretrained("abeja/gpt-neox-japanese-2.7b")
-
- >>> inputs = tokenizer("日本語のGPT-neoxがHugging Faceで使えます😀", return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> last_hidden_states = outputs.last_hidden_state
- ```
- """
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- use_cache = use_cache if use_cache is not None else self.config.use_cache
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- batch_size, seq_length = input_shape
-
- if past_key_values is None:
- past_key_values = tuple([None] * self.config.num_hidden_layers)
-
- # Attention mask.
- if attention_mask is not None:
- if not batch_size > 0:
- raise ValueError("batch_size has to be defined and > 0")
- attention_mask = attention_mask.view(batch_size, -1)
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask = attention_mask[:, None, None, :]
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and -10000.0 for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
- attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
-
- if inputs_embeds is None:
- inputs_embeds = self.embed_in(input_ids)
-
- hidden_states = inputs_embeds
-
- presents = () if use_cache else None
- all_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
- for i, (layer, layer_past) in enumerate(zip(self.layers, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
- outputs = layer(
- hidden_states,
- attention_mask=attention_mask,
- head_mask=head_mask[i],
- layer_past=layer_past,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
- if output_attentions:
- all_attentions = all_attentions + (outputs[2 if use_cache else 1],)
-
- hidden_states = self.final_layer_norm(hidden_states)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)
-
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- )
-
-
-@add_start_docstrings(
- """GPTNeoXJapanese Model with a `language modeling` head on top for Classifier Model fine-tuning.""",
- GPT_NEOX_JAPANESE_START_DOCSTRING,
-)
-class GPTNeoXJapaneseForCausalLM(GPTNeoXJapanesePreTrainedModel):
- _tied_weights_keys = ["embed_out.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.config = config
-
- self.gpt_neox_japanese = GPTNeoXJapaneseModel(config)
- self.embed_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
- # Initialize weights and apply final processing
- self.post_init()
-
- def get_output_embeddings(self):
- return self.embed_out
-
- def set_output_embeddings(self, new_embeddings):
- self.embed_out = new_embeddings
-
- @add_start_docstrings_to_model_forward(GPT_NEOX_JAPANESE_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
- Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
- `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
- `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. The two additional tensors are
- only required when the model is used as a decoder in a Sequence to Sequence model.
-
- Contains pre-computed hidden-states (key and values in the self-attention blocks that can be used (see
- `past_key_values` input) to speed up sequential decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
- `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
- ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, GPTNeoXJapaneseForCausalLM, GPTNeoXJapaneseConfig
- >>> import torch
-
- >>> tokenizer = AutoTokenizer.from_pretrained("abeja/gpt-neox-japanese-2.7b")
- >>> config = GPTNeoXJapaneseConfig.from_pretrained("abeja/gpt-neox-japanese-2.7b")
- >>> config.is_decoder = True
- >>> model = GPTNeoXJapaneseForCausalLM.from_pretrained("abeja/gpt-neox-japanese-2.7b", config=config)
-
- >>> inputs = tokenizer("日本語のGPT-neoxがHugging Faceで使えます😀", return_tensors="pt")
- >>> outputs = model(**inputs)
-
- >>> prediction_logits = outputs.logits
- ```
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.gpt_neox_japanese(
- input_ids,
- attention_mask=attention_mask,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- past_key_values=past_key_values,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- lm_logits = self.embed_out(hidden_states)
-
- lm_loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
-
- # we are doing next-token prediction; shift prediction scores and input ids by one
- shift_logits = lm_logits[:, :-1, :].contiguous()
- labels = labels[:, 1:].contiguous()
- loss_fct = CrossEntropyLoss()
- lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
-
- if not return_dict:
- output = (lm_logits,) + outputs[1:]
- return ((lm_loss,) + output) if lm_loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=lm_loss,
- logits=lm_logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs):
- input_shape = input_ids.shape
-
- # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
- if attention_mask is None:
- attention_mask = input_ids.new_ones(input_shape)
-
- # cut decoder_input_ids if past is used
- if past_key_values and past_key_values[0] is not None:
- input_ids = input_ids[:, -1:]
-
- return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values}
-
- def _reorder_cache(self, past_key_values, beam_idx):
- reordered_past = ()
- for layer_past in past_key_values:
- reordered_past += (
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past[:2])
- + layer_past[2:],
- )
- return reordered_past
diff --git a/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py b/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py
deleted file mode 100644
index fd0fe796dcab02ee6a1b1972a456fb24aa367495..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_neox_japanese/tokenization_gpt_neox_japanese.py
+++ /dev/null
@@ -1,368 +0,0 @@
-# coding=utf-8
-# Copyright 2022 ABEJA, Inc. and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for GPTNeoXJapanese."""
-import collections
-import json
-import os
-import re
-from typing import Optional, Tuple
-
-import numpy as np
-
-from ...tokenization_utils_fast import PreTrainedTokenizer
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "emoji_file": "emoji.json"}
-
-
-def load_vocab_and_emoji(vocab_file, emoji_file):
- """Loads a vocabulary file and emoji file into a dictionary."""
- with open(emoji_file, "r", encoding="utf-8") as f:
- emoji = json.loads(f.read())
-
- vocab = collections.OrderedDict()
- raw_vocab = collections.OrderedDict()
- ids_to_tokens = collections.OrderedDict()
- with open(vocab_file, "r", encoding="utf-8") as f:
- token = f.readlines()
- token = [[t.rstrip("\n")] if (t == "," or "," not in t) else t.rstrip("\n").split(",") for t in token]
- for idx, b in enumerate(token):
- ids_to_tokens[idx] = b
- raw_vocab[",".join(b)] = idx
- for wd in b:
- vocab[wd] = idx
-
- return vocab, raw_vocab, ids_to_tokens, emoji
-
-
-class GPTNeoXJapaneseTokenizer(PreTrainedTokenizer):
- """
- This tokenizer inherits from [`PreTrainedTokenizer`] and is based on Japanese special Sub-Word-Encoding that is
- used in this repository (https://github.com/tanreinama/Japanese-BPEEncoder_V2). Check the repository for details.
- Japanese has a relatively large vocabulary and there is no separation between words. Furthermore, the language is a
- combination of hiragana, katakana, and kanji, and variants such as "1" and "①" are often used. In order to cope
- with these, this tokenizer has the following features
- - Subword-by-subword segmentation, which is intermediate between byte strings and morphological analysis.
- - BPEs are created for each Kanji, Hiragana, and Katakana character, and there are no BPEs that cross character
- types, such as Kanji + Hiragana or Hiragana + Katakana.
- - All-byte encoding that does not require .
- - Independent of UTF codes such as 2-byte and 3-byte characters
- - Conversion of heterographs to the same token_id
- - Emoji and Emoticon are grouped into 12 types as special tags.
-
- Example:
-
- ```python
- >>> from transformers import GPTNeoXJapaneseTokenizer
-
- >>> tokenizer = GPTNeoXJapaneseTokenizer.from_pretrained("abeja/gpt-neox-japanese-2.7b")
- >>> # You can confirm both 慶応 and 慶應 are encoded to 17749
- >>> tokenizer("吾輩は猫である🐯。実は慶応(慶應)大学出身")["input_ids"]
- [30014, 26883, 26638, 27228, 25, 26650, 31732, 31679, 27809, 26638, 17749, 31592, 17749, 31593, 321, 1281]
-
- >>> # Both 慶応 and 慶應 are decoded to 慶応
- >>> tokenizer.decode(tokenizer("吾輩は猫である🐯。実は慶応(慶應)大学出身")["input_ids"])
- '吾輩は猫である🐯。実は慶応(慶応)大学出身'
- ```
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- emoji_file (`str`):
- File containing the emoji.
- unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead.
- pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The token used for padding
- bos_token (`str`, *optional*, defaults to `"<|startoftext|>"`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The end of sequence token.
- do_clean_text (`bool`, *optional*, defaults to `False`):
- Whether or not to clean text for URL, EMAIL, TEL, Japanese DATE and Japanese PRICE.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- emoji_file,
- unk_token="<|endoftext|>",
- pad_token="<|endoftext|>",
- bos_token="<|startoftext|>",
- eos_token="<|endoftext|>",
- do_clean_text=False,
- **kwargs,
- ):
- if not os.path.isfile(vocab_file):
- raise ValueError(
- f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
- " model use `tokenizer = GPTNeoXJapaneseokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- if not os.path.isfile(emoji_file):
- raise ValueError(
- f"Can't find a emoji file at path '{emoji_file}'. To load the emoji information from a Google"
- " pretrained model use `tokenizer = GPTNeoXJapaneseokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- self.do_clean_text = do_clean_text
- self.vocab, self.raw_vocab, self.ids_to_tokens, self.emoji = load_vocab_and_emoji(vocab_file, emoji_file)
- self.subword_tokenizer = SubWordJapaneseTokenizer(
- vocab=self.vocab, ids_to_tokens=self.ids_to_tokens, emoji=self.emoji
- )
- super().__init__(
- unk_token=unk_token,
- pad_token=pad_token,
- bos_token=bos_token,
- eos_token=eos_token,
- do_clean_text=do_clean_text,
- **kwargs,
- )
-
- @property
- def vocab_size(self):
- # self.vocab contains support for character fluctuation unique to Japanese, and has a large number of vocab
- return len(self.raw_vocab)
-
- def get_vocab(self):
- return dict(self.raw_vocab, **self.added_tokens_encoder)
-
- def _tokenize(self, text):
- return self.subword_tokenizer.tokenize(text, clean=self.do_clean_text)
-
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.vocab.get(token, self.vocab.get(self.unk_token))
-
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.subword_tokenizer.convert_id_to_token(index)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- out_string = "".join(tokens).strip()
- return out_string
-
- @property
- def default_chat_template(self):
- """
- A simple chat template that just adds BOS/EOS tokens around messages while discarding role information.
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- return (
- "{% for message in messages %}"
- "{{ bos_token + eos_token + message.content + eos_token }}"
- "{% endfor %}"
- "{% if add_generation_prompt %} {{ bos_token + eos_token }} {% endif %}"
- )
-
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- index = 0
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- emoji_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["emoji_file"]
- )
- else:
- vocab_file = (
- (filename_prefix + "-" if filename_prefix else "") + save_directory + VOCAB_FILES_NAMES["vocab_file"]
- )
- emoji_file = (
- (filename_prefix + "-" if filename_prefix else "") + save_directory + VOCAB_FILES_NAMES["emoji_file"]
- )
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token_index, token in self.ids_to_tokens.items():
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(",".join(token) + "\n")
- index += 1
- with open(emoji_file, "w", encoding="utf-8") as writer:
- json.dump(self.emoji, writer)
- return vocab_file, emoji_file
-
-
-class SubWordJapaneseTokenizer(object):
- """
- https://github.com/tanreinama/Japanese-BPEEncoder_V2 This tokenizer class is under MIT Lisence according to the
- original repository.
-
- MIT License
-
- Copyright (c) 2020 tanreinama
-
- Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
- documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
- rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
- permit persons to whom the Software is furnished to do so, subject to the following conditions:
-
- The above copyright notice and this permission notice shall be included in all copies or substantial portions of
- the Software.
-
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
- THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
- TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- SOFTWARE.
- """
-
- def __init__(self, vocab, ids_to_tokens, emoji):
- self.vocab = vocab # same as swe
- self.ids_to_tokens = ids_to_tokens # same as bpe
- self.emoji = emoji
- self.maxlen = np.max([len(w) for w in self.vocab.keys()])
- self.content_repatter1 = re.compile(r"(https?|ftp)(:\/\/[-_\.!~*\'()a-zA-Z0-9;\/?:\@&=\+$,%#]+)")
- self.content_repatter2 = re.compile(r"[A-Za-z0-9\._+]*@[\-_0-9A-Za-z]+(\.[A-Za-z]+)*")
- self.content_repatter3 = re.compile(r"[\(]{0,1}[0-9]{2,4}[\)\-\(]{0,1}[0-9]{2,4}[\)\-]{0,1}[0-9]{3,4}")
- self.content_repatter4 = re.compile(
- r"([12]\d{3}[/\-年])*(0?[1-9]|1[0-2])[/\-月]((0?[1-9]|[12][0-9]|3[01])日?)*(\d{1,2}|:|\d{1,2}時|\d{1,2}分|\(日\)|\(月\)|\(火\)|\(水\)|\(木\)|\(金\)|\(土\)|㈰|㈪|㈫|㈬|㈭|㈮|㈯)*"
- )
- self.content_repatter5 = re.compile(
- r"(明治|大正|昭和|平成|令和|㍾|㍽|㍼|㍻|\u32ff)\d{1,2}年(0?[1-9]|1[0-2])月(0?[1-9]|[12][0-9]|3[01])日(\d{1,2}|:|\d{1,2}時|\d{1,2}分|\(日\)|\(月\)|\(火\)|\(水\)|\(木\)|\(金\)|\(土\)|㈰|㈪|㈫|㈬|㈭|㈮|㈯)*"
- )
- self.content_repatter6 = re.compile(
- r"((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*億)*((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*万)*((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*千)*(0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*(千円|万円|千万円|円|千ドル|万ドル|千万ドル|ドル|千ユーロ|万ユーロ|千万ユーロ|ユーロ)+(\(税込\)|\(税抜\)|\+tax)*"
- )
- keisen = "─━│┃┄┅┆┇┈┉┊┋┌┍┎┏┐┑┒┓└┕┖┗┘┙┚┛├┝┞┟┠┡┢┣┤┥┦┧┨┩┪┫┬┭┮┯┰┱┲┳┴┵┶┷┸┹┺┻┼┽┾┿╀╁╂╃╄╅╆╇╈╉╊╋╌╍╎╏═║╒╓╔╕╖╗╘╙╚╛╜╝╞╟╠╡╢╣╤╥╦╧╨╩╪╫╬╭╮╯╰╱╲╳╴╵╶╷╸╹╺╻╼╽╾╿"
- blocks = "▀▁▂▃▄▅▆▇█▉▊▋▌▍▎▏▐░▒▓▔▕▖▗▘▙▚▛▜▝▞▟"
- self.content_trans1 = str.maketrans({k: "" for k in keisen + blocks})
-
- def __len__(self):
- return len(self.ids_to_tokens)
-
- def clean_text(self, content):
- content = self.content_repatter1.sub("", content)
- content = self.content_repatter2.sub("", content)
- content = self.content_repatter3.sub("", content)
- content = self.content_repatter4.sub("", content)
- content = self.content_repatter5.sub("", content)
- content = self.content_repatter6.sub("", content)
- content = content.translate(self.content_trans1)
- while "" in content:
- content = content.replace("", "")
- return content
-
- def tokenize(self, text, clean=False):
- text = text.replace(" ", "")
- text = text.replace(" ", "")
- text = text.replace("\r\n", "
")
- text = text.replace("\n", "
")
- text = text.replace("\r", "
")
- text = text.replace("\t", "")
- text = text.replace("—", "ー")
- text = text.replace("−", "ー")
- for k, v in self.emoji["emoji"].items():
- if k in text:
- text = text.replace(k, v)
- if clean:
- text = self.clean_text(text)
-
- def check_simbol(x):
- e = x.encode()
- if len(x) == 1 and len(e) == 2:
- c = (int(e[0]) << 8) + int(e[1])
- if (
- (c >= 0xC2A1 and c <= 0xC2BF)
- or (c >= 0xC780 and c <= 0xC783)
- or (c >= 0xCAB9 and c <= 0xCBBF)
- or (c >= 0xCC80 and c <= 0xCDA2)
- ):
- return True
- return False
-
- def checku2e(x):
- e = x.encode()
- if len(x) == 1 and len(e) == 3:
- c = (int(e[0]) << 16) + (int(e[1]) << 8) + int(e[2])
- if c >= 0xE28080 and c <= 0xE2B07F:
- return True
- return False
-
- pos = 0
- result = []
- while pos < len(text):
- end = min(len(text), pos + self.maxlen + 1) if text[pos] == "<" else pos + 3
- candidates = [] # (token_id, token, pos)
- for e in range(end, pos, -1):
- wd = text[pos:e]
- if wd in self.vocab:
- if wd[0] == "<" and len(wd) > 2:
- candidates = [(self.vocab[wd], wd, e)]
- break
- else:
- candidates.append((self.vocab[wd], wd, e))
- if len(candidates) > 0:
- # the smallest token_id is adopted
- _, wd, e = sorted(candidates, key=lambda x: x[0])[0]
- result.append(wd)
- pos = e
- else:
- end = pos + 1
- wd = text[pos:end]
- if check_simbol(wd):
- result.append("")
- elif checku2e(wd):
- result.append("")
- else:
- for i in wd.encode("utf-8"):
- result.append("<|byte%d|>" % i)
- pos = end
- return result
-
- def convert_id_to_token(self, index, breakline="\n"):
- words = []
- byte_tokens = []
- word = self.ids_to_tokens[index][0]
- if word[:6] == "<|byte" and word[-2:] == "|>":
- byte_tokens.append(int(word[6:-2]))
- else:
- if len(byte_tokens) > 0:
- words.append(bytearray(byte_tokens).decode("utf-8", errors="replace"))
- byte_tokens = []
- if word[:7] == "<|emoji" and word[-2:] == "|>":
- words.append(self.emoji["emoji_inv"][word])
- elif word == "":
- words.append(" ")
- elif word == "
":
- words.append(breakline)
- elif word == "":
- words.append("\t")
- elif word == "":
- words.append("▀")
- elif word == "":
- words.append("ǀ")
- elif word == "":
- words.append("‖")
- else:
- words.append(word)
- if len(byte_tokens) > 0:
- words.append(bytearray(byte_tokens).decode("utf-8", errors="replace"))
- text = "".join(words)
- return text
diff --git a/transformers/models/gpt_sw3/__init__.py b/transformers/models/gpt_sw3/__init__.py
deleted file mode 100644
index e7c08f0e27e747ea5468e0f9f014df4225dbd424..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_sw3/__init__.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available
-
-
-_import_structure = {}
-
-try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["tokenization_gpt_sw3"] = ["GPTSw3Tokenizer"]
-
-
-if TYPE_CHECKING:
- try:
- if not is_sentencepiece_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .tokenization_gpt_sw3 import GPTSw3Tokenizer
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gpt_sw3/__pycache__/__init__.cpython-310.pyc b/transformers/models/gpt_sw3/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index b744b9452ce73122bc8049343a141e173c6009f4..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_sw3/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_sw3/__pycache__/convert_megatron_to_pytorch.cpython-310.pyc b/transformers/models/gpt_sw3/__pycache__/convert_megatron_to_pytorch.cpython-310.pyc
deleted file mode 100644
index d27ac53f521ca376b324ec1db05d18f1d898ed77..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_sw3/__pycache__/convert_megatron_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_sw3/__pycache__/tokenization_gpt_sw3.cpython-310.pyc b/transformers/models/gpt_sw3/__pycache__/tokenization_gpt_sw3.cpython-310.pyc
deleted file mode 100644
index 7cc693138109ba1a03b91ecbebd0f3f3f960142e..0000000000000000000000000000000000000000
Binary files a/transformers/models/gpt_sw3/__pycache__/tokenization_gpt_sw3.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gpt_sw3/convert_megatron_to_pytorch.py b/transformers/models/gpt_sw3/convert_megatron_to_pytorch.py
deleted file mode 100644
index 5562efa287475be8786c28845124795951f6bfa6..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_sw3/convert_megatron_to_pytorch.py
+++ /dev/null
@@ -1,197 +0,0 @@
-# Copyright 2022 The HuggingFace Inc. team and the AI-Sweden team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" Convert GPT-SW3 megatron checkpoints to pytorch"""
-
-import argparse
-import os
-from os.path import isfile
-
-import torch
-
-from transformers import GPT2Config
-
-
-def recursive_print(name, val, spaces=0):
- # Format the message.
- if name is None:
- msg = None
- else:
- fmt = "." * max(0, spaces - 2) + "# {:" + str(50 - spaces) + "s}"
- msg = fmt.format(name)
-
- # Print and recurse (if needed).
- if isinstance(val, dict):
- if msg is not None:
- print(msg)
- for k in val.keys():
- recursive_print(k, val[k], spaces + 2)
- elif isinstance(val, torch.Tensor):
- print(msg, ":", val.size())
- else:
- print(msg, ":", val)
-
-
-def fix_query_key_value_ordering(param, num_splits, num_heads, hidden_size):
- # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
- # for compatibility with later versions of NVIDIA Megatron-LM.
- # The inverse operation is performed inside Megatron-LM to read checkpoints:
- # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
- # If param is the weight tensor of the self-attention block, the returned tensor
- # will have to be transposed one more time to be read by HuggingFace GPT2.
- input_shape = param.size()
- # other versions store [num_heads * num_splits * hidden_size, :]
- saved_shape = (num_heads, num_splits, hidden_size) + input_shape[1:]
- param = param.view(*saved_shape)
- param = param.transpose(0, 1).contiguous()
- param = param.view(*input_shape)
- return param
-
-
-def convert_megatron_checkpoint(sd_megatron, config):
- """
- Converts a Megatron checkpoint to a HuggingFace GPT-SW3 checkpoint.
- """
- n_positions = config.n_positions
- layers = config.n_layer
- vocab_size = config.vocab_size
- heads = config.n_head
- hidden_size_per_head = config.n_embd // config.n_head
-
- word_embeddings = sd_megatron["model.language_model.embedding.word_embeddings.weight"][:vocab_size, :]
- sd_hf = {
- "transformer.wte.weight": word_embeddings,
- "transformer.wpe.weight": sd_megatron["model.language_model.embedding.position_embeddings.weight"],
- "transformer.ln_f.weight": sd_megatron["model.language_model.encoder.final_layernorm.weight"],
- "transformer.ln_f.bias": sd_megatron["model.language_model.encoder.final_layernorm.bias"],
- }
-
- pf = "model.language_model.encoder.layers."
- for i in range(layers):
- causal_mask = torch.tril(torch.ones((n_positions, n_positions), dtype=torch.bool))
- causal_mask = causal_mask.view(1, 1, n_positions, n_positions)
- sd_hf[f"transformer.h.{i}.attn.bias"] = causal_mask
- sd_hf[f"transformer.h.{i}.attn.masked_bias"] = torch.tensor(-1e4, dtype=torch.bfloat16)
-
- sd_hf[f"transformer.h.{i}.ln_1.weight"] = sd_megatron[f"{pf}{i}.input_layernorm.weight"]
- sd_hf[f"transformer.h.{i}.ln_1.bias"] = sd_megatron[f"{pf}{i}.input_layernorm.bias"]
-
- val1 = sd_megatron[f"{pf}{i}.self_attention.query_key_value.weight"]
- val1 = fix_query_key_value_ordering(val1, 3, heads, hidden_size_per_head)
- sd_hf[f"transformer.h.{i}.attn.c_attn.weight"] = val1.transpose(0, 1).contiguous()
-
- val2 = sd_megatron[f"{pf}{i}.self_attention.query_key_value.bias"]
- val2 = fix_query_key_value_ordering(val2, 3, heads, hidden_size_per_head)
- sd_hf[f"transformer.h.{i}.attn.c_attn.bias"] = val2
-
- sd_hf[f"transformer.h.{i}.attn.c_proj.weight"] = sd_megatron[f"{pf}{i}.self_attention.dense.weight"].transpose(
- 0, 1
- )
- sd_hf[f"transformer.h.{i}.attn.c_proj.bias"] = sd_megatron[f"{pf}{i}.self_attention.dense.bias"]
- sd_hf[f"transformer.h.{i}.ln_2.weight"] = sd_megatron[f"{pf}{i}.post_attention_layernorm.weight"]
- sd_hf[f"transformer.h.{i}.ln_2.bias"] = sd_megatron[f"{pf}{i}.post_attention_layernorm.bias"]
- sd_hf[f"transformer.h.{i}.mlp.c_fc.weight"] = sd_megatron[f"{pf}{i}.mlp.dense_h_to_4h.weight"].transpose(0, 1)
- sd_hf[f"transformer.h.{i}.mlp.c_fc.bias"] = sd_megatron[f"{pf}{i}.mlp.dense_h_to_4h.bias"]
- sd_hf[f"transformer.h.{i}.mlp.c_proj.weight"] = sd_megatron[f"{pf}{i}.mlp.dense_4h_to_h.weight"].transpose(
- 0, 1
- )
- sd_hf[f"transformer.h.{i}.mlp.c_proj.bias"] = sd_megatron[f"{pf}{i}.mlp.dense_4h_to_h.bias"]
-
- # For LM head, transformers' wants the matrix to weight embeddings.
- sd_hf["lm_head.weight"] = word_embeddings
-
- return sd_hf
-
-
-def copy_config(config_hf, config_megatron):
- """Copy the config from Megatron to hf."""
- config_hf.vocab_size = 64000
- config_hf.n_positions = config_megatron["encoder_seq_length"]
- config_hf.n_embd = config_megatron["hidden_size"]
- config_hf.n_layer = config_megatron["num_layers"]
- config_hf.n_head = config_megatron["num_attention_heads"]
- config_hf.n_inner = config_megatron["ffn_hidden_size"]
- config_hf.activation_function = "gelu"
- config_hf.resid_pdrop = 0.1
- config_hf.embd_pdrop = 0.1
- config_hf.attn_pdrop = 0.1
- config_hf.layer_norm_epsilon = config_megatron["layernorm_epsilon"] # 1e-5
- config_hf.initializer_range = config_megatron["init_method_std"] # 0.02
- config_hf.apply_query_key_layer_scaling = config_megatron["apply_query_key_layer_scaling"] # True
- config_hf.normalize_attention_scores = True
- config_hf.use_cache = True
-
- # This identifies the 6.7B (7B) model which uses a different tokenizer
- if config_megatron["hidden_size"] == 4096:
- config_hf.bos_token_id = 1 # <|endoftext|>
- config_hf.eos_token_id = 1 # <|endoftext|>
- config_hf.pad_token_id = 0 #
- else:
- config_hf.bos_token_id = 2 #
- config_hf.eos_token_id = 3 # <|endoftext|>
- config_hf.pad_token_id = 0 #
-
- return config_hf
-
-
-def main(args):
- print(args)
-
- checkpoint_path = args.checkpoint_path
- save_path = args.save_path
- if isfile(checkpoint_path):
- raise FileNotFoundError(f"ERROR! could not find file {checkpoint_path}")
-
- # Load the model.
- checkpoint = torch.load(checkpoint_path, map_location="cpu")
-
- # Load the config.
- config_megatron = checkpoint["hyper_parameters"]["cfg"]
- config_hf = GPT2Config()
- config_hf = copy_config(config_hf=config_hf, config_megatron=config_megatron)
- config_hf.architectures = ["GPT2LMHeadModel"]
-
- sd_megatron = checkpoint["state_dict"]
-
- # Convert.
- print("Converting")
- sd_hf = convert_megatron_checkpoint(sd_megatron, config_hf)
-
- # Print the structure of converted state dict.
- if args.print_checkpoint_structure:
- recursive_print(None, sd_hf)
-
- config_hf.tokenizer_class = "GPTSw3Tokenizer"
-
- # Store the config to file.
- print("Saving config")
- config_hf.save_pretrained(save_path)
-
- # Store the state_dict to file.
- output_checkpoint_file = os.path.join(save_path, "pytorch_model.bin")
- print(f'Saving checkpoint to "{output_checkpoint_file}"')
- torch.save(sd_hf, output_checkpoint_file)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "--checkpoint_path",
- type=str,
- required=True,
- help="e.g. megatron_gpt--val_loss=2.42-step=38000-consumed_samples=54720000",
- )
- parser.add_argument("--save_path", type=str, required=True, help="e.g. /home/user/gpt-sw3/hf")
- parser.add_argument("--print-checkpoint-structure", action="store_true")
- _args = parser.parse_args()
- main(_args)
diff --git a/transformers/models/gpt_sw3/tokenization_gpt_sw3.py b/transformers/models/gpt_sw3/tokenization_gpt_sw3.py
deleted file mode 100644
index 7bb2e51f04a0785e59050527c9d8da26ebd03ef7..0000000000000000000000000000000000000000
--- a/transformers/models/gpt_sw3/tokenization_gpt_sw3.py
+++ /dev/null
@@ -1,318 +0,0 @@
-"""The tokenizer used by the GPT-SW3 models."""
-
-import os
-import re
-import unicodedata
-from shutil import copyfile
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import sentencepiece as spm
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...utils import is_torch_available, logging
-
-
-if is_torch_available():
- import torch
-
-
-logger = logging.get_logger(__name__)
-VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}
-
-
-class GPTSw3Tokenizer(PreTrainedTokenizer):
- """
- Construct an GPTSw3 tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
-
- This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
- this superclass for more information regarding those methods.
-
- Example usage:
- ```python
- >>> from transformers import GPTSw3Tokenizer
-
- >>> tokenizer = GPTSw3Tokenizer.from_pretrained("AI-Sweden-Models/gpt-sw3-126m")
- >>> tokenizer("Svenska är kul!")["input_ids"]
- [1814, 377, 3617, 63504]
- ```
-
- Args:
- vocab_file (`str`):
- [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
- contains the vocabulary necessary to instantiate a tokenizer.
- do_lower_case (`bool`, *optional*, defaults to `False`):
- Whether or not to lowercase the input when tokenizing.
- remove_space (`bool`, *optional*, defaults to `False`):
- Whether or not to strip the text when tokenizing (removing excess spaces before and after the string).
- keep_accents (`bool`, *optional*, defaults to `False`):
- Whether or not to keep accents when tokenizing.
- pad_token (`str`, *optional*):
- The token used for padding, for example when batching sequences of different lengths. If not provided, will
- default to '' or '' depending on model size.
- unk_token (`str`, *optional*):
- The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
- token instead. If not provided, will default to ''.
- eos_token (`str`, *optional*):
- The end of sequence token seen during pretraining. If not provided, will default to '<|endoftext|>'
- bos_token (`str`, *optional*):
- The beginning of sequence token that can be used for downstream task, was not seen during pretraining. If
- not provided, will default to '' or '<|endoftext|>', depending on model size.
- sp_model_kwargs (`dict`, *optional*):
- Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
- SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
- to set:
-
- - `enable_sampling`: Enable subword regularization.
- - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
-
- - `nbest_size = {0,1}`: No sampling is performed.
- - `nbest_size > 1`: samples from the nbest_size results.
- - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
- using forward-filtering-and-backward-sampling algorithm.
-
- - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
- BPE-dropout.
-
- Attributes:
- sp_model (`SentencePieceProcessor`):
- The *SentencePiece* processor that is used for every conversion (string, tokens and IDs).
- whitespaces (`set`):
- The whitespaces that are replaced in the whitespace normalization in preprocessing.
- non_printing_characters_re (`Pattern`):
- The compiled regular expression to remove non-printing characters in preprocessing.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask"]
-
- def __init__(
- self,
- vocab_file,
- do_lower_case=False,
- remove_space=False,
- keep_accents=False,
- pad_token=None,
- unk_token=None,
- eos_token=None,
- bos_token=None,
- sp_model_kwargs: Optional[Dict[str, Any]] = None,
- **kwargs,
- ) -> None:
- self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
-
- name_or_path = kwargs.get("name_or_path")
- if name_or_path is None:
- logger.warning(
- "name_or_path not provided, will work for all GPTSw3 models except gpt-sw3-7b,"
- " you are testing the model, this can safely be ignored"
- )
- name_or_path = "None"
-
- # Default definitions for our 2 tokenizer versions, with None-checks to enable proper testing
- eos_token = "<|endoftext|>" if eos_token is None else eos_token
- unk_token = "" if unk_token is None else unk_token
- if "gpt-sw3-7b" in name_or_path:
- pad_token = unk_token if pad_token is None else pad_token
- bos_token = eos_token if bos_token is None else bos_token
- else:
- pad_token = "" if pad_token is None else pad_token
- bos_token = "" if bos_token is None else bos_token
-
- self.do_lower_case = do_lower_case
- self.remove_space = remove_space
- self.keep_accents = keep_accents
- self.vocab_file = vocab_file
-
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(vocab_file)
-
- # Used for whitespace normalization in input texts
- # fmt : off
- self.whitespaces = {" ", " ", " ", " ", " ", " ", " ", " ", " ", " ", "", ""}
- # fmt : on
-
- # Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
- self.non_printing_characters_re = re.compile(
- f"[{''.join(map(chr, list(range(0, 9)) + list(range(11, 32)) + list(range(127, 160)) + [160, 173, 8203]))}]"
- )
-
- super().__init__(
- do_lower_case=do_lower_case,
- remove_space=remove_space,
- keep_accents=keep_accents,
- bos_token=bos_token,
- eos_token=eos_token,
- unk_token=unk_token,
- pad_token=pad_token,
- sp_model_kwargs=self.sp_model_kwargs,
- **kwargs,
- )
-
- # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.__getstate__
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sp_model"] = None
- return state
-
- # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.__setstate__
- def __setstate__(self, d):
- self.__dict__ = d
-
- # for backward compatibility
- if not hasattr(self, "sp_model_kwargs"):
- self.sp_model_kwargs = {}
-
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(self.vocab_file)
-
- @property
- # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.vocab_size
- def vocab_size(self) -> int:
- return len(self.sp_model)
-
- def preprocess_text(self, text: str) -> str:
- """
- Returns the preprocessed text. This procedure is identical to what was used when training the tokenizer.
- """
-
- # Remove non-printing characters
- text = self.non_printing_characters_re.sub("", text)
-
- # Normalize whitespaces
- text = "".join([char if char not in self.whitespaces else " " for char in text])
-
- # NFC Unicode normalization
- text = unicodedata.normalize("NFC", text)
- return text
-
- def _tokenize(self, text: str, **kwargs) -> List[str]:
- text = self.preprocess_text(text)
- return self.sp_model.encode(text, out_type=str)
-
- def _convert_token_to_id(self, token: str) -> int:
- """Converts a token (str) to an id (int) using the vocab."""
- return self.sp_model.PieceToId(token)
-
- def _convert_id_to_token(self, index: int) -> str:
- """Converts an index (int) to a token (str) using the vocab."""
- return self.sp_model.IdToPiece(index)
-
- @staticmethod
- def clean_up_tokenization(out_string: str) -> str:
- """Returns the input string, this function is overridden to remove the default clean up."""
- return out_string
-
- def convert_tokens_to_string(self, tokens: List[str]) -> str:
- """Converts a sequence of tokens (strings) to a single string. Special tokens remain intact."""
- current_sub_tokens = []
- out_string = ""
- prev_is_special = False
- for token in tokens:
- # make sure that special tokens are not decoded using sentencepiece model
- if token in self.all_special_tokens:
- # TODO: Check if this is needed, as it ensures that decode(encode(doc)) != doc by adding extra whitespace in the decoded document
- if not prev_is_special:
- out_string += " "
-
- out_string += self.sp_model.decode(current_sub_tokens) + token
- prev_is_special = True
- current_sub_tokens = []
- else:
- current_sub_tokens.append(token)
- prev_is_special = False
- out_string += self.sp_model.decode(current_sub_tokens)
-
- return out_string
-
- # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.get_vocab
- def get_vocab(self) -> Dict[str, int]:
- vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
- vocab.update(self.added_tokens_encoder)
- return vocab
-
- # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- if not os.path.isdir(save_directory):
- logger.error(f"Vocabulary path ({save_directory}) should be a directory")
- return
- out_vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
-
- if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
- copyfile(self.vocab_file, out_vocab_file)
- elif not os.path.isfile(self.vocab_file):
- with open(out_vocab_file, "wb") as fi:
- content_spiece_model = self.sp_model.serialized_model_proto()
- fi.write(content_spiece_model)
-
- return (out_vocab_file,)
-
- def encode_fast(
- self, text: Union[str, List[str]], return_tensors: Union[str, bool] = False
- ) -> Union[List[int], List[List[int]], "torch.Tensor"]:
- """
- Encodes a text or batch of texts to token ids using preprocessing and the raw SP tokenizer. This has reduced
- functionality but is often much faster.
-
- Does NOT handle special tokens correctly, these can manually be added as ids afterwards.
-
- Does NOT support padding, these can manually be added as ids afterwards.
-
- Use default HuggingFace tokenization methods for full functionality.
-
- Args:
- text (`str` or `List[str]`): One or several text(s) to convert to token ids.
- return_tensors (`str` or `bool`): Returns PyTorch tensors if set to True or "pt"
-
- Returns:
- `List[int]`, `List[List[int]]`, or `torch.Tensor`: The encoded text(s) as token ids.
- """
-
- if isinstance(text, str):
- text = self.preprocess_text(text)
- token_ids = self.sp_model.encode(text)
- else:
- text = [self.preprocess_text(t) for t in text]
- token_ids = self.sp_model.encode(text)
-
- if return_tensors is True or return_tensors == "pt":
- token_ids = torch.tensor(token_ids)
-
- return token_ids
-
- def decode_fast(self, token_ids: Union[int, List[int]]) -> str:
- """
- Encodes a text or batch of texts to token ids using preprocessing and the raw SP tokenizer. This has reduced
- functionality but is often much faster.
-
- Args:
- token_ids (`int` or `List[int]`): Encoded token or text as token id(s).
-
- Returns:
- `str`: Decoded text
- """
-
- return self.sp_model.decode(token_ids)
-
- @property
- def default_chat_template(self):
- """
- This chat template formats messages like an instant messenger chat log, with "User:" and "Bot:" strings
- preceding messages. BOS tokens are added between all messages.
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- return (
- "{{ eos_token }}{{ bos_token }}"
- "{% for message in messages %}"
- "{% if message['role'] == 'user' %}{{ 'User: ' + message['content']}}"
- "{% else %}{{ 'Bot: ' + message['content']}}{% endif %}"
- "{{ message['text'] }}{{ bos_token }}"
- "{% endfor %}"
- "Bot:"
- )
diff --git a/transformers/models/gptj/__init__.py b/transformers/models/gptj/__init__.py
deleted file mode 100644
index 4e59ed4706204894516b966975dbbb88d462ab29..0000000000000000000000000000000000000000
--- a/transformers/models/gptj/__init__.py
+++ /dev/null
@@ -1,112 +0,0 @@
-# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_flax_available,
- is_tf_available,
- is_torch_available,
-)
-
-
-_import_structure = {"configuration_gptj": ["GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTJConfig", "GPTJOnnxConfig"]}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_gptj"] = [
- "GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GPTJForCausalLM",
- "GPTJForQuestionAnswering",
- "GPTJForSequenceClassification",
- "GPTJModel",
- "GPTJPreTrainedModel",
- ]
-
-try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_tf_gptj"] = [
- "TFGPTJForCausalLM",
- "TFGPTJForQuestionAnswering",
- "TFGPTJForSequenceClassification",
- "TFGPTJModel",
- "TFGPTJPreTrainedModel",
- ]
-
-try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_flax_gptj"] = [
- "FlaxGPTJForCausalLM",
- "FlaxGPTJModel",
- "FlaxGPTJPreTrainedModel",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_gptj import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTJConfig, GPTJOnnxConfig
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_gptj import (
- GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST,
- GPTJForCausalLM,
- GPTJForQuestionAnswering,
- GPTJForSequenceClassification,
- GPTJModel,
- GPTJPreTrainedModel,
- )
-
- try:
- if not is_tf_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_tf_gptj import (
- TFGPTJForCausalLM,
- TFGPTJForQuestionAnswering,
- TFGPTJForSequenceClassification,
- TFGPTJModel,
- TFGPTJPreTrainedModel,
- )
-
- try:
- if not is_flax_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gptj/__pycache__/__init__.cpython-310.pyc b/transformers/models/gptj/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index 8e9ba7837af42491362a3e73ca211486eef70cd2..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptj/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptj/__pycache__/configuration_gptj.cpython-310.pyc b/transformers/models/gptj/__pycache__/configuration_gptj.cpython-310.pyc
deleted file mode 100644
index 0d1862bb025724cb91fcb56e909d4cf8e050669d..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptj/__pycache__/configuration_gptj.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptj/__pycache__/modeling_flax_gptj.cpython-310.pyc b/transformers/models/gptj/__pycache__/modeling_flax_gptj.cpython-310.pyc
deleted file mode 100644
index 70a7b207e5a065ddc20fd45183127aea714f17d1..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptj/__pycache__/modeling_flax_gptj.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptj/__pycache__/modeling_gptj.cpython-310.pyc b/transformers/models/gptj/__pycache__/modeling_gptj.cpython-310.pyc
deleted file mode 100644
index d62774fa27c50eb1ba7bb3be7f181aeee4d07a82..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptj/__pycache__/modeling_gptj.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptj/__pycache__/modeling_tf_gptj.cpython-310.pyc b/transformers/models/gptj/__pycache__/modeling_tf_gptj.cpython-310.pyc
deleted file mode 100644
index f3cc0d1a8dd66522b4697433985ee86e9e50a712..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptj/__pycache__/modeling_tf_gptj.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptj/configuration_gptj.py b/transformers/models/gptj/configuration_gptj.py
deleted file mode 100644
index 56d6042764a19a53f9181fed5a43f4c789361f8b..0000000000000000000000000000000000000000
--- a/transformers/models/gptj/configuration_gptj.py
+++ /dev/null
@@ -1,218 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPT-J model configuration"""
-from collections import OrderedDict
-from typing import Any, List, Mapping, Optional
-
-from ... import PreTrainedTokenizer, TensorType, is_torch_available
-from ...configuration_utils import PretrainedConfig
-from ...onnx import OnnxConfigWithPast, PatchingSpec
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GPTJConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GPTJModel`]. It is used to instantiate a GPT-J
- model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
- defaults will yield a similar configuration to that of the GPT-J
- [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/gpt-j-6B) architecture. Configuration objects inherit from
- [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`]
- for more information.
-
- Args:
- vocab_size (`int`, *optional*, defaults to 50400):
- Vocabulary size of the GPT-J model. Defines the number of different tokens that can be represented by the
- `inputs_ids` passed when calling [`GPTJModel`].
- n_positions (`int`, *optional*, defaults to 2048):
- The maximum sequence length that this model might ever be used with. Typically set this to something large
- just in case (e.g., 512 or 1024 or 2048).
- n_embd (`int`, *optional*, defaults to 4096):
- Dimensionality of the embeddings and hidden states.
- n_layer (`int`, *optional*, defaults to 28):
- Number of hidden layers in the Transformer encoder.
- n_head (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- rotary_dim (`int`, *optional*, defaults to 64):
- Number of dimensions in the embedding that Rotary Position Embedding is applied to.
- n_inner (`int`, *optional*, defaults to None):
- Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
- activation_function (`str`, *optional*, defaults to `"gelu_new"`):
- Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new"]`.
- resid_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- embd_pdrop (`int`, *optional*, defaults to 0.1):
- The dropout ratio for the embeddings.
- attn_pdrop (`float`, *optional*, defaults to 0.1):
- The dropout ratio for the attention.
- layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
- The epsilon to use in the layer normalization layers.
- initializer_range (`float`, *optional*, defaults to 0.02):
- The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models).
-
- Example:
-
- ```python
- >>> from transformers import GPTJModel, GPTJConfig
-
- >>> # Initializing a GPT-J 6B configuration
- >>> configuration = GPTJConfig()
-
- >>> # Initializing a model from the configuration
- >>> model = GPTJModel(configuration)
-
- >>> # Accessing the model configuration
- >>> configuration = model.config
- ```"""
-
- model_type = "gptj"
- attribute_map = {
- "max_position_embeddings": "n_positions",
- "hidden_size": "n_embd",
- "num_attention_heads": "n_head",
- "num_hidden_layers": "n_layer",
- }
-
- def __init__(
- self,
- vocab_size=50400,
- n_positions=2048,
- n_embd=4096,
- n_layer=28,
- n_head=16,
- rotary_dim=64,
- n_inner=None,
- activation_function="gelu_new",
- resid_pdrop=0.0,
- embd_pdrop=0.0,
- attn_pdrop=0.0,
- layer_norm_epsilon=1e-5,
- initializer_range=0.02,
- use_cache=True,
- bos_token_id=50256,
- eos_token_id=50256,
- tie_word_embeddings=False,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.n_positions = n_positions
- self.n_embd = n_embd
- self.n_layer = n_layer
- self.n_head = n_head
- self.n_inner = n_inner
- self.rotary_dim = rotary_dim
- self.activation_function = activation_function
- self.resid_pdrop = resid_pdrop
- self.embd_pdrop = embd_pdrop
- self.attn_pdrop = attn_pdrop
- self.layer_norm_epsilon = layer_norm_epsilon
- self.initializer_range = initializer_range
- self.use_cache = use_cache
-
- self.bos_token_id = bos_token_id
- self.eos_token_id = eos_token_id
-
- super().__init__(
- bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
- )
-
-
-# Copied from transformers.models.gpt2.configuration_gpt2.GPT2OnnxConfig
-class GPTJOnnxConfig(OnnxConfigWithPast):
- def __init__(
- self,
- config: PretrainedConfig,
- task: str = "default",
- patching_specs: List[PatchingSpec] = None,
- use_past: bool = False,
- ):
- super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past)
- if not getattr(self._config, "pad_token_id", None):
- # TODO: how to do that better?
- self._config.pad_token_id = 0
-
- @property
- def inputs(self) -> Mapping[str, Mapping[int, str]]:
- common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
- if self.use_past:
- self.fill_with_past_key_values_(common_inputs, direction="inputs")
- common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"}
- else:
- common_inputs["attention_mask"] = {0: "batch", 1: "sequence"}
-
- return common_inputs
-
- @property
- def num_layers(self) -> int:
- return self._config.n_layer
-
- @property
- def num_attention_heads(self) -> int:
- return self._config.n_head
-
- def generate_dummy_inputs(
- self,
- tokenizer: PreTrainedTokenizer,
- batch_size: int = -1,
- seq_length: int = -1,
- is_pair: bool = False,
- framework: Optional[TensorType] = None,
- ) -> Mapping[str, Any]:
- common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
- tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
- )
-
- # We need to order the input in the way they appears in the forward()
- ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]})
-
- # Need to add the past_keys
- if self.use_past:
- if not is_torch_available():
- raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
- else:
- import torch
-
- batch, seqlen = common_inputs["input_ids"].shape
- # Not using the same length for past_key_values
- past_key_values_length = seqlen + 2
- past_shape = (
- batch,
- self.num_attention_heads,
- past_key_values_length,
- self._config.hidden_size // self.num_attention_heads,
- )
- ordered_inputs["past_key_values"] = [
- (torch.zeros(past_shape), torch.zeros(past_shape)) for _ in range(self.num_layers)
- ]
-
- ordered_inputs["attention_mask"] = common_inputs["attention_mask"]
- if self.use_past:
- mask_dtype = ordered_inputs["attention_mask"].dtype
- ordered_inputs["attention_mask"] = torch.cat(
- [ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1
- )
-
- return ordered_inputs
-
- @property
- def default_onnx_opset(self) -> int:
- return 13
diff --git a/transformers/models/gptj/modeling_flax_gptj.py b/transformers/models/gptj/modeling_flax_gptj.py
deleted file mode 100644
index 9f0d4d6e86000384544fa2873690b09d34a050a2..0000000000000000000000000000000000000000
--- a/transformers/models/gptj/modeling_flax_gptj.py
+++ /dev/null
@@ -1,718 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The EleutherAI and The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from functools import partial
-from typing import Optional, Tuple
-
-import flax.linen as nn
-import jax
-import jax.numpy as jnp
-import numpy as np
-from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
-from flax.linen import combine_masks, make_causal_mask
-from flax.linen.attention import dot_product_attention_weights
-from flax.traverse_util import flatten_dict, unflatten_dict
-from jax import lax
-
-from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput
-from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring
-from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
-from .configuration_gptj import GPTJConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "gptj"
-_CONFIG_FOR_DOC = "GPTJConfig"
-
-
-GPTJ_START_DOCSTRING = r"""
-
- This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a Flax Linen
- [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a
- regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.
-
- Finally, this model supports inherent JAX features such as:
-
- - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
-
- Parameters:
- config ([`GPTJConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.
- dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
- The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
- `jax.numpy.bfloat16` (on TPUs).
-
- This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
- specified all the computation will be performed with the given `dtype`.
-
- **Note that this only specifies the dtype of the computation and does not influence the dtype of model
- parameters.**
-
- If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and
- [`~FlaxPreTrainedModel.to_bf16`].
-"""
-
-GPTJ_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`numpy.ndarray` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length`. Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
- past_key_values (`Dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):
- Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
- auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-
-def create_sinusoidal_positions(num_pos, dim):
- inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2) / dim))
- sinusoid_inp = np.einsum("i , j -> i j", np.arange(num_pos), inv_freq).astype("float32")
- sin, cos = np.sin(sinusoid_inp), np.cos(sinusoid_inp)
-
- sentinel = dim // 2 + dim % 2
- out = np.zeros((num_pos, dim))
- out[:, 0:sentinel] = sin
- out[:, sentinel:] = cos
-
- return jnp.array(out)
-
-
-def rotate_every_two(tensor):
- rotate_half_tensor = jnp.stack((-tensor[:, :, :, 1::2], tensor[:, :, :, ::2]), axis=-1)
- rotate_half_tensor = rotate_half_tensor.reshape(rotate_half_tensor.shape[:-2] + (-1,))
- return rotate_half_tensor
-
-
-def apply_rotary_pos_emb(tensor, sincos):
- sin_pos, cos_pos = sincos
- sin_pos = sin_pos[:, :, None, :].repeat(2, 3)
- cos_pos = cos_pos[:, :, None, :].repeat(2, 3)
- return (tensor * cos_pos) + (rotate_every_two(tensor) * sin_pos)
-
-
-class FlaxGPTJAttention(nn.Module):
- config: GPTJConfig
- dtype: jnp.dtype = jnp.float32
- causal: bool = True
- is_cross_attention: bool = False
-
- def setup(self):
- config = self.config
- self.embed_dim = config.hidden_size
- self.num_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_heads
-
- self.rotary_dim = config.rotary_dim
-
- dense = partial(
- nn.Dense,
- self.embed_dim,
- use_bias=False,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(self.config.initializer_range),
- )
-
- self.q_proj, self.k_proj, self.v_proj = dense(), dense(), dense()
- self.out_proj = dense()
-
- self.resid_dropout = nn.Dropout(rate=config.resid_pdrop)
-
- self.causal_mask = make_causal_mask(jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool")
-
- pos_embd_dim = self.rotary_dim or self.embed_dim
- self.embed_positions = create_sinusoidal_positions(config.max_position_embeddings, pos_embd_dim)
-
- def _split_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))
-
- def _merge_heads(self, hidden_states):
- return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))
-
- @nn.compact
- def _concatenate_to_cache(self, key, value, query, attention_mask):
- """
- This function takes projected key, value states from a single input token and concatenates the states to cached
- states from previous steps. This function is slighly adapted from the official Flax repository:
- https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252
- """
- # detect if we're initializing by absence of existing cache data.
- is_initialized = self.has_variable("cache", "cached_key")
- cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype)
- cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype)
- cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))
-
- if is_initialized:
- *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape
- # update key, value caches with our new 1d spatial slices
- cur_index = cache_index.value
- indices = (0,) * len(batch_dims) + (cur_index, 0, 0)
- key = lax.dynamic_update_slice(cached_key.value, key, indices)
- value = lax.dynamic_update_slice(cached_value.value, value, indices)
- cached_key.value = key
- cached_value.value = value
- num_updated_cache_vectors = query.shape[1]
- cache_index.value = cache_index.value + num_updated_cache_vectors
- # causal mask for cached decoder self-attention: our single query position should only attend to those key
- # positions that have already been generated and cached, not the remaining zero elements.
- pad_mask = jnp.broadcast_to(
- jnp.arange(max_length) < cur_index + num_updated_cache_vectors,
- tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),
- )
- attention_mask = combine_masks(pad_mask, attention_mask)
- return key, value, attention_mask
-
- def __call__(
- self,
- hidden_states,
- attention_mask,
- position_ids,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- query = self.q_proj(hidden_states)
- key = self.k_proj(hidden_states)
- value = self.v_proj(hidden_states)
-
- query = self._split_heads(query)
- key = self._split_heads(key)
- value = self._split_heads(value)
-
- sincos = jnp.take(self.embed_positions, position_ids, axis=0)
- sincos = jnp.split(sincos, 2, axis=-1)
- if self.rotary_dim is not None:
- k_rot = key[:, :, :, : self.rotary_dim]
- k_pass = key[:, :, :, self.rotary_dim :]
-
- q_rot = query[:, :, :, : self.rotary_dim]
- q_pass = query[:, :, :, self.rotary_dim :]
-
- k_rot = apply_rotary_pos_emb(k_rot, sincos)
- q_rot = apply_rotary_pos_emb(q_rot, sincos)
-
- key = jnp.concatenate([k_rot, k_pass], axis=-1)
- query = jnp.concatenate([q_rot, q_pass], axis=-1)
- else:
- key = apply_rotary_pos_emb(key, sincos)
- query = apply_rotary_pos_emb(query, sincos)
-
- query_length, key_length = query.shape[1], key.shape[1]
-
- if self.has_variable("cache", "cached_key"):
- mask_shift = self.variables["cache"]["cache_index"]
- max_decoder_length = self.variables["cache"]["cached_key"].shape[1]
- causal_mask = lax.dynamic_slice(
- self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)
- )
- else:
- causal_mask = self.causal_mask[:, :, :query_length, :key_length]
-
- batch_size = hidden_states.shape[0]
- causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])
-
- attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)
- attention_mask = combine_masks(attention_mask, causal_mask)
-
- dropout_rng = None
- if not deterministic and self.config.attn_pdrop > 0.0:
- dropout_rng = self.make_rng("dropout")
-
- # During fast autoregressive decoding, we feed one position at a time,
- # and cache the keys and values step by step.
- if self.has_variable("cache", "cached_key") or init_cache:
- key, value, attention_mask = self._concatenate_to_cache(key, value, query, attention_mask)
-
- # transform boolean mask into float mask
- attention_bias = lax.select(
- attention_mask > 0,
- jnp.full(attention_mask.shape, 0.0).astype(self.dtype),
- jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),
- )
-
- # usual dot product attention
- attn_weights = dot_product_attention_weights(
- query,
- key,
- bias=attention_bias,
- dropout_rng=dropout_rng,
- dropout_rate=self.config.attn_pdrop,
- deterministic=deterministic,
- dtype=self.dtype,
- precision=None,
- )
-
- attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value)
- attn_output = self._merge_heads(attn_output)
- attn_output = self.out_proj(attn_output)
- attn_output = self.resid_dropout(attn_output, deterministic=deterministic)
-
- outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)
- return outputs
-
-
-class FlaxGPTJMLP(nn.Module):
- config: GPTJConfig
- intermediate_size: int
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- embed_dim = self.config.hidden_size
- kernel_init = jax.nn.initializers.normal(self.config.initializer_range)
-
- self.fc_in = nn.Dense(self.intermediate_size, dtype=self.dtype, kernel_init=kernel_init)
- self.fc_out = nn.Dense(embed_dim, dtype=self.dtype, kernel_init=kernel_init)
-
- self.act = ACT2FN[self.config.activation_function]
- self.dropout = nn.Dropout(rate=self.config.resid_pdrop)
-
- def __call__(self, hidden_states, deterministic: bool = True):
- hidden_states = self.fc_in(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.fc_out(hidden_states)
- hidden_states = self.dropout(hidden_states, deterministic=deterministic)
- return hidden_states
-
-
-class FlaxGPTJBlock(nn.Module):
- config: GPTJConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- hidden_size = self.config.hidden_size
- inner_dim = self.config.n_inner if self.config.n_inner is not None else 4 * hidden_size
-
- self.ln_1 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
- self.attn = FlaxGPTJAttention(self.config, dtype=self.dtype)
-
- self.mlp = FlaxGPTJMLP(self.config, inner_dim, dtype=self.dtype)
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- position_ids=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- ):
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0]
-
- feed_forward_hidden_states = self.mlp(hidden_states, deterministic=deterministic)
- # residual connection
- hidden_states = attn_output + feed_forward_hidden_states + residual
-
- return (hidden_states,) + attn_outputs[1:]
-
-
-class FlaxGPTJPreTrainedModel(FlaxPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTJConfig
- base_model_prefix = "transformer"
- module_class: nn.Module = None
-
- def __init__(
- self,
- config: GPTJConfig,
- input_shape: Tuple = (1, 1),
- seed: int = 0,
- dtype: jnp.dtype = jnp.float32,
- _do_init: bool = True,
- **kwargs,
- ):
- module = self.module_class(config=config, dtype=dtype, **kwargs)
- super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
-
- def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
- # init input tensors
- input_ids = jnp.zeros(input_shape, dtype="i4")
- attention_mask = jnp.ones_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape)
- params_rng, dropout_rng = jax.random.split(rng)
- rngs = {"params": params_rng, "dropout": dropout_rng}
-
- if self.config.add_cross_attention:
- encoder_hidden_states = jnp.zeros(input_shape + (self.config.n_embd,))
- encoder_attention_mask = attention_mask
- module_init_outputs = self.module.init(
- rngs,
- input_ids,
- attention_mask,
- position_ids,
- encoder_hidden_states,
- encoder_attention_mask,
- return_dict=False,
- )
- else:
- module_init_outputs = self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False)
-
- random_params = module_init_outputs["params"]
-
- if params is not None:
- random_params = flatten_dict(unfreeze(random_params))
- params = flatten_dict(unfreeze(params))
- for missing_key in self._missing_keys:
- params[missing_key] = random_params[missing_key]
- self._missing_keys = set()
- return freeze(unflatten_dict(params))
- else:
- return random_params
-
- def init_cache(self, batch_size, max_length):
- r"""
- Args:
- batch_size (`int`):
- batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
- max_length (`int`):
- maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
- cache.
- """
- # init input variables to retrieve cache
- input_ids = jnp.ones((batch_size, max_length))
- attention_mask = jnp.ones_like(input_ids)
- position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
-
- init_variables = self.module.init(
- jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True
- )
- return init_variables["cache"]
-
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING)
- def __call__(
- self,
- input_ids,
- attention_mask=None,
- position_ids=None,
- params: dict = None,
- past_key_values: dict = None,
- dropout_rng: jax.random.PRNGKey = None,
- train: bool = False,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ):
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- return_dict = return_dict if return_dict is not None else self.config.return_dict
-
- batch_size, sequence_length = input_ids.shape
-
- if position_ids is None:
- if past_key_values is not None:
- raise ValueError("Make sure to provide `position_ids` when passing `past_key_values`.")
-
- position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
-
- if attention_mask is None:
- attention_mask = jnp.ones((batch_size, sequence_length))
-
- # Handle any PRNG if needed
- rngs = {}
- if dropout_rng is not None:
- rngs["dropout"] = dropout_rng
-
- inputs = {"params": params or self.params}
-
- # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be changed by FlaxGPTJAttention module
- if past_key_values:
- inputs["cache"] = past_key_values
- mutable = ["cache"]
- else:
- mutable = False
-
- outputs = self.module.apply(
- inputs,
- jnp.array(input_ids, dtype="i4"),
- jnp.array(attention_mask, dtype="i4"),
- jnp.array(position_ids, dtype="i4"),
- not train,
- False,
- output_attentions,
- output_hidden_states,
- return_dict,
- rngs=rngs,
- mutable=mutable,
- )
-
- # add updated cache to model output
- if past_key_values is not None and return_dict:
- outputs, past_key_values = outputs
- outputs["past_key_values"] = unfreeze(past_key_values["cache"])
- return outputs
- elif past_key_values is not None and not return_dict:
- outputs, past_key_values = outputs
- outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:]
-
- return outputs
-
-
-class FlaxGPTJBlockCollection(nn.Module):
- config: GPTJConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.blocks = [
- FlaxGPTJBlock(self.config, name=str(i), dtype=self.dtype) for i in range(self.config.num_hidden_layers)
- ]
-
- def __call__(
- self,
- hidden_states,
- attention_mask=None,
- position_ids=None,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- all_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
-
- for block in self.blocks:
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- layer_outputs = block(
- hidden_states,
- attention_mask,
- position_ids=position_ids,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- )
- hidden_states = layer_outputs[0]
-
- if output_attentions:
- all_attentions += (layer_outputs[1],)
-
- # this contains possible `None` values - `FlaxGPTJModule` will filter them out
- outputs = (hidden_states, all_hidden_states, all_attentions)
-
- return outputs
-
-
-class FlaxGPTJModule(nn.Module):
- config: GPTJConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.embed_dim = self.config.hidden_size
-
- self.wte = nn.Embed(
- self.config.vocab_size,
- self.config.hidden_size,
- embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
- self.dropout = nn.Dropout(rate=self.config.embd_pdrop)
- self.h = FlaxGPTJBlockCollection(self.config, dtype=self.dtype)
- self.ln_f = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- position_ids,
- deterministic=True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- input_embeds = self.wte(input_ids.astype("i4"))
-
- hidden_states = self.dropout(input_embeds, deterministic=deterministic)
-
- outputs = self.h(
- hidden_states,
- attention_mask,
- position_ids=position_ids,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
- hidden_states = self.ln_f(hidden_states)
-
- if output_hidden_states:
- all_hidden_states = outputs[1] + (hidden_states,)
- outputs = (hidden_states, all_hidden_states) + outputs[2:]
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- if not return_dict:
- return tuple(v for v in outputs if v is not None)
-
- return FlaxBaseModelOutput(
- last_hidden_state=hidden_states,
- hidden_states=outputs[1],
- attentions=outputs[-1],
- )
-
-
-@add_start_docstrings(
- "The bare GPTJ Model transformer outputting raw hidden-states without any specific head on top.",
- GPTJ_START_DOCSTRING,
-)
-class FlaxGPTJModel(FlaxGPTJPreTrainedModel):
- module_class = FlaxGPTJModule
-
-
-append_call_sample_docstring(
- FlaxGPTJModel,
- _CHECKPOINT_FOR_DOC,
- FlaxCausalLMOutput,
- _CONFIG_FOR_DOC,
-)
-
-
-class FlaxGPTJForCausalLMModule(nn.Module):
- config: GPTJConfig
- dtype: jnp.dtype = jnp.float32
-
- def setup(self):
- self.transformer = FlaxGPTJModule(self.config, dtype=self.dtype)
- self.lm_head = nn.Dense(
- self.config.vocab_size,
- dtype=self.dtype,
- kernel_init=jax.nn.initializers.normal(stddev=self.config.initializer_range),
- )
-
- def __call__(
- self,
- input_ids,
- attention_mask,
- position_ids,
- deterministic: bool = True,
- init_cache: bool = False,
- output_attentions: bool = False,
- output_hidden_states: bool = False,
- return_dict: bool = True,
- ):
- outputs = self.transformer(
- input_ids,
- attention_mask,
- position_ids,
- deterministic=deterministic,
- init_cache=init_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- hidden_states = outputs[0]
-
- if self.config.tie_word_embeddings:
- shared_kernel = self.transformer.variables["params"]["wte"]["embedding"].T
- lm_logits = self.lm_head.apply({"params": {"kernel": shared_kernel}}, hidden_states)
- else:
- lm_logits = self.lm_head(hidden_states)
-
- if not return_dict:
- return (lm_logits,) + outputs[1:]
-
- return FlaxCausalLMOutput(logits=lm_logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
-
-
-@add_start_docstrings(
- """
- The GPTJ Model transformer with a language modeling head on top.
- """,
- GPTJ_START_DOCSTRING,
-)
-class FlaxGPTJForCausalLM(FlaxGPTJPreTrainedModel):
- module_class = FlaxGPTJForCausalLMModule
-
- def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jax.Array] = None):
- # initializing the cache
- batch_size, seq_length = input_ids.shape
-
- past_key_values = self.init_cache(batch_size, max_length)
- # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
- # But since GPTJ uses a causal mask, those positions are masked anyways.
- # Thus we can create a single static attention_mask here, which is more efficient for compilation
- extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
- if attention_mask is not None:
- position_ids = attention_mask.cumsum(axis=-1) - 1
- extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, attention_mask, (0, 0))
- else:
- position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length))
-
- return {
- "past_key_values": past_key_values,
- "attention_mask": extended_attention_mask,
- "position_ids": position_ids,
- }
-
- def update_inputs_for_generation(self, model_outputs, model_kwargs):
- model_kwargs["past_key_values"] = model_outputs.past_key_values
- model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1
- return model_kwargs
-
-
-append_call_sample_docstring(
- FlaxGPTJForCausalLM,
- _CHECKPOINT_FOR_DOC,
- FlaxCausalLMOutput,
- _CONFIG_FOR_DOC,
-)
diff --git a/transformers/models/gptj/modeling_gptj.py b/transformers/models/gptj/modeling_gptj.py
deleted file mode 100644
index 3c6ddac4ecf4ca057a17c7fb555912cfda29ad7e..0000000000000000000000000000000000000000
--- a/transformers/models/gptj/modeling_gptj.py
+++ /dev/null
@@ -1,1427 +0,0 @@
-# coding=utf-8
-# Copyright 2021 The EleutherAI and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch GPT-J model."""
-
-import warnings
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.fx
-import torch.nn.functional as F
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
-from ...activations import ACT2FN
-from ...modeling_outputs import (
- BaseModelOutputWithPast,
- CausalLMOutputWithPast,
- QuestionAnsweringModelOutput,
- SequenceClassifierOutputWithPast,
-)
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_flash_attn_2_available,
- is_flash_attn_greater_or_equal_2_10,
- is_torch_fx_proxy,
- logging,
-)
-from ...utils.model_parallel_utils import assert_device_map, get_device_map
-from .configuration_gptj import GPTJConfig
-
-
-if is_flash_attn_2_available():
- from flash_attn import flash_attn_func, flash_attn_varlen_func
- from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "hf-internal-testing/tiny-random-gptj"
-_REAL_CHECKPOINT_FOR_DOC = "EleutherAI/gpt-j-6B"
-_CONFIG_FOR_DOC = "GPTJConfig"
-
-
-from ..deprecated._archive_maps import GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
-def _get_unpad_data(attention_mask):
- seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
- indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
- max_seqlen_in_batch = seqlens_in_batch.max().item()
- cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
- return (
- indices,
- cu_seqlens,
- max_seqlen_in_batch,
- )
-
-
-def create_sinusoidal_positions(num_pos: int, dim: int) -> torch.Tensor:
- inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64) / dim))
- sinusoid_inp = torch.einsum("i , j -> i j", torch.arange(num_pos, dtype=torch.int64).float(), inv_freq).float()
- return torch.cat((torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)), dim=1)
-
-
-@torch.fx.wrap
-def get_embed_positions(embed_positions, position_ids):
- return embed_positions.to(position_ids.device).repeat(position_ids.shape[0], 1, 1)
-
-
-def rotate_every_two(x: torch.Tensor) -> torch.Tensor:
- x1 = x[:, :, :, ::2]
- x2 = x[:, :, :, 1::2]
- x = torch.stack((-x2, x1), dim=-1)
- return x.flatten(-2) # in einsum notation: rearrange(x, '... d j -> ... (d j)')
-
-
-def apply_rotary_pos_emb(tensor: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
- sin = torch.repeat_interleave(sin[:, :, None, :], 2, 3)
- cos = torch.repeat_interleave(cos[:, :, None, :], 2, 3)
- return (tensor * cos) + (rotate_every_two(tensor) * sin)
-
-
-class GPTJAttention(nn.Module):
- def __init__(self, config):
- super().__init__()
- self.config = config
- max_positions = config.max_position_embeddings
- self.register_buffer(
- "bias",
- torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
- 1, 1, max_positions, max_positions
- ),
- persistent=False,
- )
- self.register_buffer("masked_bias", torch.tensor(-1e9), persistent=False)
-
- self.attn_dropout = nn.Dropout(config.attn_pdrop)
- self.resid_dropout = nn.Dropout(config.resid_pdrop)
-
- self.is_causal = True
-
- self.embed_dim = config.hidden_size
- self.num_attention_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_attention_heads
- if self.head_dim * self.num_attention_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and"
- f" `num_attention_heads`: {self.num_attention_heads})."
- )
- self.scale_attn = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype())
-
- self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
- self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
- self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
- self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
- self.rotary_dim = config.rotary_dim
- pos_embd_dim = self.rotary_dim or self.embed_dim
- self.embed_positions = create_sinusoidal_positions(max_positions, pos_embd_dim)
-
- def _split_heads(self, tensor, num_attention_heads, attn_head_size, rotary):
- """
- Splits hidden dim into attn_head_size and num_attention_heads
- """
- new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
- tensor = tensor.view(new_shape)
- if rotary:
- return tensor
- if len(tensor.shape) == 5:
- return tensor.permute(0, 1, 3, 2, 4) # (batch, blocks, head, block_length, head_features)
- elif len(tensor.shape) == 4:
- return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)
- else:
- raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}")
-
- def _merge_heads(self, tensor, num_attention_heads, attn_head_size):
- """
- Merges attn_head_size dim and num_attn_heads dim into hidden dim
- """
- if len(tensor.shape) == 5:
- tensor = tensor.permute(0, 1, 3, 2, 4).contiguous()
- elif len(tensor.shape) == 4:
- tensor = tensor.permute(0, 2, 1, 3).contiguous()
- else:
- raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}")
- new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,)
- return tensor.view(new_shape)
-
- def _attn(
- self,
- query,
- key,
- value,
- attention_mask=None,
- head_mask=None,
- ):
- # compute causal mask from causal mask buffer
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
-
- # Keep the attention weights computation in fp32 to avoid overflow issues
- query = query.to(torch.float32)
- key = key.to(torch.float32)
-
- attn_weights = torch.matmul(query, key.transpose(-1, -2))
-
- mask_value = torch.finfo(attn_weights.dtype).min
- # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
- # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
- mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
- attn_weights = torch.where(causal_mask, attn_weights, mask_value)
-
- attn_weights = attn_weights / self.scale_attn
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
- attn_weights = attn_weights.to(value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def _get_embed_positions(self, position_ids):
- embed_positions = self.embed_positions
- if embed_positions.device != position_ids.device:
- embed_positions = embed_positions.to(position_ids.device)
- self.embed_positions = embed_positions
- return embed_positions.repeat(position_ids.shape[0], 1, 1)
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[
- Tuple[torch.Tensor, Tuple[torch.Tensor]],
- Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
- ]:
- query = self.q_proj(hidden_states)
- key = self.k_proj(hidden_states)
- value = self.v_proj(hidden_states)
-
- query = self._split_heads(query, self.num_attention_heads, self.head_dim, True)
- key = self._split_heads(key, self.num_attention_heads, self.head_dim, True)
- value = self._split_heads(value, self.num_attention_heads, self.head_dim, False)
-
- if is_torch_fx_proxy(position_ids) or torch.jit.is_tracing():
- # The logic to conditionally copy to GPU could not be traced, so we do this
- # every time in the torch.fx case
- embed_positions = get_embed_positions(self.embed_positions, position_ids)
- else:
- embed_positions = self._get_embed_positions(position_ids)
-
- repeated_position_ids = position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1])
- sincos = torch.gather(embed_positions, 1, repeated_position_ids)
- sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1)
-
- if self.rotary_dim is not None:
- k_rot = key[:, :, :, : self.rotary_dim]
- k_pass = key[:, :, :, self.rotary_dim :]
-
- q_rot = query[:, :, :, : self.rotary_dim]
- q_pass = query[:, :, :, self.rotary_dim :]
-
- k_rot = apply_rotary_pos_emb(k_rot, sin, cos)
- q_rot = apply_rotary_pos_emb(q_rot, sin, cos)
-
- key = torch.cat([k_rot, k_pass], dim=-1)
- query = torch.cat([q_rot, q_pass], dim=-1)
- else:
- key = apply_rotary_pos_emb(key, sin, cos)
- query = apply_rotary_pos_emb(query, sin, cos)
-
- key = key.permute(0, 2, 1, 3)
- query = query.permute(0, 2, 1, 3)
-
- if layer_past is not None:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- if use_cache is True:
- # Note that this cast is quite ugly, but is not implemented before ROPE as the original codebase keeps the key in float32 all along the computation.
- # Reference: https://github.com/kingoflolz/mesh-transformer-jax/blob/f8315e3003033b23f21d78361b288953064e0e76/mesh_transformer/layers.py#L128
- present = (key.to(hidden_states.dtype), value)
- else:
- present = None
-
- # compute self-attention: V x Softmax(QK^T)
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
- attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim)
- attn_output = self.out_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs # a, present, (attentions)
-
-
-class GPTJFlashAttention2(GPTJAttention):
- """
- GPTJ flash attention module. This module inherits from `GPTJAttention` as the weights of the module stays
- untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
- flash attention and deal with padding tokens in case the input contains any of them.
- """
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
-
- # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
- # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
- # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
- self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
-
- def forward(
- self,
- hidden_states: torch.FloatTensor,
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[
- Tuple[torch.Tensor, Tuple[torch.Tensor]],
- Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
- ]:
- query = self.q_proj(hidden_states)
- key = self.k_proj(hidden_states)
- value = self.v_proj(hidden_states)
-
- query = self._split_heads(query, self.num_attention_heads, self.head_dim, True)
- key = self._split_heads(key, self.num_attention_heads, self.head_dim, True)
- value = self._split_heads(value, self.num_attention_heads, self.head_dim, False)
-
- if is_torch_fx_proxy(position_ids) or torch.jit.is_tracing():
- # The logic to conditionally copy to GPU could not be traced, so we do this
- # every time in the torch.fx case
- embed_positions = get_embed_positions(self.embed_positions, position_ids)
- else:
- embed_positions = self._get_embed_positions(position_ids)
-
- repeated_position_ids = position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1])
- sincos = torch.gather(embed_positions, 1, repeated_position_ids)
- sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1)
-
- if self.rotary_dim is not None:
- k_rot = key[:, :, :, : self.rotary_dim]
- k_pass = key[:, :, :, self.rotary_dim :]
-
- q_rot = query[:, :, :, : self.rotary_dim]
- q_pass = query[:, :, :, self.rotary_dim :]
-
- k_rot = apply_rotary_pos_emb(k_rot, sin, cos)
- q_rot = apply_rotary_pos_emb(q_rot, sin, cos)
-
- key = torch.cat([k_rot, k_pass], dim=-1)
- query = torch.cat([q_rot, q_pass], dim=-1)
- else:
- key = apply_rotary_pos_emb(key, sin, cos)
- query = apply_rotary_pos_emb(query, sin, cos)
-
- # tanspose to have the desired shape
- # before transpose: batch_size x seq_length x num_attention_heads x head_dim
- # after transpose: batch_size x num_attention_heads x seq_length x head_dim
- key = key.permute(0, 2, 1, 3)
- query = query.permute(0, 2, 1, 3)
- # value: batch_size x num_attention_heads x seq_length x head_dim
-
- if layer_past is not None:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = torch.cat((past_key, key), dim=-2)
- value = torch.cat((past_value, value), dim=-2)
-
- if use_cache is True:
- # Note that this cast is quite ugly, but is not implemented before ROPE as the original codebase keeps the key in float32 all along the computation.
- # Reference: https://github.com/kingoflolz/mesh-transformer-jax/blob/f8315e3003033b23f21d78361b288953064e0e76/mesh_transformer/layers.py#L128
- present = (key.to(hidden_states.dtype), value)
- else:
- present = None
-
- # The Flash attention requires the input to have the shape
- # batch_size x seq_length x head_dim x hidden_dim
- # therefore we need to keep the original shape for query and key, and reshape value
- # to have the correct shape.
- key = key.permute(0, 2, 1, 3).contiguous()
- query = query.permute(0, 2, 1, 3).contiguous()
- value = value.permute(0, 2, 1, 3).contiguous()
-
- # In PEFT, usually we cast the layer norms in float32 for training stability reasons
- # therefore the input hidden states gets silently casted in float32. Hence, we need
- # cast them back in the correct dtype just to be sure everything works as expected.
- # This might slowdown training & inference so it is recommended to not cast the LayerNorms
- # in fp32. (LlamaRMSNorm handles it correctly)
-
- input_dtype = query.dtype
- if input_dtype == torch.float32:
- if torch.is_autocast_enabled():
- target_dtype = torch.get_autocast_gpu_dtype()
- # Handle the case where the model is quantized
- elif hasattr(self.config, "_pre_quantization_dtype"):
- target_dtype = self.config._pre_quantization_dtype
- else:
- target_dtype = self.q_proj.weight.dtype
-
- logger.warning_once(
- f"The input hidden states seems to be silently casted in float32, this might be related to"
- f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
- f" {target_dtype}."
- )
-
- query = query.to(target_dtype)
- key = key.to(target_dtype)
- value = value.to(target_dtype)
-
- attention_dropout = self.config.attn_pdrop if self.training else 0.0 # attn_pdrop in gptj
-
- query_length = query.shape[1]
-
- # Compute attention
- attn_weights = self._flash_attention_forward(
- query,
- key,
- value,
- attention_mask,
- query_length,
- dropout=attention_dropout,
- )
-
- # Reshape outputs
- attn_output = attn_weights.reshape(
- attn_weights.shape[0], attn_weights.shape[1], attn_weights.shape[2] * attn_weights.shape[3]
- )
- attn_output = self.out_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
- def _flash_attention_forward(
- self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
- ):
- """
- Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
- first unpad the input, then computes the attention scores and pad the final attention scores.
-
- Args:
- query_states (`torch.Tensor`):
- Input query states to be passed to Flash Attention API
- key_states (`torch.Tensor`):
- Input key states to be passed to Flash Attention API
- value_states (`torch.Tensor`):
- Input value states to be passed to Flash Attention API
- attention_mask (`torch.Tensor`):
- The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
- position of padding tokens and 1 for the position of non-padding tokens.
- dropout (`float`):
- Attention dropout
- softmax_scale (`float`, *optional*):
- The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
- """
- if not self._flash_attn_uses_top_left_mask:
- causal = self.is_causal
- else:
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
- causal = self.is_causal and query_length != 1
-
- # Contains at least one padding token in the sequence
- if attention_mask is not None:
- batch_size = query_states.shape[0]
- query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
- query_states, key_states, value_states, attention_mask, query_length
- )
-
- cu_seqlens_q, cu_seqlens_k = cu_seq_lens
- max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
- attn_output_unpad = flash_attn_varlen_func(
- query_states,
- key_states,
- value_states,
- cu_seqlens_q=cu_seqlens_q,
- cu_seqlens_k=cu_seqlens_k,
- max_seqlen_q=max_seqlen_in_batch_q,
- max_seqlen_k=max_seqlen_in_batch_k,
- dropout_p=dropout,
- softmax_scale=softmax_scale,
- causal=causal,
- )
-
- attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
- else:
- attn_output = flash_attn_func(
- query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
- )
-
- return attn_output
-
- # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input with num_heads->num_attention_heads
- def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
- indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
- batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
- key_layer = index_first_axis(
- key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- value_layer = index_first_axis(
- value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
- )
- if query_length == kv_seq_len:
- query_layer = index_first_axis(
- query_layer.reshape(batch_size * kv_seq_len, self.num_attention_heads, head_dim), indices_k
- )
- cu_seqlens_q = cu_seqlens_k
- max_seqlen_in_batch_q = max_seqlen_in_batch_k
- indices_q = indices_k
- elif query_length == 1:
- max_seqlen_in_batch_q = 1
- cu_seqlens_q = torch.arange(
- batch_size + 1, dtype=torch.int32, device=query_layer.device
- ) # There is a memcpy here, that is very bad.
- indices_q = cu_seqlens_q[:-1]
- query_layer = query_layer.squeeze(1)
- else:
- # The -q_len: slice assumes left padding.
- attention_mask = attention_mask[:, -query_length:]
- query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
- return (
- query_layer,
- key_layer,
- value_layer,
- indices_q,
- (cu_seqlens_q, cu_seqlens_k),
- (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
- )
-
-
-GPTJ_ATTENTION_CLASSES = {
- "eager": GPTJAttention,
- "flash_attention_2": GPTJFlashAttention2,
-}
-
-
-class GPTJMLP(nn.Module):
- def __init__(self, intermediate_size, config): # in MLP: intermediate_size= 4 * embed_dim
- super().__init__()
- embed_dim = config.n_embd
-
- self.fc_in = nn.Linear(embed_dim, intermediate_size)
- self.fc_out = nn.Linear(intermediate_size, embed_dim)
-
- self.act = ACT2FN[config.activation_function]
- self.dropout = nn.Dropout(config.resid_pdrop)
-
- def forward(self, hidden_states: Optional[torch.FloatTensor]) -> torch.FloatTensor:
- hidden_states = self.fc_in(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.fc_out(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
-
-class GPTJBlock(nn.Module):
- def __init__(self, config):
- super().__init__()
- inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd
- self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
- self.attn = GPTJ_ATTENTION_CLASSES[config._attn_implementation](config)
- self.mlp = GPTJMLP(inner_dim, config)
-
- def forward(
- self,
- hidden_states: Optional[torch.FloatTensor],
- layer_past: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states=hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
- outputs = attn_outputs[1:]
-
- feed_forward_hidden_states = self.mlp(hidden_states)
- hidden_states = attn_output + feed_forward_hidden_states + residual
-
- if use_cache:
- outputs = (hidden_states,) + outputs
- else:
- outputs = (hidden_states,) + outputs[1:]
-
- return outputs # hidden_states, present, (attentions)
-
-
-class GPTJPreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTJConfig
- base_model_prefix = "transformer"
- is_parallelizable = True
- supports_gradient_checkpointing = True
- _no_split_modules = ["GPTJBlock"]
- _skip_keys_device_placement = "past_key_values"
- _supports_flash_attn_2 = True
-
- def __init__(self, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- def _init_weights(self, module):
- """Initialize the weights."""
- if isinstance(module, (nn.Linear,)):
- # Slightly different from Mesh Transformer JAX which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
- elif isinstance(module, nn.LayerNorm):
- module.bias.data.zero_()
- module.weight.data.fill_(1.0)
-
-
-GPTJ_START_DOCSTRING = r"""
- This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
- it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
- behavior.
-
- Parameters:
- config ([`GPTJConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPTJ_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `({0})`):
- Indices of input sequence tokens in the vocabulary.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
- [`PreTrainedTokenizer.__call__`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.n_positions - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`torch.FloatTensor` of shape `(num_attention_heads,)` or `(n_layer, num_attention_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_dim)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-"""
-
-PARALLELIZE_DOCSTRING = r"""
- This is an experimental feature and is a subject to change at a moment's notice. Uses a device map to distribute
- attention modules of the model across several devices. If no device map is given, it will evenly distribute blocks
- across all devices.
-
- Args:
- device_map (`Dict[int, list]`, optional, defaults to None):
- A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always
- automatically mapped to the first device (for esoteric reasons). That means that the first device should
- have fewer attention modules mapped to it than other devices. For reference, the GPT-J models have the
- following number of attention modules:
-
- - gpt-j-6B: 28
-
- Example:
-
- ```python
- # Here is an example of a device map on a machine with 4 GPUs using gpt-j-6B, which has a total of 28 attention modules:
- model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
- device_map = {
- 0: [0, 1, 2, 3, 4, 5, 6],
- 1: [7, 8, 9, 10, 11, 12, 13],
- 2: [14, 15, 16, 17, 18, 19, 20],
- 3: [21, 22, 23, 24, 25, 26, 27],
- }
- model.parallelize(device_map)
- ```
-"""
-
-DEPARALLELIZE_DOCSTRING = r"""
- Moves the model to CPU from a model parallel state.
-
- Example:
-
- ```python
- # On a 4 GPU machine with gpt-j-6B:
- model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
- device_map = {
- 0: [0, 1, 2, 3, 4, 5, 6],
- 1: [7, 8, 9, 10, 11, 12, 13],
- 2: [14, 15, 16, 17, 18, 19, 20],
- 3: [21, 22, 23, 24, 25, 26, 27],
- }
- model.parallelize(device_map) # Splits the model across several devices
- model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache()
- ```
-"""
-
-
-@add_start_docstrings(
- "The bare GPT-J Model transformer outputting raw hidden-states without any specific head on top.",
- GPTJ_START_DOCSTRING,
-)
-class GPTJModel(GPTJPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
-
- self.embed_dim = config.n_embd
- self.vocab_size = config.vocab_size
- self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
- self.drop = nn.Dropout(config.embd_pdrop)
- self.h = nn.ModuleList([GPTJBlock(config) for _ in range(config.n_layer)])
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
- self.gradient_checkpointing = False
-
- # Initialize weights and apply final processing
- self.post_init()
-
- self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
-
- @add_start_docstrings(PARALLELIZE_DOCSTRING)
- def parallelize(self, device_map=None):
- warnings.warn(
- "`GPTJModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your"
- " model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own"
- " `device_map` but it needs to be a dictionary module_name to device, so for instance {'h.0': 0, 'h.1': 1,"
- " ...}",
- FutureWarning,
- )
- # Check validity of device_map
- self.device_map = (
- get_device_map(len(self.h), range(torch.cuda.device_count())) if device_map is None else device_map
- )
- assert_device_map(self.device_map, len(self.h))
- self.model_parallel = True
- self.first_device = "cpu" if "cpu" in self.device_map.keys() else "cuda:" + str(min(self.device_map.keys()))
- self.last_device = "cuda:" + str(max(self.device_map.keys()))
- self.wte = self.wte.to(self.first_device)
- # Load onto devices
- for k, v in self.device_map.items():
- for block in v:
- cuda_device = "cuda:" + str(k)
- self.h[block] = self.h[block].to(cuda_device)
- # ln_f to last
- self.ln_f = self.ln_f.to(self.last_device)
-
- @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
- def deparallelize(self):
- warnings.warn(
- "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.",
- FutureWarning,
- )
- self.model_parallel = False
- self.device_map = None
- self.first_device = "cpu"
- self.last_device = "cpu"
- self.wte = self.wte.to("cpu")
- for index in range(len(self.h)):
- self.h[index] = self.h[index].to("cpu")
- self.ln_f = self.ln_f.to("cpu")
- torch.cuda.empty_cache()
-
- def get_input_embeddings(self):
- return self.wte
-
- def set_input_embeddings(self, new_embeddings):
- self.wte = new_embeddings
-
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=BaseModelOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, BaseModelOutputWithPast]:
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
- output_hidden_states = (
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
- input_shape = input_ids.size()
- input_ids = input_ids.view(-1, input_shape[-1])
- batch_size = input_ids.shape[0]
- elif inputs_embeds is not None:
- input_shape = inputs_embeds.size()[:-1]
- batch_size = inputs_embeds.shape[0]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- device = input_ids.device if input_ids is not None else inputs_embeds.device
-
- if token_type_ids is not None:
- token_type_ids = token_type_ids.view(-1, input_shape[-1])
-
- if past_key_values is None:
- past_length = 0
- past_key_values = tuple([None] * len(self.h))
- else:
- past_length = past_key_values[0][0].size(-2)
-
- if position_ids is None:
- position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
- position_ids = position_ids.unsqueeze(0)
-
- if not self._use_flash_attention_2:
- # Attention mask.
- if attention_mask is not None:
- if batch_size <= 0:
- raise ValueError("batch_size has to be defined and > 0")
- attention_mask = attention_mask.view(batch_size, -1)
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask = attention_mask[:, None, None, :]
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and the dtype's smallest value for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
- attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x num_attention_heads x N x N
- # head_mask has shape n_layer x batch x num_attention_heads x N x N
- head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-
- if inputs_embeds is None:
- inputs_embeds = self.wte(input_ids)
-
- hidden_states = inputs_embeds
-
- if token_type_ids is not None:
- token_type_embeds = self.wte(token_type_ids)
- hidden_states = hidden_states + token_type_embeds
-
- hidden_states = self.drop(hidden_states)
-
- output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- presents = () if use_cache else None
- all_self_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- # Model parallel
- if self.model_parallel:
- torch.cuda.set_device(hidden_states.device)
- # Ensure layer_past is on same device as hidden_states (might not be correct)
- if layer_past is not None:
- layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past)
- # Ensure that attention_mask is always on the same device as hidden_states
- if attention_mask is not None:
- attention_mask = attention_mask.to(hidden_states.device)
- if isinstance(head_mask, torch.Tensor):
- head_mask = head_mask.to(hidden_states.device)
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if self.gradient_checkpointing and self.training:
- outputs = self._gradient_checkpointing_func(
- block.__call__,
- hidden_states,
- None,
- attention_mask,
- position_ids,
- head_mask[i],
- use_cache,
- output_attentions,
- )
- else:
- outputs = block(
- hidden_states=hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask[i],
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
-
- hidden_states = outputs[0]
- if use_cache is True:
- presents = presents + (outputs[1],)
-
- if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
-
- # Model Parallel: If it's the last layer for that device, put things on the next device
- if self.model_parallel:
- for k, v in self.device_map.items():
- if i == v[-1] and "cuda:" + str(k) != self.last_device:
- hidden_states = hidden_states.to("cuda:" + str(k + 1))
-
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = hidden_states.view(output_shape)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
-
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_self_attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT-J Model transformer with a language modeling head on top.
- """,
- GPTJ_START_DOCSTRING,
-)
-class GPTJForCausalLM(GPTJPreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config):
- super().__init__(config)
- self.transformer = GPTJModel(config)
- self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings(PARALLELIZE_DOCSTRING)
- def parallelize(self, device_map=None):
- warnings.warn(
- "`GPTJForCausalLM.parallelize` is deprecated and will be removed in v5 of Transformers, you should load"
- " your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own"
- " `device_map` but it needs to be a dictionary module_name to device, so for instance {'transformer.h.0':"
- " 0, 'transformer.h.1': 1, ...}",
- FutureWarning,
- )
- self.device_map = (
- get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
- if device_map is None
- else device_map
- )
- assert_device_map(self.device_map, len(self.transformer.h))
- self.transformer.parallelize(self.device_map)
- self.lm_head = self.lm_head.to(self.transformer.first_device)
- self.model_parallel = True
-
- @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
- def deparallelize(self):
- warnings.warn(
- "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.",
- FutureWarning,
- )
- self.transformer.deparallelize()
- self.transformer = self.transformer.to("cpu")
- self.lm_head = self.lm_head.to("cpu")
- self.model_parallel = False
- torch.cuda.empty_cache()
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # Omit tokens covered by past_key_values
- if past_key_values:
- past_length = past_key_values[0][0].shape[2]
-
- # Some generation methods already pass only the last input ID
- if input_ids.shape[1] > past_length:
- remove_prefix_length = past_length
- else:
- # Default to old behavior: keep only final ID
- remove_prefix_length = input_ids.shape[1] - 1
-
- input_ids = input_ids[:, remove_prefix_length:]
- if token_type_ids is not None:
- token_type_ids = token_type_ids[:, -input_ids.shape[1] :]
-
- attention_mask = kwargs.get("attention_mask", None)
- position_ids = kwargs.get("position_ids", None)
-
- if attention_mask is not None and position_ids is None:
- # create position_ids on the fly for batch generation
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -input_ids.shape[1] :]
-
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
- if inputs_embeds is not None and past_key_values is None:
- model_inputs = {"inputs_embeds": inputs_embeds}
- else:
- model_inputs = {"input_ids": input_ids}
-
- model_inputs.update(
- {
- "past_key_values": past_key_values,
- "use_cache": kwargs.get("use_cache"),
- "position_ids": position_ids,
- "attention_mask": attention_mask,
- "token_type_ids": token_type_ids,
- }
- )
-
- return model_inputs
-
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=CausalLMOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
-
- # Set device for model parallelism
- if self.model_parallel:
- torch.cuda.set_device(self.transformer.first_device)
- hidden_states = hidden_states.to(self.lm_head.weight.device)
-
- # make sure sampling in fp16 works correctly and
- # compute loss in fp32 to match with mesh-tf version
- # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179
- lm_logits = self.lm_head(hidden_states).to(torch.float32)
-
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
- loss = loss.to(hidden_states.dtype)
-
- if not return_dict:
- output = (lm_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=lm_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- @staticmethod
- def _reorder_cache(
- past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
- ) -> Tuple[Tuple[torch.Tensor]]:
- """
- This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
- [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
- beam_idx at every generation step.
- """
- return tuple(
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
- for layer_past in past_key_values
- )
-
-
-@add_start_docstrings(
- """
- The GPT-J Model transformer with a sequence classification head on top (linear layer).
-
- [`GPTJForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT, GPT-2, GPT-Neo) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- GPTJ_START_DOCSTRING,
-)
-class GPTJForSequenceClassification(GPTJPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = GPTJModel(config)
- self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint="ydshieh/tiny-random-gptj-for-sequence-classification",
- output_type=SequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- transformer_outputs = self.transformer(
- input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
-
- if input_ids is not None:
- batch_size = input_ids.shape[0]
- else:
- batch_size = inputs_embeds.shape[0]
-
- if self.config.pad_token_id is None and batch_size != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
- sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
- sequence_lengths = sequence_lengths % input_ids.shape[-1]
- sequence_lengths = sequence_lengths.to(logits.device)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
-
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
- loss = None
- if labels is not None:
- labels = labels.to(pooled_logits.device)
- if self.config.problem_type is None:
- if self.num_labels == 1:
- self.config.problem_type = "regression"
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
- self.config.problem_type = "single_label_classification"
- else:
- self.config.problem_type = "multi_label_classification"
-
- if self.config.problem_type == "regression":
- loss_fct = MSELoss()
- if self.num_labels == 1:
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
- else:
- loss = loss_fct(pooled_logits, labels)
- elif self.config.problem_type == "single_label_classification":
- loss_fct = CrossEntropyLoss()
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
- elif self.config.problem_type == "multi_label_classification":
- loss_fct = BCEWithLogitsLoss()
- loss = loss_fct(pooled_logits, labels)
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return SequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
-
-@add_start_docstrings(
- """
- The GPT-J Model transformer with a span classification head on top for extractive question-answering tasks like
- SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- GPTJ_START_DOCSTRING,
-)
-class GPTJForQuestionAnswering(GPTJPreTrainedModel):
- def __init__(self, config):
- super().__init__(config)
- self.num_labels = config.num_labels
- self.transformer = GPTJModel(config)
- self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
-
- # Model parallel
- self.model_parallel = False
- self.device_map = None
-
- # Initialize weights and apply final processing
- self.post_init()
-
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=QuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
- )
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- start_positions: Optional[torch.LongTensor] = None,
- end_positions: Optional[torch.LongTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- ) -> Union[Tuple, QuestionAnsweringModelOutput]:
- r"""
- start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
- outputs = self.transformer(
- input_ids,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- )
-
- sequence_output = outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = logits.split(1, dim=-1)
- start_logits = start_logits.squeeze(-1).contiguous()
- end_logits = end_logits.squeeze(-1).contiguous()
-
- total_loss = None
- if start_positions is not None and end_positions is not None:
- # If we are on multi-GPU, split add a dimension
- if len(start_positions.size()) > 1:
- start_positions = start_positions.squeeze(-1).to(start_logits.device)
- if len(end_positions.size()) > 1:
- end_positions = end_positions.squeeze(-1).to(end_logits.device)
- # sometimes the start/end positions are outside our model inputs, we ignore these terms
- ignored_index = start_logits.size(1)
- start_positions = start_positions.clamp(0, ignored_index)
- end_positions = end_positions.clamp(0, ignored_index)
-
- loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
- start_loss = loss_fct(start_logits, start_positions)
- end_loss = loss_fct(end_logits, end_positions)
- total_loss = (start_loss + end_loss) / 2
-
- if not return_dict:
- output = (start_logits, end_logits) + outputs[2:]
- return ((total_loss,) + output) if total_loss is not None else output
-
- return QuestionAnsweringModelOutput(
- loss=total_loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
diff --git a/transformers/models/gptj/modeling_tf_gptj.py b/transformers/models/gptj/modeling_tf_gptj.py
deleted file mode 100644
index 5c315b5b66f04929a48f809de5ec0994d40561a7..0000000000000000000000000000000000000000
--- a/transformers/models/gptj/modeling_tf_gptj.py
+++ /dev/null
@@ -1,1099 +0,0 @@
-# coding=utf-8
-# Copyright 2022 The EleutherAI and HuggingFace Teams. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" TF 2.0 GPT-J model."""
-
-from __future__ import annotations
-
-from typing import Optional, Tuple, Union
-
-import numpy as np
-import tensorflow as tf
-
-from ...activations_tf import get_tf_activation
-from ...file_utils import (
- add_code_sample_docstrings,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
-)
-from ...modeling_tf_outputs import (
- TFBaseModelOutputWithPast,
- TFCausalLMOutputWithPast,
- TFQuestionAnsweringModelOutput,
- TFSequenceClassifierOutputWithPast,
-)
-from ...modeling_tf_utils import (
- TFCausalLanguageModelingLoss,
- TFModelInputType,
- TFPreTrainedModel,
- TFQuestionAnsweringLoss,
- TFSequenceClassificationLoss,
- TFSharedEmbeddings,
- get_initializer,
- keras,
- keras_serializable,
- unpack_inputs,
-)
-from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
-from ...utils import logging
-from .configuration_gptj import GPTJConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CHECKPOINT_FOR_DOC = "EleutherAI/gpt-j-6B"
-_CONFIG_FOR_DOC = "GPTJConfig"
-
-
-def create_sinusoidal_positions(num_pos: int, dim: int) -> tf.Tensor:
- inv_freq = tf.cast(1.0 / (10000 ** (tf.range(0, dim, 2) / dim)), tf.float32)
- sinusoid_inp = tf.cast(tf.einsum("i , j -> i j", tf.range(num_pos, dtype=tf.float32), inv_freq), tf.float32)
- sin, cos = tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)
- out = tf.concat((sin, cos), axis=1)
- return out
-
-
-def rotate_every_two(x: tf.Tensor) -> tf.Tensor:
- rotate_half_tensor = tf.stack((-x[:, :, :, 1::2], x[:, :, :, ::2]), axis=-1)
- new_shape = shape_list(rotate_half_tensor)[:-2] + [tf.math.reduce_prod(shape_list(rotate_half_tensor)[-2:])]
- rotate_half_tensor = tf.reshape(rotate_half_tensor, new_shape)
- return rotate_half_tensor
-
-
-def apply_rotary_pos_emb(tensor: tf.Tensor, sincos: tf.Tensor) -> tf.Tensor:
- sin_pos, cos_pos = sincos
- sin_pos = tf.repeat(sin_pos[:, :, None, :], 2, 3)
- cos_pos = tf.repeat(cos_pos[:, :, None, :], 2, 3)
- return (tensor * cos_pos) + (rotate_every_two(tensor) * sin_pos)
-
-
-class TFGPTJAttention(keras.layers.Layer):
- def __init__(self, config: GPTJConfig, **kwargs):
- super().__init__(**kwargs)
-
- self.embed_dim = config.hidden_size
- self.num_attention_heads = config.num_attention_heads
- self.head_dim = self.embed_dim // self.num_attention_heads
- if self.head_dim * self.num_attention_heads != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and"
- f" `num_attention_heads`: {self.num_attention_heads})."
- )
- self.scale_attn = self.head_dim**0.5
- self.rotary_dim = config.rotary_dim
-
- self.attn_dropout = keras.layers.Dropout(config.attn_pdrop)
- self.resid_dropout = keras.layers.Dropout(config.resid_pdrop)
-
- self.q_proj = keras.layers.Dense(
- self.embed_dim,
- use_bias=False,
- kernel_initializer=get_initializer(config.initializer_range),
- name="q_proj",
- )
- self.k_proj = keras.layers.Dense(
- self.embed_dim,
- use_bias=False,
- kernel_initializer=get_initializer(config.initializer_range),
- name="k_proj",
- )
- self.v_proj = keras.layers.Dense(
- self.embed_dim,
- use_bias=False,
- kernel_initializer=get_initializer(config.initializer_range),
- name="v_proj",
- )
- self.out_proj = keras.layers.Dense(
- self.embed_dim,
- use_bias=False,
- kernel_initializer=get_initializer(config.initializer_range),
- name="out_proj",
- )
-
- self.max_positions = config.max_position_embeddings
- self.lower_triangle_mask = tf.reshape(
- tf.cast(tf.experimental.numpy.tril(tf.ones((self.max_positions, self.max_positions))), tf.int8),
- (1, 1, self.max_positions, self.max_positions),
- )
- pos_embd_dim = self.rotary_dim or self.embed_dim
- self.embed_positions = create_sinusoidal_positions(self.max_positions, pos_embd_dim)
-
- def get_causal_mask(self, key_length, query_length) -> tf.Tensor:
- return tf.cast(self.lower_triangle_mask[:, :, key_length - query_length : key_length, :key_length], tf.bool)
-
- @staticmethod
- def get_masked_bias(dtype: tf.DType) -> tf.Tensor:
- return tf.cast(tf.constant(-1e9), dtype)
-
- def _split_heads(self, hidden_states: tf.Tensor, rotary: bool) -> tf.Tensor:
- """
- Splits hidden dim into attn_head_size and num_attention_heads
- """
- new_shape = shape_list(hidden_states)[:-1] + [self.num_attention_heads, self.head_dim]
- hidden_states = tf.reshape(hidden_states, new_shape)
- if rotary:
- return hidden_states
- if len(shape_list(hidden_states)) == 4:
- return tf.transpose(hidden_states, (0, 2, 1, 3)) # (batch, head, seq_length, head_features)
- if len(shape_list(hidden_states)) == 5:
- return tf.transpose(hidden_states, (0, 1, 3, 2, 4)) # (batch, blocks, head, block_length, head_features)
- raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(shape_list(hidden_states))}")
-
- def _merge_heads(self, hidden_states: tf.Tensor) -> tf.Tensor:
- """
- Merges attn_head_size dim and num_attn_heads dim into hidden dim
- """
- if len(shape_list(hidden_states)) == 4:
- hidden_states = tf.transpose(hidden_states, (0, 2, 1, 3))
- elif len(shape_list(hidden_states)) == 5:
- hidden_states = tf.transpose(hidden_states, (0, 1, 3, 2, 4))
- else:
- raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(shape_list(hidden_states))}")
- new_shape = shape_list(hidden_states)[:-2] + [self.num_attention_heads * self.head_dim]
- return tf.reshape(hidden_states, new_shape)
-
- def _attn(
- self,
- query: tf.Tensor,
- key: tf.Tensor,
- value: tf.Tensor,
- attention_mask: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- ) -> Tuple[tf.Tensor, tf.Tensor]:
- # compute causal mask from causal mask buffer
- query_length, key_length = shape_list(query)[-2], shape_list(key)[-2]
- causal_mask = self.get_causal_mask(key_length, query_length)
-
- # Keep the attention weights computation in fp32 to avoid overflow issues
- query = tf.cast(query, tf.float32)
- key = tf.cast(key, tf.float32)
-
- attn_weights = tf.matmul(query, key, transpose_b=True)
- attn_weights = tf.where(causal_mask, attn_weights, self.get_masked_bias(attn_weights.dtype))
-
- attn_weights = attn_weights / self.scale_attn
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = stable_softmax(attn_weights, axis=-1)
- attn_weights = tf.cast(attn_weights, value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- # Mask heads if we want to
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = tf.matmul(attn_weights, value)
-
- return attn_output, attn_weights
-
- def call(
- self,
- hidden_states: tf.Tensor,
- layer_past: Optional[Tuple[tf.Tensor, tf.Tensor]] = None,
- attention_mask: tf.Tensor | None = None,
- position_ids: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- use_cache: bool = False,
- output_attentions: bool = False,
- ):
- query = self.q_proj(hidden_states)
- key = self.k_proj(hidden_states)
- value = self.v_proj(hidden_states)
-
- query = self._split_heads(query, True)
- key = self._split_heads(key, True)
- value = self._split_heads(value, False)
-
- sincos = tf.cast(tf.gather(self.embed_positions, position_ids, axis=0), hidden_states.dtype)
- sincos = tf.split(sincos, 2, axis=-1)
- if self.rotary_dim is not None:
- k_rot = key[:, :, :, : self.rotary_dim]
- k_pass = key[:, :, :, self.rotary_dim :]
-
- q_rot = query[:, :, :, : self.rotary_dim]
- q_pass = query[:, :, :, self.rotary_dim :]
-
- k_rot = apply_rotary_pos_emb(k_rot, sincos)
- q_rot = apply_rotary_pos_emb(q_rot, sincos)
-
- key = tf.concat((k_rot, k_pass), axis=-1)
- query = tf.concat((q_rot, q_pass), axis=-1)
- else:
- key = apply_rotary_pos_emb(key, sincos)
- query = apply_rotary_pos_emb(query, sincos)
-
- key = tf.transpose(key, (0, 2, 1, 3))
- query = tf.transpose(query, (0, 2, 1, 3))
-
- if layer_past is not None:
- past_key = layer_past[0]
- past_value = layer_past[1]
- key = tf.concat((past_key, key), axis=-2)
- value = tf.concat((past_value, value), axis=-2)
-
- if use_cache is True:
- present = (key, value)
- else:
- present = None
-
- # compute self-attention: V x Softmax(QK^T)
- attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
-
- attn_output = self._merge_heads(attn_output)
- attn_output = self.out_proj(attn_output)
- attn_output = self.resid_dropout(attn_output)
-
- outputs = (attn_output, present)
- if output_attentions:
- outputs += (attn_weights,)
-
- return outputs # a, present, (attentions)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "q_proj", None) is not None:
- with tf.name_scope(self.q_proj.name):
- self.q_proj.build([None, None, self.embed_dim])
- if getattr(self, "k_proj", None) is not None:
- with tf.name_scope(self.k_proj.name):
- self.k_proj.build([None, None, self.embed_dim])
- if getattr(self, "v_proj", None) is not None:
- with tf.name_scope(self.v_proj.name):
- self.v_proj.build([None, None, self.embed_dim])
- if getattr(self, "out_proj", None) is not None:
- with tf.name_scope(self.out_proj.name):
- self.out_proj.build([None, None, self.embed_dim])
-
-
-class TFGPTJMLP(keras.layers.Layer):
- def __init__(self, intermediate_size: int, config: GPTJConfig, **kwargs):
- super().__init__(**kwargs)
- embed_dim = config.n_embd
-
- self.fc_in = keras.layers.Dense(
- intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="fc_in"
- )
- self.fc_out = keras.layers.Dense(
- embed_dim, kernel_initializer=get_initializer(config.initializer_range), name="fc_out"
- )
-
- self.act = get_tf_activation(config.activation_function)
- self.dropout = keras.layers.Dropout(config.embd_pdrop)
- self.embed_dim = config.n_embd
- self.intermediate_size = intermediate_size
-
- def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
- hidden_states = self.fc_in(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.fc_out(hidden_states)
- hidden_states = self.dropout(hidden_states)
- return hidden_states
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "fc_in", None) is not None:
- with tf.name_scope(self.fc_in.name):
- self.fc_in.build([None, None, self.embed_dim])
- if getattr(self, "fc_out", None) is not None:
- with tf.name_scope(self.fc_out.name):
- self.fc_out.build([None, None, self.intermediate_size])
-
-
-class TFGPTJBlock(keras.layers.Layer):
- def __init__(self, config: GPTJConfig, **kwargs):
- super().__init__(**kwargs)
- inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd
- self.ln_1 = keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="ln_1")
- self.attn = TFGPTJAttention(config, name="attn")
- self.mlp = TFGPTJMLP(inner_dim, config, name="mlp")
- self.config = config
-
- def call(
- self,
- hidden_states: tf.Tensor,
- layer_past: tf.Tensor | None = None,
- attention_mask: tf.Tensor | None = None,
- position_ids: tf.Tensor | None = None,
- head_mask: tf.Tensor | None = None,
- use_cache: bool = False,
- output_attentions: bool = False,
- ):
- residual = hidden_states
- hidden_states = self.ln_1(hidden_states)
- attn_outputs = self.attn(
- hidden_states=hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- ) # attn_outputs: attn_output, present, (attentions)
- attn_output = attn_outputs[0]
- outputs = attn_outputs[1:]
-
- feed_forward_hidden_states = self.mlp(hidden_states)
- hidden_states = attn_output + feed_forward_hidden_states + residual
-
- if use_cache:
- outputs = (hidden_states,) + outputs
- else:
- outputs = (hidden_states,) + outputs[1:]
- return outputs # hidden_states, present, (attentions)
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "ln_1", None) is not None:
- with tf.name_scope(self.ln_1.name):
- self.ln_1.build([None, None, self.config.n_embd])
- if getattr(self, "attn", None) is not None:
- with tf.name_scope(self.attn.name):
- self.attn.build(None)
- if getattr(self, "mlp", None) is not None:
- with tf.name_scope(self.mlp.name):
- self.mlp.build(None)
-
-
-@keras_serializable
-class TFGPTJMainLayer(keras.layers.Layer):
- config_class = GPTJConfig
-
- def __init__(self, config: GPTJConfig, *inputs, **kwargs):
- super().__init__(*inputs, **kwargs)
-
- self.config = config
- self.output_attentions = config.output_attentions
- self.output_hidden_states = config.output_hidden_states
- self.use_cache = config.use_cache
- self.return_dict = config.use_return_dict
-
- self.num_hidden_layers = config.n_layer
- self.n_embd = config.n_embd
- self.n_positions = config.n_positions
- self.initializer_range = config.initializer_range
-
- self.wte = TFSharedEmbeddings(
- config.vocab_size, config.hidden_size, initializer_range=config.initializer_range, name="wte"
- )
- self.drop = keras.layers.Dropout(config.embd_pdrop)
- self.h = [TFGPTJBlock(config, name=f"h_._{i}") for i in range(config.n_layer)]
- self.ln_f = keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="ln_f")
- self.embed_dim = config.n_embd
-
- def get_input_embeddings(self):
- return self.wte
-
- def set_input_embeddings(self, value: tf.Tensor):
- self.wte.weight = value
- self.wte.vocab_size = shape_list(value)[0]
-
- def _prune_heads(self, heads_to_prune):
- """
- Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
- """
- raise NotImplementedError
-
- @unpack_inputs
- def call(
- self,
- input_ids=None,
- past_key_values=None,
- attention_mask=None,
- token_type_ids=None,
- position_ids=None,
- head_mask=None,
- inputs_embeds=None,
- use_cache=None,
- output_attentions=None,
- output_hidden_states=None,
- return_dict=None,
- training=False,
- ) -> Union[TFBaseModelOutputWithPast, Tuple[tf.Tensor]]:
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
- elif input_ids is not None:
- input_shape = shape_list(input_ids)
- input_ids = tf.reshape(input_ids, [-1, input_shape[-1]])
- elif inputs_embeds is not None:
- input_shape = shape_list(inputs_embeds)[:-1]
- else:
- raise ValueError("You have to specify either input_ids or inputs_embeds")
-
- if past_key_values is None:
- past_length = 0
- past_key_values = [None] * len(self.h)
- else:
- past_length = shape_list(past_key_values[0][0])[-2]
-
- if position_ids is None:
- position_ids = tf.expand_dims(tf.range(past_length, input_shape[-1] + past_length), axis=0)
-
- if attention_mask is not None:
- # We create a 3D attention mask from a 2D tensor mask.
- # Sizes are [batch_size, 1, 1, to_seq_length]
- # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
- # this attention mask is more simple than the triangular masking of causal attention
- # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
- attention_mask_shape = shape_list(attention_mask)
- attention_mask = tf.reshape(attention_mask, (attention_mask_shape[0], 1, 1, attention_mask_shape[1]))
-
- # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
- # masked positions, this operation will create a tensor which is 0.0 for
- # positions we want to attend and -10000.0 for masked positions.
- # Since we are adding it to the raw scores before the softmax, this is
- # effectively the same as removing these entirely.
- one_cst = tf.constant(1.0)
- attention_mask = tf.cast(attention_mask, dtype=one_cst.dtype)
- attention_mask = tf.multiply(tf.subtract(one_cst, attention_mask), tf.constant(-10000.0))
-
- # Prepare head mask if needed
- # 1.0 in head_mask indicate we keep the head
- # attention_probs has shape bsz x n_heads x N x N
- # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
- # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
- if head_mask is not None:
- raise NotImplementedError
- else:
- head_mask = [None] * self.num_hidden_layers
- # head_mask = tf.constant([0] * self.num_hidden_layers)
-
- position_ids = tf.reshape(position_ids, [-1, shape_list(position_ids)[-1]])
-
- if inputs_embeds is None:
- check_embeddings_within_bounds(input_ids, self.wte.vocab_size)
- inputs_embeds = self.wte(input_ids, mode="embedding")
-
- if token_type_ids is not None:
- token_type_ids = tf.reshape(token_type_ids, [-1, shape_list(token_type_ids)[-1]])
- token_type_embeds = self.wte(token_type_ids, mode="embedding")
- else:
- token_type_embeds = tf.constant(0.0)
-
- token_type_embeds = tf.cast(token_type_embeds, dtype=inputs_embeds.dtype)
- hidden_states = inputs_embeds + token_type_embeds
- hidden_states = self.drop(hidden_states, training=training)
-
- output_shape = input_shape + [shape_list(hidden_states)[-1]]
-
- presents = () if use_cache else None
- all_attentions = () if output_attentions else None
- all_hidden_states = () if output_hidden_states else None
- for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (tf.reshape(hidden_states, output_shape),)
-
- outputs = block(
- hidden_states=hidden_states,
- layer_past=layer_past,
- attention_mask=attention_mask,
- position_ids=position_ids,
- head_mask=head_mask[i],
- use_cache=use_cache,
- output_attentions=output_attentions,
- training=training,
- )
-
- hidden_states = outputs[0]
- if use_cache:
- presents = presents + (outputs[1],)
-
- if output_attentions:
- all_attentions = all_attentions + (outputs[2 if use_cache else 1],)
-
- hidden_states = self.ln_f(hidden_states)
-
- hidden_states = tf.reshape(hidden_states, output_shape)
- # Add last hidden state
- if output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if output_attentions:
- # let the number of heads free (-1) so we can extract attention even after head pruning
- attention_output_shape = input_shape[:-1] + [-1] + shape_list(all_attentions[0])[-2:]
- all_attentions = tuple(tf.reshape(t, attention_output_shape) for t in all_attentions)
-
- if not return_dict:
- return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None)
-
- return TFBaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=presents,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "wte", None) is not None:
- with tf.name_scope(self.wte.name):
- self.wte.build(None)
- if getattr(self, "ln_f", None) is not None:
- with tf.name_scope(self.ln_f.name):
- self.ln_f.build([None, None, self.embed_dim])
- if getattr(self, "h", None) is not None:
- for layer in self.h:
- with tf.name_scope(layer.name):
- layer.build(None)
-
-
-class TFGPTJPreTrainedModel(TFPreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTJConfig
- base_model_prefix = "transformer"
- # names with a '.' represents the authorized unexpected/missing layers when a TF model is loaded from a PT model
- _keys_to_ignore_on_load_unexpected = [r"h.\d+.attn.bias"]
-
-
-GPTJ_START_DOCSTRING = r"""
-
- This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
- etc.)
-
- This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it
- as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and
- behavior.
-
-
-
- TensorFlow models and layers in `transformers` accept two formats as input:
-
- - having all inputs as keyword arguments (like PyTorch models), or
- - having all inputs as a list, tuple or dict in the first positional argument.
-
- The reason the second format is supported is that Keras methods prefer this format when passing inputs to models
- and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just
- pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second
- format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with
- the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first
- positional argument:
-
- - a single Tensor with `input_ids` only and nothing else: `model(input_ids)`
- - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
- `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- - a dictionary with one or several input Tensors associated to the input names given in the docstring:
- `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
-
- Note that when creating models and layers with
- [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry
- about any of this, as you can just pass inputs like you would to any other Python function!
-
-
-
- Parameters:
- config ([`GPTJConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPTJ_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`Numpy array` or `tf.Tensor` of shape `(batch_size, input_ids_length)`):
- `input_ids_length` = `sequence_length` if `past` is `None` else `past[0].shape[-2]` (`sequence_length` of
- input past key value states). Indices of input sequence tokens in the vocabulary.
-
- If `past` is used, only input IDs that do not have their past calculated should be passed as `input_ids`.
-
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.__call__`] and
- [`PreTrainedTokenizer.encode`] for details.
-
- [What are input IDs?](../glossary#input-ids)
- past_key_values (`List[tf.Tensor]` of length `config.n_layers`):
- Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model (see
- `past` output below). Can be used to speed up sequential decoding. The token ids which have their past
- given to this model should not be passed as input ids as they have already been computed.
- attention_mask (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
- 1]`:
-
- - 0 corresponds to a *sentence A* token,
- - 1 corresponds to a *sentence B* token.
-
- [What are token type IDs?](../glossary#token-type-ids)
- position_ids (`tf.Tensor` or `Numpy array` of shape `(batch_size, sequence_length)`, *optional*):
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
- config.max_position_embeddings - 1]`.
-
- [What are position IDs?](../glossary#position-ids)
- head_mask (`Numpy array` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- inputs_embeds (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
- config will be used instead.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
- used instead.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. This argument can be used
- in eager mode, in graph mode the value will always be set to True.
- training (`bool`, *optional*, defaults to `False`):
- Whether or not to use the model in training mode (some modules like dropout modules have different
- behaviors between training and evaluation).
-"""
-
-
-@add_start_docstrings(
- "The bare GPT-J Model transformer outputting raw hidden-states without any specific head on top.",
- GPTJ_START_DOCSTRING,
-)
-class TFGPTJModel(TFGPTJPreTrainedModel):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFGPTJMainLayer(config, name="transformer")
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING)
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFBaseModelOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFBaseModelOutputWithPast, Tuple[tf.Tensor]]:
- r"""
- use_cache (`bool`, *optional*, defaults to `True`):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past`). Set to `False` during training, `True` during generation
- """
-
- outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
-
- return outputs
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
-
-
-@add_start_docstrings(
- """
- The GPT-J Model transformer with a language modeling head on top.
- """,
- GPTJ_START_DOCSTRING,
-)
-class TFGPTJForCausalLM(TFGPTJPreTrainedModel, TFCausalLanguageModelingLoss):
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.transformer = TFGPTJMainLayer(config, name="transformer")
- self.lm_head = keras.layers.Dense(
- config.vocab_size, kernel_initializer=get_initializer(config.initializer_range), name="lm_head"
- )
- self.config = config
-
- def get_output_embeddings(self):
- return self.lm_head
-
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- def prepare_inputs_for_generation(self, inputs, past_key_values=None, use_cache=None, **kwargs):
- token_type_ids = kwargs.get("token_type_ids", None)
- # only last token for inputs_ids if past is defined in kwargs
- if past_key_values:
- inputs = tf.expand_dims(inputs[:, -1], -1)
- if token_type_ids is not None:
- token_type_ids = tf.expand_dims(token_type_ids[:, -1], -1)
-
- position_ids = kwargs.get("position_ids", None)
- attention_mask = kwargs.get("attention_mask", None)
-
- if attention_mask is not None and position_ids is None:
- position_ids = tf.math.cumsum(attention_mask, axis=-1, exclusive=True)
- if past_key_values:
- position_ids = tf.expand_dims(position_ids[:, -1], -1)
-
- return {
- "input_ids": inputs,
- "attention_mask": attention_mask,
- "position_ids": position_ids,
- "past_key_values": past_key_values,
- "use_cache": use_cache,
- "token_type_ids": token_type_ids,
- }
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFCausalLMOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- labels: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFCausalLMOutputWithPast, Tuple[tf.Tensor]]:
- r"""
- labels (`np.ndarray` or `tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
- `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
- are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
- """
-
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_states = transformer_outputs[0]
- lm_logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # shift labels to the left and cut last logit token
- shifted_logits = lm_logits[:, :-1]
- labels = labels[:, 1:]
- loss = self.hf_compute_loss(labels, shifted_logits)
-
- if not return_dict:
- output = (lm_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFCausalLMOutputWithPast(
- loss=loss,
- logits=lm_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "lm_head", None) is not None:
- with tf.name_scope(self.lm_head.name):
- self.lm_head.build([None, None, self.config.n_embd])
-
-
-@add_start_docstrings(
- """
- The GPT-J Model transformer with a sequence classification head on top (linear layer).
-
- [`GPTJForSequenceClassification`] uses the last token in order to do the classification, as other causal models
- (e.g. GPT, GPT-2, GPT-Neo) do.
-
- Since it does classification on the last token, it requires to know the position of the last token. If a
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
- each row of the batch).
- """,
- GPTJ_START_DOCSTRING,
-)
-class TFGPTJForSequenceClassification(TFGPTJPreTrainedModel, TFSequenceClassificationLoss):
- _keys_to_ignore_on_load_missing = [r"h.\d+.attn.masked_bias", r"h.\d+.attn.bias", r"lm_head.weight"]
-
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
- self.transformer = TFGPTJMainLayer(config, name="transformer")
- self.score = keras.layers.Dense(
- self.num_labels,
- use_bias=False,
- kernel_initializer=get_initializer(config.initializer_range),
- name="score",
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFSequenceClassifierOutputWithPast,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- labels: np.ndarray | tf.Tensor | None = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFSequenceClassifierOutputWithPast, Tuple[tf.Tensor]]:
- r"""
- labels (`np.ndarray` or `tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
- """
-
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- hidden_states = transformer_outputs[0]
- logits = self.score(hidden_states)
- logits_shape = shape_list(logits)
- in_logits = None
- if self.config.pad_token_id is None:
- sequence_lengths = -1
- else:
- if input_ids is not None:
- sequence_lengths = (
- tf.argmax(tf.cast(tf.math.equal(input_ids, self.config.pad_token_id), input_ids.dtype), axis=-1)
- - 1
- )
- sequence_lengths = tf.where(
- sequence_lengths >= 0,
- sequence_lengths,
- tf.cast(shape_list(input_ids[-1]), sequence_lengths.dtype) - 1,
- )
- in_logits = tf.gather(logits, sequence_lengths, batch_dims=1, axis=1)
- else:
- sequence_lengths = -1
- logger.warning(
- f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
- "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
- )
- loss = None
-
- if labels is not None:
- if self.config.pad_token_id is None and logits_shape[0] != 1:
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
-
- if not tf.is_tensor(sequence_lengths):
- in_logits = logits[0 : logits_shape[0], sequence_lengths]
-
- loss = self.hf_compute_loss(tf.reshape(labels, [-1]), tf.reshape(in_logits, [-1, self.num_labels]))
- pooled_logits = in_logits if in_logits is not None else logits
-
- if not return_dict:
- output = (pooled_logits,) + transformer_outputs[1:]
- return ((loss,) + output) if loss is not None else output
-
- return TFSequenceClassifierOutputWithPast(
- loss=loss,
- logits=pooled_logits,
- past_key_values=transformer_outputs.past_key_values,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "score", None) is not None:
- with tf.name_scope(self.score.name):
- self.score.build([None, None, self.config.n_embd])
-
-
-@add_start_docstrings(
- """
- The GPT-J Model transformer with a span classification head on top for extractive question-answering tasks like
- SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
- """,
- GPTJ_START_DOCSTRING,
-)
-class TFGPTJForQuestionAnswering(TFGPTJPreTrainedModel, TFQuestionAnsweringLoss):
- _keys_to_ignore_on_load_missing = [r"h.\d+.attn.masked_bias", r"h.\d+.attn.bias", r"lm_head.weight"]
-
- def __init__(self, config, *inputs, **kwargs):
- super().__init__(config, *inputs, **kwargs)
- self.num_labels = config.num_labels
- self.transformer = TFGPTJMainLayer(config, name="transformer")
- self.qa_outputs = keras.layers.Dense(
- self.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
- )
- self.config = config
-
- @unpack_inputs
- @add_start_docstrings_to_model_forward(GPTJ_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
- @add_code_sample_docstrings(
- checkpoint=_CHECKPOINT_FOR_DOC,
- output_type=TFQuestionAnsweringModelOutput,
- config_class=_CONFIG_FOR_DOC,
- )
- def call(
- self,
- input_ids: TFModelInputType | None = None,
- past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
- attention_mask: np.ndarray | tf.Tensor | None = None,
- token_type_ids: np.ndarray | tf.Tensor | None = None,
- position_ids: np.ndarray | tf.Tensor | None = None,
- head_mask: np.ndarray | tf.Tensor | None = None,
- inputs_embeds: np.ndarray | tf.Tensor | None = None,
- start_positions: np.ndarray | tf.Tensor | None = None,
- end_positions: np.ndarray | tf.Tensor | None = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- training: Optional[bool] = False,
- ) -> Union[TFQuestionAnsweringModelOutput, Tuple[tf.Tensor]]:
- r"""
- start_positions (`np.ndarray` or `tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the start of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- end_positions (`np.ndarray` or `tf.Tensor` of shape `(batch_size,)`, *optional*):
- Labels for position (index) of the end of the labelled span for computing the token classification loss.
- Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
- are not taken into account for computing the loss.
- """
-
- transformer_outputs = self.transformer(
- input_ids=input_ids,
- past_key_values=past_key_values,
- attention_mask=attention_mask,
- token_type_ids=token_type_ids,
- position_ids=position_ids,
- head_mask=head_mask,
- inputs_embeds=inputs_embeds,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- training=training,
- )
- sequence_output = transformer_outputs[0]
-
- logits = self.qa_outputs(sequence_output)
- start_logits, end_logits = tf.split(logits, 2, axis=-1)
- start_logits = tf.squeeze(start_logits, axis=-1)
- end_logits = tf.squeeze(end_logits, axis=-1)
-
- loss = None
- if start_positions is not None and end_positions is not None:
- labels = {"start_position": start_positions}
- labels["end_position"] = end_positions
- loss = self.hf_compute_loss(labels, (start_logits, end_logits))
-
- if not return_dict:
- output = (start_logits, end_logits) + transformer_outputs[2:]
- return ((loss,) + output) if loss is not None else output
-
- return TFQuestionAnsweringModelOutput(
- loss=loss,
- start_logits=start_logits,
- end_logits=end_logits,
- hidden_states=transformer_outputs.hidden_states,
- attentions=transformer_outputs.attentions,
- )
-
- def build(self, input_shape=None):
- if self.built:
- return
- self.built = True
- if getattr(self, "transformer", None) is not None:
- with tf.name_scope(self.transformer.name):
- self.transformer.build(None)
- if getattr(self, "qa_outputs", None) is not None:
- with tf.name_scope(self.qa_outputs.name):
- self.qa_outputs.build([None, None, self.config.hidden_size])
diff --git a/transformers/models/gptsan_japanese/__init__.py b/transformers/models/gptsan_japanese/__init__.py
deleted file mode 100644
index b3635ace91163577201f716c9d67e255f11ea55b..0000000000000000000000000000000000000000
--- a/transformers/models/gptsan_japanese/__init__.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# Copyright 2022 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import TYPE_CHECKING
-
-from ...utils import (
- OptionalDependencyNotAvailable,
- _LazyModule,
- is_flax_available,
- is_tf_available,
- is_torch_available,
-)
-
-
-_import_structure = {
- "configuration_gptsan_japanese": ["GPTSAN_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTSanJapaneseConfig"],
- "tokenization_gptsan_japanese": ["GPTSanJapaneseTokenizer"],
-}
-
-try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
- pass
-else:
- _import_structure["modeling_gptsan_japanese"] = [
- "GPTSAN_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST",
- "GPTSanJapaneseForConditionalGeneration",
- "GPTSanJapaneseModel",
- "GPTSanJapanesePreTrainedModel",
- ]
- _import_structure["tokenization_gptsan_japanese"] = [
- "GPTSanJapaneseTokenizer",
- ]
-
-
-if TYPE_CHECKING:
- from .configuration_gptsan_japanese import GPTSAN_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTSanJapaneseConfig
- from .tokenization_gptsan_japanese import GPTSanJapaneseTokenizer
-
- try:
- if not is_torch_available():
- raise OptionalDependencyNotAvailable()
- except OptionalDependencyNotAvailable:
- pass
- else:
- from .modeling_gptsan_japanese import (
- GPTSAN_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST,
- GPTSanJapaneseForConditionalGeneration,
- GPTSanJapaneseModel,
- GPTSanJapanesePreTrainedModel,
- )
- from .tokenization_gptsan_japanese import GPTSanJapaneseTokenizer
-
-
-else:
- import sys
-
- sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/transformers/models/gptsan_japanese/__pycache__/__init__.cpython-310.pyc b/transformers/models/gptsan_japanese/__pycache__/__init__.cpython-310.pyc
deleted file mode 100644
index b3888844448a9fbc782a643e4a07e8332cf51454..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptsan_japanese/__pycache__/__init__.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptsan_japanese/__pycache__/configuration_gptsan_japanese.cpython-310.pyc b/transformers/models/gptsan_japanese/__pycache__/configuration_gptsan_japanese.cpython-310.pyc
deleted file mode 100644
index 6df3bbd292a1bb5ea64bf852df1af098b79b395b..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptsan_japanese/__pycache__/configuration_gptsan_japanese.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptsan_japanese/__pycache__/convert_gptsan_tf_checkpoint_to_pytorch.cpython-310.pyc b/transformers/models/gptsan_japanese/__pycache__/convert_gptsan_tf_checkpoint_to_pytorch.cpython-310.pyc
deleted file mode 100644
index 779b316776b9f7143aa40cd5aa02541bc2fbacf6..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptsan_japanese/__pycache__/convert_gptsan_tf_checkpoint_to_pytorch.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptsan_japanese/__pycache__/modeling_gptsan_japanese.cpython-310.pyc b/transformers/models/gptsan_japanese/__pycache__/modeling_gptsan_japanese.cpython-310.pyc
deleted file mode 100644
index 9abd43d3bdce4315f5f65308d2b1b879648b49d9..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptsan_japanese/__pycache__/modeling_gptsan_japanese.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptsan_japanese/__pycache__/tokenization_gptsan_japanese.cpython-310.pyc b/transformers/models/gptsan_japanese/__pycache__/tokenization_gptsan_japanese.cpython-310.pyc
deleted file mode 100644
index 9c5ee1de19e93f6d41f82456fcd42fcc50d9f89f..0000000000000000000000000000000000000000
Binary files a/transformers/models/gptsan_japanese/__pycache__/tokenization_gptsan_japanese.cpython-310.pyc and /dev/null differ
diff --git a/transformers/models/gptsan_japanese/configuration_gptsan_japanese.py b/transformers/models/gptsan_japanese/configuration_gptsan_japanese.py
deleted file mode 100644
index e0a17d1c114aef10038478955fac1d244e8c9a02..0000000000000000000000000000000000000000
--- a/transformers/models/gptsan_japanese/configuration_gptsan_japanese.py
+++ /dev/null
@@ -1,156 +0,0 @@
-# coding=utf-8
-# Copyright 2023, HuggingFace Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" GPTSAN-japanese model configuration"""
-from ...configuration_utils import PretrainedConfig
-from ...utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-from ..deprecated._archive_maps import GPTSAN_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP # noqa: F401, E402
-
-
-class GPTSanJapaneseConfig(PretrainedConfig):
- r"""
- This is the configuration class to store the configuration of a [`GPTSanJapaneseModel`]. It is used to instantiate
- a GPTSANJapanese model according to the specified arguments, defining the model architecture. Instantiating a
- configuration with the defaults will yield a similar configuration to that of the GPTSANJapanese
- [Tanrei/GPTSAN-japanese](https://huggingface.co/Tanrei/GPTSAN-japanese) architecture.
-
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
- documentation from [`PretrainedConfig`] for more information.
-
- Arguments:
- vocab_size (`int`, *optional*, defaults to 36000):
- Vocabulary size of the GPTSANJapanese model. Defines the number of different tokens that can be represented
- by the `inputs_ids` passed when calling [`GPTSanJapaneseModel`].
- max_position_embeddings (`int`, *optional*, defaults to 1280):
- The maximum sequence length that this model might ever be used with. Defaults set this to 1280.
- d_model (`int`, *optional*, defaults to 1024):
- Size of the encoder layers and the pooler layer.
- d_ff (`int`, *optional*, defaults to 8192):
- Size of the intermediate feed forward layer in each `SwitchTransformersBlock`.
- d_ext (`int`, *optional*, defaults to 4096):
- Size of the intermediate feed forward layer in each Extra-layers.
- d_spout (`int`, *optional*, defaults to 128):
- Size of the `spout` vector.
- num_switch_layers (`int`, *optional*, defaults to 10):
- Number of layers in the Switch Transformer layer.
- num_ext_layers (`int`, *optional*, defaults to 0):
- Number of layers in the Extra-layers.
- num_heads (`int`, *optional*, defaults to 16):
- Number of attention heads for each attention layer in the Transformer encoder.
- num_experts (`int`, *optional*, defaults to 16):
- Number of experts for each SwitchTransformer layer.
- expert_capacity (`int`, *optional*, defaults to 128):
- Number of tokens that can be stored in each expert. If set to 1, the model will behave like a regular
- Transformer.
- dropout_rate (`float`, *optional*, defaults to 0.0):
- The ratio for all dropout layers.
- layer_norm_eps (`float`, *optional*, defaults to 1e-5):
- The epsilon used by the layer normalization layers.
- router_bias (`bool`, *optional*, defaults to `False`):
- Whether to add a bias to the router.
- router_jitter_noise (`float`, *optional*, defaults to 0.0):
- Amount of noise to add to the router. Set it to 0.0 during prediction or set small value (usually 1e-2)
- during training.
- router_dtype (`str`, *optional*, default to `"float32"`):
- The `dtype` used for the routers. It is preferable to keep the `dtype` to `"float32"` as specified in the
- *selective precision* discussion in [the paper](https://arxiv.org/abs/2101.03961).
- router_ignore_padding_tokens (`bool`, *optional*, defaults to `False`):
- Whether to ignore padding tokens when routing.
- output_hidden_states (`bool`, *optional*, default to `False`):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- output_attentions (`bool`, *optional*, defaults to `False`):
- Whether or not to return the attentions tensors of all attention layers.
- initializer_factor (`float`, *optional*, defaults to 0.002):
- A factor for initializing all weight matrices.
- output_router_logits (`bool`, *optional*, default to `False`):
- Whether or not to return the router logits of all experts.
- use_cache (`bool`, *optional*, defaults to `True`):
- Whether or not the model should return the last key/values attentions (not used by all models)
- """
-
- model_type = "gptsan-japanese"
- keys_to_ignore_at_inference = [
- "past_key_values",
- ]
- attribute_map = {
- "hidden_size": "d_model",
- "num_attention_heads": "num_heads",
- "num_hidden_layers": "num_layers",
- }
-
- def __init__(
- self,
- vocab_size=36000,
- max_position_embeddings=1280,
- d_model=1024,
- d_ff=8192,
- d_ext=4096,
- d_spout=128,
- num_switch_layers=10,
- num_ext_layers=0,
- num_heads=16,
- num_experts=16,
- expert_capacity=128,
- dropout_rate=0.0,
- layer_norm_epsilon=1e-5,
- router_bias=False,
- router_jitter_noise=0.0,
- router_dtype="float32",
- router_ignore_padding_tokens=False,
- output_hidden_states=False,
- output_attentions=False,
- initializer_factor=0.002,
- output_router_logits=False,
- use_cache=True,
- separator_token_id=35998,
- pad_token_id=35995,
- eos_token_id=35999,
- **kwargs,
- ):
- self.vocab_size = vocab_size
- self.max_position_embeddings = max_position_embeddings
- self.d_model = d_model
- self.d_ff = d_ff
- self.d_ext = d_ext
- self.d_spout = d_spout
- self.num_switch_layers = num_switch_layers
- self.num_ext_layers = num_ext_layers
- self.num_layers = num_switch_layers + num_ext_layers
- self.num_heads = num_heads
- self.num_experts = num_experts
- self.expert_capacity = expert_capacity
- self.dropout_rate = dropout_rate
- self.layer_norm_epsilon = layer_norm_epsilon
- self.router_bias = router_bias
- self.router_jitter_noise = router_jitter_noise
- self.router_dtype = router_dtype
- self.router_ignore_padding_tokens = router_ignore_padding_tokens
- self.output_hidden_states = output_hidden_states
- self.output_attentions = output_attentions
- self.initializer_factor = initializer_factor
- self.output_router_logits = output_router_logits
- self.use_cache = use_cache
-
- super().__init__(
- separator_token_id=separator_token_id,
- pad_token_id=pad_token_id,
- eos_token_id=eos_token_id,
- **kwargs,
- )
diff --git a/transformers/models/gptsan_japanese/convert_gptsan_tf_checkpoint_to_pytorch.py b/transformers/models/gptsan_japanese/convert_gptsan_tf_checkpoint_to_pytorch.py
deleted file mode 100644
index a84d000d44390fe6ae821fb1cdfba968d40a2b93..0000000000000000000000000000000000000000
--- a/transformers/models/gptsan_japanese/convert_gptsan_tf_checkpoint_to_pytorch.py
+++ /dev/null
@@ -1,181 +0,0 @@
-# coding=utf-8
-# Copyright 2023 The HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Convert GPTSANJapanese checkpoints from the original repository to pytorch model."""
-
-import argparse
-import json
-import os
-from collections import OrderedDict
-
-import numpy as np
-import tensorflow as tf
-import torch
-
-
-def convert_tf_gptsan_to_pt(args):
- parameter_file = os.path.join(args.tf_model_dir, "parameters.json")
- params = json.loads(open(parameter_file).read())
- if not params:
- raise ValueError(
- f"It seems that the json file at {parameter_file} is empty. Make sure you have a correct json file."
- )
- if not args.output.endswith(".pt"):
- args.output = args.output + ".pt"
- new_state = OrderedDict()
- with tf.device("/CPU:0"):
- reader = tf.train.load_checkpoint(args.tf_model_dir)
- shapes = reader.get_variable_to_shape_map()
- for key_name in shapes.keys():
- vnp = reader.get_tensor(key_name).astype(np.float16)
- if key_name.endswith("/adam_m") or key_name.endswith("/adam_v"):
- continue
- if key_name.startswith("pasts/"):
- if key_name.startswith("pasts/mlp"):
- player = int(key_name[9])
- elif key_name.startswith("pasts/out"):
- player = 8
- name = "model.sqout.%d.weight" % (player * 2) # enter to nn.Sequencial with Tanh, so 2 at a time
- state = vnp.transpose([1, 0]).copy() # Mesh-Tensorflow is a diagonal matrix
- new_state[name] = torch.tensor(state)
- elif key_name.startswith("model/moe"):
- player = int(key_name[9:].split("/")[0])
- if key_name.endswith("/switch_gating/kernel"):
- name = "model.blocks.%d.feed_forward.mlp.router.classifier.weight" % player
- state = vnp.transpose([1, 0]).copy() # Mesh-Tensorflow is a diagonal matrix
- new_state[name] = torch.tensor(state)
- elif key_name.endswith("/softmlp/kernel"):
- name = "model.blocks.%d.feed_forward.soft_bypass_mlp.weight" % player
- state = vnp.transpose([1, 0]).copy() # Mesh-Tensorflow is a diagonal matrix
- new_state[name] = torch.tensor(state)
- elif key_name.endswith("/wo/kernel") or key_name.endswith("/wi/kernel"):
- nlayer = key_name[-9:-7]
- for i in range(16):
- name = "model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight" % (player, i, nlayer)
- state = (
- vnp[i].transpose([1, 0]).copy()
- ) # In Mesh-Tensorflow, it is one array, so it is divided
- new_state[name] = torch.tensor(state)
- elif key_name.startswith("model/mlp"):
- player = int(key_name[9:].split("/")[0])
- if key_name.endswith("/p1/kernel"):
- name = "model.blocks.%d.feed_forward.mlp.wi.weight" % player
- state = vnp.transpose([1, 0]).copy() # Mesh-Tensorflow is a diagonal matrix
- new_state[name] = torch.tensor(state)
- elif key_name.endswith("/p1/bias"):
- name = "model.blocks.%d.feed_forward.mlp.wi.bias" % player
- state = vnp.copy() # same because it is one dimensional
- new_state[name] = torch.tensor(state)
- elif key_name.endswith("/p2/kernel"):
- name = "model.blocks.%d.feed_forward.mlp.wo.weight" % player
- state = vnp.transpose([1, 0]).copy() # Mesh-Tensorflow is a diagonal matrix
- new_state[name] = torch.tensor(state)
- elif key_name.endswith("/p2/bias"):
- name = "model.blocks.%d.feed_forward.mlp.wo.bias" % player
- state = vnp.copy() # same because it is one dimensional
- new_state[name] = torch.tensor(state)
- elif key_name.startswith("model/ln"):
- player = int(key_name[8:].split("/")[0])
- if key_name.endswith("/b"):
- name = "model.blocks.%d.feed_forward.norm.bias" % player
- state = vnp.copy() # same because it is one dimensional
- new_state[name] = torch.tensor(state)
- elif key_name.endswith("/g"):
- name = "model.blocks.%d.feed_forward.norm.weight" % player
- state = vnp.copy() # same because it is one dimensional
- new_state[name] = torch.tensor(state)
- elif key_name.startswith("model/att"):
- player = int(key_name[9:].split("/")[0])
- if key_name.endswith("/qkv/kernel"):
- state = vnp.copy() # Compute same dimension as Mesh-tensorflow using einsum
- state_q = state[:, 0, :, :]
- state_k = state[:, 1, :, :]
- state_v = state[:, 2, :, :]
- state_q = (
- state_q.reshape([state_q.shape[0], state_q.shape[1] * state_q.shape[2]])
- .transpose([1, 0])
- .copy()
- ) # Mesh-Tensorflow is a diagonal matrix
- state_k = (
- state_k.reshape([state_k.shape[0], state_k.shape[1] * state_k.shape[2]])
- .transpose([1, 0])
- .copy()
- ) # Mesh-Tensorflow is a diagonal matrix
- state_v = (
- state_v.reshape([state_v.shape[0], state_v.shape[1] * state_v.shape[2]])
- .transpose([1, 0])
- .copy()
- ) # Mesh-Tensorflow is a diagonal matrix
- name = "model.blocks.%d.self_attn.self_attn.q_proj.weight" % player
- new_state[name] = torch.tensor(state_q)
- name = "model.blocks.%d.self_attn.self_attn.k_proj.weight" % player
- new_state[name] = torch.tensor(state_k)
- name = "model.blocks.%d.self_attn.self_attn.v_proj.weight" % player
- new_state[name] = torch.tensor(state_v)
- elif key_name.endswith("/o/kernel"):
- name = "model.blocks.%d.self_attn.self_attn.out_proj.weight" % player
- state = (
- vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]]).transpose([1, 0]).copy()
- ) # Mesh-Tensorflow is a diagonal matrix
- new_state[name] = torch.tensor(state)
- elif key_name.startswith("model/an"):
- player = int(key_name[8:].split("/")[0])
- if key_name.endswith("/b"):
- name = "model.blocks.%d.self_attn.norm.bias" % player
- state = vnp.copy() # same because it is one dimensional
- new_state[name] = torch.tensor(state)
- elif key_name.endswith("/g"):
- name = "model.blocks.%d.self_attn.norm.weight" % player
- state = vnp.copy() # same because it is one dimensional
- new_state[name] = torch.tensor(state)
- elif (
- key_name.startswith("model/wte")
- or key_name.startswith("model/wpe")
- or key_name.startswith("model/ete")
- ):
- nlayer = {"wte": "embed_tokens", "wpe": "position_embeddings", "ete": "extra_position_embeddings"}[
- key_name[-3:]
- ]
- name = "model.%s.weight" % nlayer
- state = vnp.copy() # same in embedded
- new_state[name] = torch.tensor(state)
- if key_name.startswith("model/wte"):
- name = "lm_head.weight"
- state = vnp.copy() # same in embedded
- new_state[name] = torch.tensor(state)
- elif key_name.startswith("model/wob"):
- name = "final_logits_bias"
- state = vnp.copy() # same in embedded
- state = state.reshape((1, -1))
- new_state[name] = torch.tensor(state)
- elif key_name == "model/dense/kernel":
- name = "model.last_project.weight"
- state = vnp.transpose([1, 0]).copy() # Mesh-Tensorflow is a diagonal matrix
- new_state[name] = torch.tensor(state)
- elif key_name == "model/dense_1/bias":
- name = "model.last_project.bias"
- state = vnp.copy() # same because it is one dimensional
- new_state[name] = torch.tensor(state)
- torch.save(new_state, args.output)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser(
- description="model converter.", formatter_class=argparse.ArgumentDefaultsHelpFormatter
- )
- parser.add_argument("--tf_model_dir", metavar="PATH", type=str, required=True, help="import model")
- parser.add_argument("--output", metavar="PATH", type=str, required=True, help="output model")
- args = parser.parse_args()
- convert_tf_gptsan_to_pt(args)
diff --git a/transformers/models/gptsan_japanese/modeling_gptsan_japanese.py b/transformers/models/gptsan_japanese/modeling_gptsan_japanese.py
deleted file mode 100644
index 59252bc567a4624a058d65c1e07270c185b8a71b..0000000000000000000000000000000000000000
--- a/transformers/models/gptsan_japanese/modeling_gptsan_japanese.py
+++ /dev/null
@@ -1,1343 +0,0 @@
-# coding=utf-8
-# Copyright 2023 Toshiyuki Sakamoto(tanreinama) and HuggingFace Inc. team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-""" PyTorch GPTSANJapanese model."""
-
-
-import copy
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-
-from ...activations import ACT2FN
-from ...modeling_outputs import MoECausalLMOutputWithPast, MoEModelOutputWithPastAndCrossAttentions
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
- DUMMY_INPUTS,
- DUMMY_MASK,
- add_start_docstrings,
- add_start_docstrings_to_model_forward,
- is_torch_fx_proxy,
- logging,
-)
-from .configuration_gptsan_japanese import GPTSanJapaneseConfig
-
-
-logger = logging.get_logger(__name__)
-
-_CONFIG_FOR_DOC = "GPTSanJapaneseConfig"
-_CHECKPOINT_FOR_DOC = "Tanrei/GPTSAN-japanese"
-
-####################################################
-# This dict contains ids and associated url
-# for the pretrained weights provided with the models
-####################################################
-
-from ..deprecated._archive_maps import GPTSAN_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST # noqa: F401, E402
-
-
-# Copied from transformers.models.switch_transformers.modeling_switch_transformers.router_z_loss_func
-def router_z_loss_func(router_logits: torch.Tensor) -> float:
- r"""
- Compute the router z-loss implemented in PyTorch.
-
- The router z-loss was introduced in [Designing Effective Sparse Expert Models](https://arxiv.org/abs/2202.08906).
- It encourages router logits to remain small in an effort to improve stability.
-
- Args:
- router_logits (`float`):
- Input logits of shape [batch_size, sequence_length, num_experts]
-
- Returns:
- Scalar router z-loss.
- """
- num_groups, tokens_per_group, _ = router_logits.shape
- log_z = torch.logsumexp(router_logits, dim=-1)
- z_loss = log_z**2
- return torch.sum(z_loss) / (num_groups * tokens_per_group)
-
-
-# Copied from transformers.models.switch_transformers.modeling_switch_transformers.load_balancing_loss_func
-def load_balancing_loss_func(router_probs: torch.Tensor, expert_indices: torch.Tensor) -> float:
- r"""
- Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
-
- See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
- function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
- experts is too unbalanced.
-
- Args:
- router_probs (`torch.Tensor`):
- Probability assigned to each expert per token. Shape: [batch_size, seqeunce_length, num_experts].
- expert_indices (`torch.Tensor`):
- Indices tensor of shape [batch_size, seqeunce_length] identifying the selected expert for a given token.
-
- Returns:
- The auxiliary loss.
- """
- num_experts = router_probs.shape[-1]
-
- # cast the expert indices to int64, otherwise one-hot encoding will fail
- if expert_indices.dtype != torch.int64:
- expert_indices = expert_indices.to(torch.int64)
-
- if len(expert_indices.shape) == 2:
- expert_indices = expert_indices.unsqueeze(2)
-
- expert_mask = torch.nn.functional.one_hot(expert_indices, num_experts)
-
- # For a given token, determine if it was routed to a given expert.
- expert_mask = torch.max(expert_mask, axis=-2).values
-
- # cast to float32 otherwise mean will fail
- expert_mask = expert_mask.to(torch.float32)
- tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)
-
- router_prob_per_group_and_expert = torch.mean(router_probs, axis=-2)
- return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert) * (num_experts**2)
-
-
-class GPTSanJapaneseDenseActDense(nn.Module):
- """
- FFN Layer for Switch Transformer and Extra layers
-
- GPTSAN can mix Switch Transformer layers and normal Transformer layers This class is used as Expert in Switch
- Transformer layers and as FFN in regular Transformer layers. RELU is used in the Switch Transformer layer, and
- Swish is used in the normal Transformer layer, so there is a choice of which is used in the argument.
-
- """
-
- def __init__(self, config: GPTSanJapaneseConfig, ext_layer=False):
- super().__init__()
- d_inter = config.d_ext if ext_layer else config.d_ff
- self.wi = nn.Linear(config.d_model, d_inter, bias=ext_layer)
- self.wo = nn.Linear(d_inter, config.d_model, bias=ext_layer)
- self.dropout = nn.Identity() if ext_layer else nn.Dropout(config.dropout_rate)
- self.act = ACT2FN["swish" if ext_layer else "relu"]
-
- def forward(self, hidden_states):
- r"""
- Args:
- hidden_states (`torch.Tensor`) :
- [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
- Returns:
- torch.Tensor[num_groups, tokens_per_group, hidden_dim]
-
- """
- hidden_states = self.wi(hidden_states)
- hidden_states = self.act(hidden_states)
- hidden_states = self.dropout(hidden_states)
- hidden_states = self.wo(hidden_states)
- return hidden_states
-
-
-# Copied from transformers.models.switch_transformers.modeling_switch_transformers.SwitchTransformersTop1Router with SwitchTransformers->GPTSanJapanese
-class GPTSanJapaneseTop1Router(nn.Module):
- """
- Router using tokens choose top-1 experts assignment.
-
- This router uses the same mechanism as in Switch Transformer (https://arxiv.org/abs/2101.03961) and V-MoE
- (https://arxiv.org/abs/2106.05974): tokens choose their top experts. Items are sorted by router_probs and then
- routed to their choice of expert until the expert's expert_capacity is reached. **There is no guarantee that each
- token is processed by an expert**, or that each expert receives at least one token.
-
- """
-
- def __init__(self, config: GPTSanJapaneseConfig):
- super().__init__()
- self.num_experts = config.num_experts
- self.expert_capacity = config.expert_capacity
- self.classifier = nn.Linear(config.hidden_size, self.num_experts, bias=config.router_bias)
- self.jitter_noise = config.router_jitter_noise
- self.ignore_padding_tokens = config.router_ignore_padding_tokens
- self.dtype = getattr(torch, config.router_dtype)
-
- def _compute_router_probabilities(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
- r"""
- Computes router probabilities from input hidden states.
-
- Args:
- hidden_states (`torch.Tensor`):
- (batch_size, sequence_length, hidden_dim) from which router probabilities are computed.
- Returns:
- router_probabilities (`torch.Tensor`):
- Tensor of shape (batch_size, sequence_length, num_experts) corresponding to the probabilities for each
- token and expert. Used for routing tokens to experts.
- router_logits (`torch.Tensor`):
- Logits tensor of shape (batch_size, sequence_length, num_experts) corresponding to raw router logits.
- This is used later for computing router z-loss.
- """
- # float32 is used to ensure stability. See the discussion of "selective precision" in
- # https://arxiv.org/abs/2101.03961.
- # We also store the previous dtype to cast back the output to the previous dtype
- self.input_dtype = hidden_states.dtype
- hidden_states = hidden_states.to(self.dtype)
-
- if self.training and self.jitter_noise > 0:
- # Multiply the token inputs by the uniform distribution - adding some noise
- hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)
-
- # Shape: [num_groups, tokens_per_group, num_experts]
- self._cast_classifier()
- router_logits = self.classifier(hidden_states)
-
- # Apply Softmax and cast back to the original `dtype`
- router_probabilities = nn.functional.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype)
- return router_probabilities, router_logits
-
- def _cast_classifier(self):
- r"""
- `bitsandbytes` `Linear8bitLt` layers does not support manual casting Therefore we need to check if they are an
- instance of the `Linear8bitLt` class by checking special attributes.
- """
- if not (hasattr(self.classifier, "SCB") or hasattr(self.classifier, "CB")):
- self.classifier = self.classifier.to(self.dtype)
-
- def forward(self, hidden_states: torch.Tensor) -> Tuple:
- r"""
- Generic forward function for every Router class. Each Router expects to have the same input hidden states
- (`hidden_states`) corresponding to the hidden states for each token, the `expert_capacity` corresponding to the
- number of tokens the Router will send to each expert, some Routers can send up to few tokens to each expert.
-
- Each Router works as the following: it expects the hidden states for each token, gets the `router_probs` and
- `router_logits` from the `router_weights`. This will assign for each token, the raw probability to be assigned
- to an expert. Then each Router class will have to define its own `_compute_routing_instructions`.
-
- Args:
- hidden_states (`torch.Tensor`) :
- [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
- Returns:
- Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`] Tuple containing the expert index, the router probs
- and the router logits. The router probabilities and logits are required to compute the loss.
- """
- router_probs, router_logits = self._compute_router_probabilities(hidden_states)
-
- expert_index = torch.argmax(router_probs, dim=-1)
- expert_index = torch.nn.functional.one_hot(expert_index, num_classes=self.num_experts)
-
- # Mask tokens outside expert capacity. Sum over each sequence
- token_priority = torch.cumsum(expert_index, dim=-2)
- # mask if the token routed to to the expert will overflow
- expert_capacity_mask = token_priority <= self.expert_capacity
- expert_index = expert_index * expert_capacity_mask
-
- router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(-1)
- return expert_index, router_probs, router_logits
-
-
-# Copied from transformers.models.switch_transformers.modeling_switch_transformers.SwitchTransformersSparseMLP with SwitchTransformers->GPTSanJapanese
-class GPTSanJapaneseSparseMLP(nn.Module):
- r"""
- Implementation of the Switch Transformers Sparse MLP module.
- """
-
- def __init__(self, config: GPTSanJapaneseConfig, expert_class: nn.Module = GPTSanJapaneseDenseActDense):
- super().__init__()
- # Step 1: Get the correct router according to its class
- self.router = GPTSanJapaneseTop1Router(config)
-
- # Step 2: Get the experts
- self.experts = nn.ModuleDict()
- for idx in range(config.num_experts):
- self.experts[f"expert_{idx}"] = expert_class(config)
-
- def forward(self, hidden_states):
- r"""
- Hold on, this will be slightly tricky to understand In the correct order, a MoE layer does the following:
-
- 1- Gets the `router_mask` from the router. The shape of the mask is `(batch_size, sequence_length, num_expert)`
- and corresponds to the argmax of the `router_probs`. The probabilities are needed in the computation of the
- hidden states : they are broadcasted to the hidden states values (can be interpreted as a scaling factor).
-
- 2- Dispatch the tokens to its associated experts. We do a classic for loop over the experts and assign for each
- expert the corresponding hidden states.
-
- """
- # Step 1: Get the router_mask from the router as wel as the probabilities
- router_mask, router_probs, router_logits = self.router(hidden_states)
- expert_index = torch.argmax(router_mask, dim=-1)
-
- # The routers introduced might not always map all the tokens, to a router, which means that some hidden states
- # can be unchanged from one layer to another. That is why the hidden states are cloned before updating only the seleced ones.
-
- next_states = hidden_states.clone()
- for idx, expert in enumerate(self.experts.values()):
- token_indices = router_mask[:, :, idx].bool()
- next_states[token_indices] = expert(hidden_states[token_indices]).to(next_states.dtype)
-
- hidden_states = router_probs * next_states
- return hidden_states, (router_logits, expert_index)
-
-
-class GPTSanJapaneseLayerSparseFF(nn.Module):
- r"""
- Switch Transformers Feed Forward layer module. This is a wrapper around the Mixture of Experts module.
-
- Parameters:
- config : ([`GPTSanJapaneseConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
- """
-
- def __init__(self, config: GPTSanJapaneseConfig):
- super().__init__()
- self.mlp = GPTSanJapaneseSparseMLP(config)
- self.soft_bypass_mlp = nn.Linear(config.d_model, config.d_model, bias=False)
- self.norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon)
-
- def forward(self, hidden_states, output_router_logits):
- r"""
- Args:
- hidden_states (`torch.Tensor`) :
- [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
- output_router_logits (`bool`) :
- output experts router output.
- Returns:
- torch.Tensor[num_groups, tokens_per_group, hidden_dim]
-
- """
- forwarded_states, router_tuple = self.mlp(hidden_states)
- forwarded_states += torch.tanh(self.soft_bypass_mlp(hidden_states))
- output = hidden_states + self.norm(forwarded_states)
-
- if output_router_logits and router_tuple is not None:
- return output, router_tuple
- else:
- return output
-
-
-class GPTSanJapaneseLayerDenseFF(nn.Module):
- r"""
- Extra Transformers Feed Forward layer module.
-
- Parameters:
- config : ([`GPTSanJapaneseConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
- """
-
- def __init__(self, config: GPTSanJapaneseConfig):
- super().__init__()
- # Check if it is a sparse layer, if not then it is a dense layer
- self.mlp = GPTSanJapaneseDenseActDense(config, ext_layer=True)
- self.norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon)
-
- def forward(self, hidden_states):
- r"""
- Args:
- hidden_states (`torch.Tensor`) :
- [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
- Returns:
- torch.Tensor[num_groups, tokens_per_group, hidden_dim]
-
- """
- forwarded_states = self.mlp(hidden_states)
- output = hidden_states + self.norm(forwarded_states)
- return output
-
-
-# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->GPTSanJapanese
-class GPTSanJapaneseAttention(nn.Module):
- """Multi-headed attention from 'Attention Is All You Need' paper"""
-
- def __init__(
- self,
- embed_dim: int,
- num_heads: int,
- dropout: float = 0.0,
- is_decoder: bool = False,
- bias: bool = True,
- is_causal: bool = False,
- config: Optional[GPTSanJapaneseConfig] = None,
- ):
- super().__init__()
- self.embed_dim = embed_dim
- self.num_heads = num_heads
- self.dropout = dropout
- self.head_dim = embed_dim // num_heads
- self.config = config
-
- if (self.head_dim * num_heads) != self.embed_dim:
- raise ValueError(
- f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
- f" and `num_heads`: {num_heads})."
- )
- self.scaling = self.head_dim**-0.5
- self.is_decoder = is_decoder
- self.is_causal = is_causal
-
- self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
- self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
-
- def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
- return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- key_value_states: Optional[torch.Tensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.Tensor] = None,
- layer_head_mask: Optional[torch.Tensor] = None,
- output_attentions: bool = False,
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel"""
-
- # if key_value_states are provided this layer is used as a cross-attention layer
- # for the decoder
- is_cross_attention = key_value_states is not None
-
- bsz, tgt_len, _ = hidden_states.size()
-
- # get query proj
- query_states = self.q_proj(hidden_states) * self.scaling
- # get key, value proj
- # `past_key_value[0].shape[2] == key_value_states.shape[1]`
- # is checking that the `sequence_length` of the `past_key_value` is the same as
- # the provided `key_value_states` to support prefix tuning
- if (
- is_cross_attention
- and past_key_value is not None
- and past_key_value[0].shape[2] == key_value_states.shape[1]
- ):
- # reuse k,v, cross_attentions
- key_states = past_key_value[0]
- value_states = past_key_value[1]
- elif is_cross_attention:
- # cross_attentions
- key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
- value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
- elif past_key_value is not None:
- # reuse k, v, self_attention
- key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
- value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
- else:
- # self_attention
- key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
- value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-
- if self.is_decoder:
- # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
- # Further calls to cross_attention layer can then reuse all cross-attention
- # key/value_states (first "if" case)
- # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
- # all previous decoder key/value_states. Further calls to uni-directional self-attention
- # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
- # if encoder bi-directional self-attention `past_key_value` is always `None`
- past_key_value = (key_states, value_states)
-
- proj_shape = (bsz * self.num_heads, -1, self.head_dim)
- query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
- key_states = key_states.reshape(*proj_shape)
- value_states = value_states.reshape(*proj_shape)
-
- src_len = key_states.size(1)
- attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
-
- if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
- raise ValueError(
- f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
- f" {attn_weights.size()}"
- )
-
- if attention_mask is not None:
- if attention_mask.size() != (bsz, 1, tgt_len, src_len):
- raise ValueError(
- f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
- )
- attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- if layer_head_mask is not None:
- if layer_head_mask.size() != (self.num_heads,):
- raise ValueError(
- f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
- f" {layer_head_mask.size()}"
- )
- attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
- attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
- if output_attentions:
- # this operation is a bit awkward, but it's required to
- # make sure that attn_weights keeps its gradient.
- # In order to do so, attn_weights have to be reshaped
- # twice and have to be reused in the following
- attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
- attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
- else:
- attn_weights_reshaped = None
-
- attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-
- attn_output = torch.bmm(attn_probs, value_states)
-
- if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
-
- attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
- attn_output = attn_output.transpose(1, 2)
-
- # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
- # partitioned across GPUs when using tensor-parallelism.
- attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
-
- attn_output = self.out_proj(attn_output)
-
- return attn_output, attn_weights_reshaped, past_key_value
-
-
-class GPTSanJapaneseLayerSelfAttention(nn.Module):
- """
- Self Attention and Normalization Unit
- """
-
- def __init__(self, config, has_relative_attention_bias=False):
- super().__init__()
- self.self_attn = GPTSanJapaneseAttention(
- embed_dim=config.d_model,
- num_heads=config.num_heads,
- is_decoder=True,
- bias=has_relative_attention_bias,
- )
- self.norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon)
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
- r"""
- Self-attention and normalize block.
-
- Args:
- hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
- if the model is configured as a decoder.
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up
- decoding. If `past_key_values` are used, the user can optionally input only the last
- `decoder_input_ids` (those that don't have their past key value states given to this model) of shape
- `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used
- in the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- head_mask (`numpy.ndarray` of shape `({0})`, `optional):
- Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- Returns:
- Tuple[torch.Tensor[num_groups, tokens_per_group, hidden_dim],...]
- """
- # Self Attention
- # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
- self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
- # add present self-attn cache to positions 1,2 of present_key_value tuple
- atten_out = self.self_attn(
- hidden_states=hidden_states,
- past_key_value=self_attn_past_key_value,
- attention_mask=(1 - attention_mask) * torch.finfo(hidden_states.dtype).min,
- layer_head_mask=head_mask,
- output_attentions=output_attentions,
- )
- if output_attentions:
- attn_weights = (atten_out[1],)
- else:
- attn_weights = ()
-
- attention_output = atten_out[0]
-
- hidden = hidden_states + self.norm(attention_output)
-
- if use_cache:
- outputs = (hidden, atten_out[2]) # hidden, present, (attentions)
- else:
- outputs = (hidden,) # hidden, (attentions)
-
- return outputs + attn_weights
-
-
-class GPTSanJapaneseBlock(nn.Module):
- """
- Self Attention and FFN Unit
- """
-
- def __init__(self, config, ext_layer=False):
- super().__init__()
- self.self_attn = GPTSanJapaneseLayerSelfAttention(config)
- self.feed_forward = GPTSanJapaneseLayerDenseFF(config) if ext_layer else GPTSanJapaneseLayerSparseFF(config)
-
- def forward(
- self,
- hidden_states: Optional[Tuple[torch.FloatTensor]],
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- output_attentions: Optional[bool] = False,
- output_router_tuple: Optional[bool] = False,
- ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
- r"""
- GPTSAN transformer block.
-
- Args:
- hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
- if the model is configured as a decoder.
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up
- decoding. If `past_key_values` are used, the user can optionally input only the last
- `decoder_input_ids` (those that don't have their past key value states given to this model) of shape
- `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used
- in the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- head_mask (`numpy.ndarray` of shape `({0})`, `optional):
- Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
-
- - 1 indicates the head is **not masked**,
- - 0 indicates the head is **masked**.
-
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- output_attentions (`bool`) :
- output attention probabirities.
- output_router_tuple:
- output experts router logits and expert id.
- Returns:
- Tuple[torch.Tensor[num_groups, tokens_per_group, hidden_dim],...]
- """
- atten_out = self.self_attn(
- hidden_states=hidden_states,
- past_key_value=past_key_value,
- attention_mask=attention_mask,
- head_mask=head_mask,
- use_cache=use_cache,
- output_attentions=output_attentions,
- )
- attention_output = atten_out[0]
-
- if isinstance(self.feed_forward, GPTSanJapaneseLayerSparseFF):
- sparse_out = self.feed_forward(attention_output, output_router_tuple)
- if output_router_tuple:
- hidden, router_tuple = sparse_out
- else:
- hidden = sparse_out
- else:
- hidden = self.feed_forward(attention_output)
-
- outputs = (hidden,) + atten_out[1:]
-
- if isinstance(self.feed_forward, GPTSanJapaneseLayerSparseFF) and output_router_tuple:
- outputs += (router_tuple,)
-
- return outputs
-
-
-class GPTSanJapanesePreTrainedModel(PreTrainedModel):
- """
- An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
- models.
- """
-
- config_class = GPTSanJapaneseConfig
- base_model_prefix = "gptsan_japanese"
- supports_gradient_checkpointing = False
- _no_split_modules = ["GPTSanJapaneseBlock"]
- _skip_keys_device_placement = "past_key_values"
-
- @property
- def dummy_inputs(self):
- input_ids = torch.tensor(DUMMY_INPUTS)
- input_mask = torch.tensor(DUMMY_MASK)
- dummy_inputs = {
- "input_ids": input_ids,
- "attention_mask": input_mask,
- }
- return dummy_inputs
-
- def _init_weights(self, module):
- """Initialize the weights"""
- factor = self.config.initializer_factor # Used for testing weights initialization
- if isinstance(module, nn.LayerNorm):
- module.weight.data.fill_(factor * 1.0)
- module.bias.data.zero_()
- elif isinstance(module, nn.Linear):
- module.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))
- if hasattr(module, "bias") and module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=factor * 1.0)
- elif isinstance(module, GPTSanJapaneseModel):
- # Mesh TensorFlow embeddings initialization
- # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624
- module.embed_tokens.weight.data.normal_(mean=0.0, std=factor * 1.0)
- module.position_embeddings.weight.data.normal_(mean=0.0, std=factor * 1.0)
- if hasattr(module, "extra_position_embeddings") and module.extra_position_embeddings is not None:
- module.extra_position_embeddings.weight.data.normal_(mean=0.0, std=factor * 1.0)
- elif isinstance(module, (GPTSanJapaneseModel, GPTSanJapaneseForConditionalGeneration)):
- # Mesh TensorFlow embeddings initialization
- # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624
- module.final_logits_bias.data.normal_(mean=0.0, std=factor * 1.0)
- if hasattr(module, "lm_head") and not self.config.tie_word_embeddings:
- module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0)
- elif isinstance(module, GPTSanJapaneseDenseActDense):
- # Mesh TensorFlow FF initialization
- # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56
- # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89
- module.wi.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))
- if hasattr(module.wi, "bias") and module.wi.bias is not None:
- module.wi.bias.data.zero_()
- module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5))
- if hasattr(module.wo, "bias") and module.wo.bias is not None:
- module.wo.bias.data.zero_()
- elif isinstance(module, GPTSanJapaneseAttention):
- # Multi-headed attention
- d_model = self.config.d_model
- key_value_proj_dim = self.config.d_model
- n_heads = self.config.num_heads
- module.k_proj.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))
- module.v_proj.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))
- module.q_proj.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))
- module.out_proj.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5))
- elif isinstance(module, GPTSanJapaneseSparseMLP):
- # Mesh TensorFlow attention initialization to avoid scaling before softmax
- # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136
- d_model = self.config.d_model
- key_value_proj_dim = self.config.d_model
- n_heads = self.config.num_heads
- module.router.classifier.weight.data.normal_(mean=0.0, std=factor * 1)
- for idx in range(self.config.num_experts):
- module.experts[f"expert_{idx}"].wi.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5))
- module.experts[f"expert_{idx}"].wo.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5))
-
- # Copied from transformers.models.t5.modeling_t5.T5PreTrainedModel._shift_right
- def _shift_right(self, input_ids):
- decoder_start_token_id = self.config.decoder_start_token_id
- pad_token_id = self.config.pad_token_id
-
- if decoder_start_token_id is None:
- raise ValueError(
- "self.model.config.decoder_start_token_id has to be defined. In T5 it is usually set to the pad_token_id. "
- "See T5 docs for more information."
- )
-
- # shift inputs to the right
- if is_torch_fx_proxy(input_ids):
- # Item assignment is not supported natively for proxies.
- shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id)
- shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1)
- else:
- shifted_input_ids = input_ids.new_zeros(input_ids.shape)
- shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
- shifted_input_ids[..., 0] = decoder_start_token_id
-
- if pad_token_id is None:
- raise ValueError("self.model.config.pad_token_id has to be defined.")
- # replace possible -100 values in labels by `pad_token_id`
- shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
-
- return shifted_input_ids
-
-
-GPTSAN_JAPANESE_START_DOCSTRING = r"""
-
- The [GPTSAN-japanese](https://github.com/tanreinama/GPTSAN) model was proposed in General-purpose Swich transformer
- based Japanese language model
-
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
- and behavior.
-
- Parameters:
- config ([`GPTSanJapaneseConfig`]): Model configuration class with all the parameters of the model.
- Initializing with a config file does not load the weights associated with the model, only the
- configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
-"""
-
-GPTSAN_JAPANESE_INPUTS_DOCSTRING = r"""
- Args:
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
- Indices of input sequence tokens in the vocabulary. GPTSAN-japanese is a model that generates sentence
- continuations or predicts tokens at mask positions. Special tokens required for inputs to the model are
- automatically appended.
- attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **not masked**,
- - 0 for tokens that are **masked**.
-
- [What are attention masks?](../glossary#attention-mask)
- token_type_ids (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
- An input that masks the Prefix part in the Prefix-LM input. Mask values selected in `[0, 1]`:
-
- - 1 for tokens that are **prefix** input,
- - 0 for tokens that are **not-prefix** input.
- spout (`torch.Tensor` of shape `(batch_size, config.d_spout)`):
- This vector is transformed through an 8-layer FFN and can be used instead of `past_key_values`.
- past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
- Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
-
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
- Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
- `past_key_values`).
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
- model's internal embedding lookup matrix.
- decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*):
- Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded
- representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be
- input (see `past_key_values`). This is useful if you want more control over how to convert
- `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
- tensors for more detail.
- output_hidden_states (`bool`, *optional*):
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
- more detail.
- return_dict (`bool`, *optional*):
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
- router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`):
- Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
- Router logits of the decoder model, useful to compute the auxiliary loss for Mixture of Experts models.
-"""
-
-
-@add_start_docstrings(
- "The bare GPTSAN-japanese Model transformer outputting raw hidden-states without any specific head on top.",
- GPTSAN_JAPANESE_START_DOCSTRING,
-)
-class GPTSanJapaneseModel(GPTSanJapanesePreTrainedModel):
- def __init__(self, config: GPTSanJapaneseConfig):
- super().__init__(config)
- self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.d_model)
- self.config = copy.deepcopy(config)
- self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
- self.last_project = nn.Linear(config.d_model, config.d_model, bias=True)
- self.act = ACT2FN["swish"]
-
- self.blocks = torch.nn.ModuleList([])
- for _ in range(config.num_switch_layers):
- self.blocks.append(GPTSanJapaneseBlock(config))
- for _ in range(config.num_ext_layers):
- self.blocks.append(GPTSanJapaneseBlock(config, ext_layer=True))
-
- if config.num_ext_layers > 0:
- self.extra_position_embeddings = nn.Embedding(config.max_position_embeddings, config.d_model)
-
- if config.d_spout:
- spouts = []
- for _ in range(8):
- spouts.append(nn.Linear(config.d_spout, config.d_spout, bias=False))
- spouts.append(nn.Tanh())
- spouts.append(nn.Linear(config.d_spout, config.num_layers * 2 * config.d_model, bias=False))
- self.spout = nn.Sequential(*spouts)
-
- self.post_init()
-
- def get_input_embeddings(self):
- return self.embed_tokens
-
- def set_input_embeddings(self, new_embeddings):
- self.embed_tokens = new_embeddings
-
- @add_start_docstrings_to_model_forward(GPTSAN_JAPANESE_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.FloatTensor] = None,
- spout: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- output_router_logits: Optional[bool] = None,
- num_precontext: Optional[torch.LongTensor] = None,
- ) -> Union[MoEModelOutputWithPastAndCrossAttentions, Tuple[torch.FloatTensor]]:
- r"""
- num_precontext (`torch.LongTensor` of shape `(batch_size,1)`):
- length of `hybrid` input tokens in the input. Tokens up to this length refer to both front and back like
- BERT, tokens after that refer only to front like GPT. see also:
- https://github.com/tanreinama/GPTSAN/blob/main/report/model.md
-
- Returns:
- `MoEModelOutputWithPastAndCrossAttentions` or `tuple` if `return_dict` returns
- MoEModelOutputWithPastAndCrossAttentions insted of tuple
- """
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- device = self.position_embeddings.weight.device
- if input_ids is None:
- input_ids = torch.zeros([1, 1]).int().to(device) # dummy for input_ids was None
- num_pasts_contexts = 0
- num_batch = input_ids.shape[0]
- pasts_or_spout_value = None
- if past_key_values is not None:
- num_pasts_contexts = past_key_values[0][0].shape[2]
- elif self.config.d_spout and spout is not None:
- # `spout` is a special input vector specific to GPTSAN
- # This controls the output by projecting embedded information such as the class of sentences during learning.
- # It should passed instead of the first past_key_value.
- # See the original GPTSAN repository for details
- num_pasts_contexts += 1
-
- # If there is an attention_mask, increase first one for spout
- if self.config.d_spout and spout is not None and attention_mask is not None:
- attention_mask_with_spout = torch.ones(num_batch, attention_mask.shape[1] + 1, device=device)
- attention_mask_with_spout[:, 1:] -= 1 - attention_mask # 1st token should be spout
- attention_mask = attention_mask_with_spout # update attention_mask
-
- if num_precontext is not None:
- # `num_precontext` is the number of tokens that refer to each other in prefix-lm
- # created per batch, so dimension of num_precontext should be [batch, 1]
- if not (
- len(num_precontext.shape) == 2 and num_precontext.shape[1] == 1
- ): # num_precontext Should be [batch,1]
- raise ValueError("num_precontext should be [batch, 1] size.")
- num_precontext = torch.reshape(num_precontext, [-1])
- else:
- num_precontext = torch.zeros([num_batch]).int().to(device)
-
- num_input_contexts = input_ids.shape[1]
- num_output_contexts = num_input_contexts + num_pasts_contexts
-
- hidden_states = self.embed_tokens(input_ids)
-
- if past_key_values is not None:
- pasts_or_spout_value = past_key_values
- elif self.config.d_spout and spout is not None:
- # Make vector from `spout` of GPTSAN to the same shape as past_key_values
- pasts_or_spout_value = self.spout(spout) # projecting `spout` vector
- pasts_or_spout_value = torch.reshape(
- pasts_or_spout_value,
- [
- num_batch,
- self.config.num_layers,
- 2,
- self.config.num_heads,
- num_pasts_contexts,
- self.config.d_model // self.config.num_heads,
- ],
- )
- pasts_or_spout_value = torch.split(pasts_or_spout_value, [1] * self.config.num_layers, dim=1)
- # make same shape as past_key_values
- pasts_or_spout_value = tuple(
- tuple([b.squeeze(1) for b in torch.split(a.squeeze(1), [1, 1], dim=1)]) for a in pasts_or_spout_value
- )
- else:
- pasts_or_spout_value = [None] * self.config.num_layers
-
- # Token position considering spout and pasts
- token_position = torch.arange(num_input_contexts).to(device) + num_pasts_contexts
-
- if attention_mask is None:
- attention_mask = torch.ones(num_batch, num_input_contexts, device=device)
-
- # positions for get position_embeddings
- gather_position = (
- (
- torch.zeros((num_batch, self.config.d_model, num_input_contexts)).to(device)
- + token_position.unsqueeze(0)
- )
- .transpose(1, 2)
- .long()
- )
- # When padding with padding_side="left", zeros line up on the left side of attention_mask, so position_embeddings is shifted accordingly
- gather_position -= (1 - attention_mask).argmin(dim=-1).unsqueeze(1).unsqueeze(2)
- gather_position = torch.clip(gather_position, num_pasts_contexts, self.config.max_position_embeddings - 1)
-
- # attention_mask is applied per batch
- for i in range(num_batch):
- hidden_states[i] += torch.gather(self.position_embeddings.weight, dim=0, index=gather_position[i])
-
- # Create a mask to be used when making the prefix Input length of Prefix-LM variable
- causal_mask = (
- torch.tril(torch.ones((num_output_contexts, num_output_contexts), dtype=torch.uint8))
- .view(1, 1, num_output_contexts, num_output_contexts)
- .to(device)
- )
- prefix_lm_mask = causal_mask[:, :, -num_input_contexts:, :]
- if token_type_ids is not None:
- token_type_ids = token_type_ids.unsqueeze(1).unsqueeze(2)
- prefix_lm_mask = ((prefix_lm_mask + token_type_ids) > 0).float()
- # Marge prefix_lm_mask and attention_mask
- extended_attention_mask = prefix_lm_mask * attention_mask.unsqueeze(1).unsqueeze(2)
-
- # Prepare head mask if needed
- if head_mask is not None:
- head_mask = self.get_head_mask(
- head_mask, self.config.num_switch_layers + self.config.num_ext_layers
- ) # n_layer x batch x n_heads x N x N
-
- # outputs
- present_key_value_states = () if self.config.use_cache or use_cache else None
- all_hidden_states = () if self.config.output_hidden_states or output_hidden_states else None
- all_attentions = () if self.config.output_attentions or output_attentions else None
- all_router_probs = () if self.config.output_router_logits or output_router_logits else None
-
- for layer, past in enumerate(pasts_or_spout_value):
- if layer == self.config.num_switch_layers:
- if self.config.num_ext_layers > 0:
- # extra_position_embeddings are extra position embeddings that are only created when extending the model with code from the original GPTSAN repository. Not used in the default model.
- # However, it is created when you create an additional layer and partially train only that location.
- # Therefore, convert_gptsan_tf_checkpoint_to_pytorch.py is used when converting and loading models created in the original GPTSAN repository.
- for i in range(num_batch):
- hidden_states[i] += torch.gather(
- self.extra_position_embeddings.weight, dim=0, index=gather_position[i]
- )
-
- output_router_tuple = (
- self.config.output_router_logits or output_router_logits
- ) and layer < self.config.num_switch_layers
- block_output = self.blocks[layer](
- hidden_states=hidden_states,
- past_key_value=past,
- attention_mask=extended_attention_mask,
- head_mask=head_mask,
- use_cache=self.config.use_cache or use_cache,
- output_attentions=self.config.output_attentions or output_attentions,
- output_router_tuple=output_router_tuple,
- )
-
- outpos = 0
- hidden_states = block_output[outpos]
- if self.config.output_hidden_states or output_hidden_states:
- all_hidden_states += (hidden_states,)
- if self.config.use_cache or use_cache:
- outpos += 1
- present = block_output[outpos]
- present_key_value_states += (present,)
- if self.config.output_attentions or output_attentions:
- outpos += 1
- attention_probs = block_output[outpos]
- all_attentions += (attention_probs,)
- if output_router_tuple:
- outpos += 1
- router_tuple = block_output[outpos]
- all_router_probs.append(router_tuple[0])
-
- hidden_states = self.last_project(hidden_states)
- hidden_states = self.act(hidden_states)
-
- if self.config.output_hidden_states or output_hidden_states:
- all_hidden_states = all_hidden_states + (hidden_states,)
-
- if not return_dict:
- return tuple(
- v
- for v in [
- hidden_states,
- present_key_value_states,
- all_hidden_states,
- all_attentions,
- all_router_probs,
- ]
- if v is not None
- )
-
- return MoEModelOutputWithPastAndCrossAttentions(
- last_hidden_state=hidden_states,
- past_key_values=present_key_value_states,
- hidden_states=all_hidden_states,
- attentions=all_attentions,
- router_probs=all_router_probs,
- )
-
-
-@add_start_docstrings(
- "The bare GPTSAN-japanese Model with a language modeling head.",
- GPTSAN_JAPANESE_START_DOCSTRING,
-)
-class GPTSanJapaneseForConditionalGeneration(GPTSanJapanesePreTrainedModel):
- _tied_weights_keys = ["lm_head.weight"]
-
- def __init__(self, config: GPTSanJapaneseConfig):
- super().__init__(config)
- self.model = GPTSanJapaneseModel(config)
- self.register_buffer("final_logits_bias", torch.zeros([1, config.vocab_size]))
- self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
- if not self.config.torchscript:
- self.lm_head.weight = self.model.embed_tokens.weight
-
- @add_start_docstrings_to_model_forward(GPTSAN_JAPANESE_INPUTS_DOCSTRING)
- def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.FloatTensor] = None,
- token_type_ids: Optional[torch.FloatTensor] = None,
- spout: Optional[torch.FloatTensor] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- head_mask: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = False,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- output_router_logits: Optional[bool] = None,
- labels: Optional[torch.LongTensor] = None,
- ) -> Union[Tuple[torch.FloatTensor], MoECausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
- Labels for computing the sequence classification loss. Indices should be in `[-100, 0, ...,
- config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for
- labels in `[0, ..., config.vocab_size]`
-
- Returns:
- `MoECausalLMOutputWithPast` or `tuple` if `return_dict` returns MoECausalLMOutputWithPast insted of tuple
-
- Example:
-
- Text Generation with regular LM Model
- ```python
- >>> from transformers import AutoModel, AutoTokenizer, trainer_utils
-
- >>> device = "cuda"
- >>> model = AutoModel.from_pretrained("Tanrei/GPTSAN-japanese").to(device)
- >>> tokenizer = AutoTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
- >>> x_token = tokenizer("織田信長は、", return_tensors="pt")
- >>> trainer_utils.set_seed(30)
- >>> input_ids = x_token.input_ids.to(device)
- >>> gen_token = model.generate(input_ids, max_new_tokens=50)
- >>> tokenizer.decode(gen_token[0])
- "織田信長は、政治・軍事の中枢まで掌握した政治家であり、日本史上類を見ない驚異的な軍事侵攻を続け..."
- ```
-
- Text Generation with Prefix-LM Model
- ```python
- >>> from transformers import AutoModel, AutoTokenizer, trainer_utils
-
- >>> device = "cuda"
- >>> model = AutoModel.from_pretrained("Tanrei/GPTSAN-japanese").to(device)
- >>> tokenizer = AutoTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
- >>> x_token = tokenizer("", prefix_text="織田信長は、", return_tensors="pt")
- >>> trainer_utils.set_seed(30)
- >>> input_ids = x_token.input_ids.to(device)
- >>> token_type_ids = x_token.token_type_ids.to(device)
- >>> gen_token = model.generate(input_ids, token_type_ids=token_type_ids, max_new_tokens=50)
- >>> tokenizer.decode(gen_token[0])
- "織田信長は、政治・外交で数々の戦果を上げるが、1568年からは、いわゆる本能寺の変で細川晴元に暗殺される..."
- ```
-
- Simultaneously Text Generation And Masked Language Model
- ```python
- >>> from transformers import AutoModel, AutoTokenizer, trainer_utils
-
- >>> device = "cuda"
- >>> model = AutoModel.from_pretrained("Tanrei/GPTSAN-japanese").to(device)
- >>> tokenizer = AutoTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
- >>> masked_sentence = "武田信玄は、<|inputmask|>時代ファンならぜひ押さえ<|inputmask|>きたい名将の一人。"
- >>> x_token = tokenizer("", prefix_text=masked_sentence, return_tensors="pt")
- >>> trainer_utils.set_seed(30)
- >>> input_ids = x_token.input_ids.to(device)
- >>> token_type_ids = x_token.token_type_ids.to(device)
- >>> out_lm_token = model.generate(input_ids, token_type_ids=token_type_ids, max_new_tokens=50)
- >>> out_mlm_token = model(input_ids, token_type_ids=token_type_ids).logits.argmax(axis=-1)
- >>> tokenizer.decode(out_mlm_token[0])
- "武田信玄は、戦国時代ファンならぜひ押さえておきたい名将の一人。"
-
- >>> tokenizer.decode(out_lm_token[0][input_ids.shape[1] :])
- "武田氏の三代に渡った武田家のひとり\n甲斐市に住む、日本史上最大の戦国大名。..."
- ```"""
- SEG_TOKEN = self.config.separator_token_id
- use_cache = use_cache or self.config.use_cache
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
- model_return_dict = True
- num_precontext = None
- if input_ids is not None:
- num_batch = input_ids.shape[0]
- num_precontext = torch.zeros([num_batch]).int().to(input_ids.device)
- where_separators = torch.where(input_ids == SEG_TOKEN)
- num_precontext[where_separators[0]] += where_separators[1]
- num_precontext = num_precontext.unsqueeze(1)
-
- outputs = self.model(
- input_ids,
- attention_mask,
- token_type_ids,
- spout,
- past_key_values,
- head_mask,
- use_cache,
- inputs_embeds,
- decoder_inputs_embeds,
- output_attentions,
- output_hidden_states,
- model_return_dict,
- output_router_logits,
- num_precontext,
- )
-
- lm_logits = self.lm_head(outputs[0])
- if lm_logits.shape[-1] == self.final_logits_bias.shape[-1]:
- lm_logits = lm_logits + self.final_logits_bias
-
- loss = None
- z_loss = None
- router_probs = None
- aux_loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
-
- loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
-
- if output_router_logits:
- # Compute the router loss (z_loss + auxiliary loss) for each router in the encoder and decoder
- router_logits, expert_indexes = self._unpack_router_logits(outputs.router_probs)
- z_loss = router_z_loss_func(router_logits)
- router_probs = nn.Softmax(dim=-1)(router_logits)
- aux_loss = load_balancing_loss_func(router_probs, expert_indexes)
-
- loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
-
- if not return_dict:
- return tuple(
- v
- for v in [
- loss,
- lm_logits,
- outputs.past_key_values,
- outputs.hidden_states,
- outputs.router_probs,
- z_loss,
- aux_loss,
- ]
- if v is not None
- )
-
- return MoECausalLMOutputWithPast(
- loss=loss,
- logits=lm_logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- router_logits=outputs.router_probs,
- z_loss=z_loss,
- aux_loss=aux_loss,
- )
-
- def prepare_inputs_for_generation(
- self,
- input_ids: torch.LongTensor,
- attention_mask: torch.FloatTensor,
- token_type_ids: Optional[torch.FloatTensor] = None,
- spout: Optional[Union[List, torch.FloatTensor]] = None,
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
- **kwargs,
- ):
- if isinstance(spout, list):
- spout = torch.tensor(spout).float()
- if input_ids is not None:
- spout = spout.to(input_ids.device)
- if past_key_values is not None:
- return {
- "input_ids": input_ids[:, -1:] if input_ids is not None else None,
- "attention_mask": attention_mask,
- "token_type_ids": token_type_ids[:, -1:] if token_type_ids is not None else None,
- "spout": spout,
- "past_key_values": past_key_values,
- }
- return {
- "input_ids": input_ids,
- "attention_mask": attention_mask,
- "token_type_ids": token_type_ids,
- "spout": spout,
- "past_key_values": None,
- }
-
- # Copied from transformers.models.switch_transformers.modeling_switch_transformers.SwitchTransformersForConditionalGeneration.prepare_decoder_input_ids_from_labels with SwitchTransformers->GPTSanJapanese
- def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
- return self._shift_right(labels)
-
- # Copied from transformers.models.mbart.modeling_mbart.MBartForConditionalGeneration.resize_token_embeddings with MBart->GPTSanJapanese
- def resize_token_embeddings(self, new_num_tokens: int, pad_to_multiple_of: Optional[int] = None) -> nn.Embedding:
- new_embeddings = super().resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
- self._resize_final_logits_bias(new_embeddings.weight.shape[0])
- return new_embeddings
-
- # Copied from transformers.models.mbart.modeling_mbart.MBartForConditionalGeneration._resize_final_logits_bias with MBart->GPTSanJapanese
- def _resize_final_logits_bias(self, new_num_tokens: int) -> None:
- old_num_tokens = self.final_logits_bias.shape[-1]
- if new_num_tokens <= old_num_tokens:
- new_bias = self.final_logits_bias[:, :new_num_tokens]
- else:
- extra_bias = torch.zeros((1, new_num_tokens - old_num_tokens), device=self.final_logits_bias.device)
- new_bias = torch.cat([self.final_logits_bias, extra_bias], dim=1)
- self.register_buffer("final_logits_bias", new_bias)
-
- def get_input_embeddings(self):
- return self.model.get_input_embeddings()
-
- def set_input_embeddings(self, new_embeddings):
- self.model.set_input_embeddings(new_embeddings)
-
- # Copied from transformers.models.switch_transformers.modeling_switch_transformers.SwitchTransformersForConditionalGeneration.set_output_embeddings with SwitchTransformers->GPTSanJapanese
- def set_output_embeddings(self, new_embeddings):
- self.lm_head = new_embeddings
-
- # Copied from transformers.models.switch_transformers.modeling_switch_transformers.SwitchTransformersForConditionalGeneration.get_output_embeddings with SwitchTransformers->GPTSanJapanese
- def get_output_embeddings(self):
- return self.lm_head
-
- # Copied from transformers.models.switch_transformers.modeling_switch_transformers.SwitchTransformersForConditionalGeneration._unpack_router_logits with SwitchTransformers->GPTSanJapanese
- def _unpack_router_logits(self, router_outputs):
- total_router_logits = []
- total_expert_indexes = []
- for router_output in router_outputs:
- if len(router_output[0].shape) > 1:
- router_logits, expert_indexes = router_output
- total_router_logits.append(router_logits)
- total_expert_indexes.append(expert_indexes)
- return torch.cat(total_router_logits, dim=1), torch.cat(total_expert_indexes, dim=1)
diff --git a/transformers/models/gptsan_japanese/tokenization_gptsan_japanese.py b/transformers/models/gptsan_japanese/tokenization_gptsan_japanese.py
deleted file mode 100644
index f9b6d7fb5871c9c47ca7789eec220a7dbeeb8ec7..0000000000000000000000000000000000000000
--- a/transformers/models/gptsan_japanese/tokenization_gptsan_japanese.py
+++ /dev/null
@@ -1,526 +0,0 @@
-# coding=utf-8
-# Copyright 2023 HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for GPTSANJapanese."""
-import collections
-import json
-import os
-import re
-from typing import List, Optional, Tuple, Union
-
-import numpy as np
-
-from ...tokenization_utils import PreTrainedTokenizer
-from ...tokenization_utils_base import (
- BatchEncoding,
- PreTokenizedInput,
- PreTokenizedInputPair,
- TextInput,
- TextInputPair,
- TruncationStrategy,
-)
-from ...utils import PaddingStrategy, logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "emoji_file": "emoji.json"}
-
-
-def load_vocab_and_emoji(vocab_file, emoji_file):
- """Loads a vocabulary file and emoji file into a dictionary."""
- with open(emoji_file, "r", encoding="utf-8") as f:
- emoji = json.loads(f.read())
-
- vocab = collections.OrderedDict()
- raw_vocab = collections.OrderedDict()
- ids_to_tokens = collections.OrderedDict()
- with open(vocab_file, "r", encoding="utf-8") as f:
- token = f.readlines()
- token = [[t.rstrip("\n")] if (t == ",\n" or "," not in t) else t.rstrip("\n").split(",") for t in token]
- for idx, b in enumerate(token):
- ids_to_tokens[idx] = b
- raw_vocab[",".join(b)] = idx
- for wd in b:
- vocab[wd] = idx
-
- return vocab, raw_vocab, ids_to_tokens, emoji
-
-
-class GPTSanJapaneseTokenizer(PreTrainedTokenizer):
- """
- This tokenizer is based on GPTNeoXJapaneseTokenizer and has the following modifications
- - Decoding byte0~byte255 tokens correctly
- - Added bagofword token handling
- - Return token_type_ids for Prefix-LM model
- The bagofword token represents a repetition of the previous token and is converted to 3 consecutive tokens when
- decoding In addition, the original Japanese special Sub-Word-Encoding has been released in this repository
- (https://github.com/tanreinama/Japanese-BPEEncoder_V2). The token_type_ids is a mask indicating the prefix input
- position of the Prefix-LM model. To specify a prefix position, specify a prefix input for prefix_text, or specify a
- sentence of the prefix part and the part after it as a text pair of batch input.
-
- Example:
-
- ```python
- >>> from transformers import GPTSanJapaneseTokenizer
-
- >>> tokenizer = GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
- >>> # You can confirm both 慶応 and 慶應 are encoded to 17750
- >>> tokenizer("吾輩は猫である🐯。実は慶応(慶應)大学出身")["input_ids"]
- [35993, 35998, 34347, 31459, 30647, 31448, 25, 30659, 35729, 35676, 32417, 30647, 17750, 35589, 17750, 35590, 321, 1281]
-
- >>> # Both 慶応 and 慶應 are decoded to 慶応
- >>> tokenizer.decode(tokenizer("吾輩は猫である🐯。実は慶応(慶應)大学出身")["input_ids"])
- '吾輩は猫である🐯。実は慶応(慶応)大学出身'
- ```
-
- Example for Prefix-LM:
-
- ```python
- >>> from transformers import GPTSanJapaneseTokenizer
-
- >>> tokenizer = GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
- >>> tokenizer("実は慶応(慶應)大学出身", prefix_text="吾輩は猫である🐯。")["input_ids"]
- [35993, 34347, 31459, 30647, 31448, 25, 30659, 35729, 35676, 35998, 32417, 30647, 17750, 35589, 17750, 35590, 321, 1281]
-
- >>> # Mask for Prefix-LM inputs
- >>> tokenizer("実は慶応(慶應)大学出身", prefix_text="吾輩は猫である🐯。")["token_type_ids"]
- [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
- ```
-
- Example for batch encode:
-
- ```python
- >>> from transformers import GPTSanJapaneseTokenizer
-
- >>> tokenizer = GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
- >>> tokenizer([["武田信玄", "は、"], ["織田信長", "の配下の、"]], padding=True)["input_ids"]
- [[35993, 35998, 8640, 25948, 35993, 35998, 30647, 35675, 35999, 35999], [35993, 35998, 10382, 9868, 35993, 35998, 30646, 9459, 30646, 35675]]
-
- >>> # Mask for Prefix-LM inputs
- >>> tokenizer([["武田信玄", "は、"], ["織田信長", "の配下の、"]], padding=True)["token_type_ids"]
- [[1, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0, 0, 0, 0, 0]]
-
- >>> # Mask for padding
- >>> tokenizer([["武田信玄", "は、"], ["織田信長", "の配下の、"]], padding=True)["attention_mask"]
- [[1, 1, 1, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
- ```
-
- Args:
- vocab_file (`str`):
- File containing the vocabulary.
- emoji_file (`str`):
- File containing the emoji.
- unk_token (`str`, *optional*, defaults to `"<|nottoken|>"`):
- The token used for unknown charactor
- pad_token (`str`, *optional*, defaults to `"<|separator|>"`):
- The token used for padding
- bos_token (`str`, *optional*, defaults to `"<|startoftext|>"`):
- The beginning of sequence token.
- eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
- The end of sequence token.
- sep_token (`str`, *optional*, defaults to `"<|segmenter|>"`):
- A special token to separate token to prefix part and general input part.
- do_clean_text (`bool`, *optional*, defaults to `False`):
- Whether or not to clean text for URL, EMAIL, TEL, Japanese DATE and Japanese PRICE.
- """
-
- vocab_files_names = VOCAB_FILES_NAMES
- model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
-
- def __init__(
- self,
- vocab_file,
- emoji_file,
- unk_token="<|nottoken|>",
- pad_token="<|separator|>",
- bos_token="<|startoftext|>",
- eos_token="<|endoftext|>",
- sep_token="<|segmenter|>",
- do_clean_text=False,
- **kwargs,
- ):
- if not os.path.isfile(vocab_file):
- raise ValueError(
- f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
- " model use `tokenizer = GPTSanJapaneseTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- if not os.path.isfile(emoji_file):
- raise ValueError(
- f"Can't find a emoji file at path '{emoji_file}'. To load the emoji information from a Google"
- " pretrained model use `tokenizer = GPTSanJapaneseTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
- )
- self.do_clean_text = do_clean_text
- self.vocab, self.raw_vocab, self.ids_to_tokens, self.emoji = load_vocab_and_emoji(vocab_file, emoji_file)
- self.subword_tokenizer = SubWordJapaneseTokenizer(
- vocab=self.vocab, ids_to_tokens=self.ids_to_tokens, emoji=self.emoji
- )
-
- super().__init__(
- unk_token=unk_token,
- pad_token=pad_token,
- bos_token=bos_token,
- eos_token=eos_token,
- sep_token=sep_token,
- do_clean_text=do_clean_text,
- **kwargs,
- )
-
- @property
- # Copied from tokenization_gpt_neox_japanese.GPTNeoXJapaneseTokenizer.vocab_size
- def vocab_size(self):
- # self.vocab contains support for character fluctuation unique to Japanese, and has a large number of vocab
- return len(self.raw_vocab)
-
- # Copied from tokenization_gpt_neox_japanese.GPTNeoXJapaneseTokenizer.get_vocab
- def get_vocab(self):
- return dict(self.raw_vocab, **self.added_tokens_encoder)
-
- # Copied from tokenization_gpt_neox_japanese.GPTNeoXJapaneseTokenizer._tokenize
- def _tokenize(self, text):
- return self.subword_tokenizer.tokenize(text, clean=self.do_clean_text)
-
- # Copied from tokenization_gpt_neox_japanese.GPTNeoXJapaneseTokenizer._convert_token_to_id
- def _convert_token_to_id(self, token):
- """Converts a token (str) in an id using the vocab."""
- return self.vocab.get(token, self.vocab.get(self.unk_token))
-
- # Copied from tokenization_gpt_neox_japanese.GPTNeoXJapaneseTokenizer._convert_id_to_token
- def _convert_id_to_token(self, index):
- """Converts an index (integer) in a token (str) using the vocab."""
- return self.subword_tokenizer.convert_id_to_token(index)
-
- def convert_tokens_to_string(self, tokens):
- """Converts a sequence of tokens (string) in a single string."""
- words = []
- byte_tokens = []
- for word in tokens:
- if word[:6] == "<|byte" and word[-2:] == "|>":
- byte_tokens.append(int(word[6:-2]))
- else:
- if len(byte_tokens) > 0:
- words.append(bytearray(byte_tokens).decode("utf-8", errors="replace"))
- byte_tokens = []
- if word[:7] == "<|emoji" and word[-2:] == "|>":
- words.append(self.emoji["emoji_inv"][word])
- elif word == "":
- words.append(" ")
- elif word == "
":
- words.append("\n")
- elif word == "":
- words.append("\t")
- elif word == "":
- words.append("▀")
- elif word == "":
- words.append("ǀ")
- elif word == "":
- words.append("‖")
- elif word == "<|bagoftoken|>":
- if len(words) > 0:
- words.append(words[-1])
- words.append(words[-1])
- words.append(words[-1])
- elif word.startswith("<|") and word.endswith("|>"):
- words.append("")
- else:
- words.append(word)
- if len(byte_tokens) > 0:
- words.append(bytearray(byte_tokens).decode("utf-8", errors="replace"))
- text = "".join(words)
- return text
-
- @property
- def default_chat_template(self):
- """
- A simple chat template that adds standard BOS, SEP and EOS tokens between messages while discarding role
- information.
- """
- logger.warning_once(
- "\nNo chat template is defined for this tokenizer - using the default template "
- f"for the {self.__class__.__name__} class. If the default is not appropriate for "
- "your model, please set `tokenizer.chat_template` to an appropriate template. "
- "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
- )
- return (
- "{% for message in messages %}"
- "{% if not loop.first %}{{ bos_token}}{% endif %}"
- "{{ sep_token }}{{ message.content }} {{ eos_token }}"
- "{% endfor %}"
- )
-
- # Copied from tokenization_gpt_neox_japanese.GPTNeoXJapaneseTokenizer.save_vocabulary
- def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
- index = 0
- if os.path.isdir(save_directory):
- vocab_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
- )
- emoji_file = os.path.join(
- save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["emoji_file"]
- )
- else:
- vocab_file = (
- (filename_prefix + "-" if filename_prefix else "") + save_directory + VOCAB_FILES_NAMES["vocab_file"]
- )
- emoji_file = (
- (filename_prefix + "-" if filename_prefix else "") + save_directory + VOCAB_FILES_NAMES["emoji_file"]
- )
- with open(vocab_file, "w", encoding="utf-8") as writer:
- for token_index, token in self.ids_to_tokens.items():
- if index != token_index:
- logger.warning(
- f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
- " Please check that the vocabulary is not corrupted!"
- )
- index = token_index
- writer.write(",".join(token) + "\n")
- index += 1
- with open(emoji_file, "w", encoding="utf-8") as writer:
- json.dump(self.emoji, writer)
- return vocab_file, emoji_file
-
- def create_token_type_ids_from_sequences(
- self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
- ) -> List[int]:
- # docstyle-ignore
- """
- The tokenizer returns token_type_ids as separators between the Prefix part and the rest.
- token_type_ids is 1 for the Prefix part and 0 for the rest of the token.
-
- Example:
- ```python
- >>> from transformers import GPTSanJapaneseTokenizer
-
- >>> tokenizer = GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
- >>> x_token = tokenizer("ｱｲｳｴ")
- >>> # input_ids: | SOT | SEG | ｱ | ｲ | ｳ | ｴ |
- >>> # token_type_ids: | 1 | 0 | 0 | 0 | 0 | 0 |
-
- >>> x_token = tokenizer("", prefix_text="ｱｲｳｴ")
- >>> # input_ids: | SOT | ｱ | ｲ | ｳ | ｴ | SEG |
- >>> # token_type_ids: | 1 | 1 | 1 | 1 | 1 | 0 |
-
- >>> x_token = tokenizer("ｳｴ", prefix_text="ｱｲ")
- >>> # input_ids: | SOT | ｱ | ｲ | SEG | ｳ | ｴ |
- >>> # token_type_ids: | 1 | 1 | 1 | 0 | 0 | 0 |
- ```"""
- prefix_len = 0
- if self.sep_token in self.vocab:
- segid = self.vocab[self.sep_token]
- if segid in token_ids_0:
- prefix_len = token_ids_0.index(segid)
- if token_ids_1 is None:
- total_len = len(token_ids_0)
- else:
- total_len = len(token_ids_0 + token_ids_1)
- return prefix_len * [1] + (total_len - prefix_len) * [0]
-
- def prepare_for_tokenization(self, text, prefix_text=None, add_sep_token=None, **kwargs):
- # GPTSAN inserts extra SEP tokens in Prefix-LM in addition to SOT for text generation.
- # SOT at the beginning of the text, and SEP at the separator between the Prefix part and the rest.
- if add_sep_token is None:
- add_sep_token = self.sep_token not in text # If insert un-prefix position explicitly
- prepared = self.bos_token if self.bos_token in self.vocab else ""
- prepared += prefix_text if prefix_text is not None else ""
- if add_sep_token:
- prepared += self.sep_token if self.sep_token in self.vocab else ""
- prepared += text
- return (prepared, kwargs)
-
- def _batch_encode_plus(
- self,
- batch_text_or_text_pairs: Union[
- List[TextInput], List[TextInputPair], List[PreTokenizedInput], List[PreTokenizedInputPair]
- ],
- add_special_tokens: bool = True,
- padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
- truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
- max_length: Optional[int] = None,
- stride: int = 0,
- is_split_into_words: bool = False,
- pad_to_multiple_of: Optional[int] = None,
- return_tensors: Optional[str] = None,
- return_token_type_ids: Optional[bool] = None,
- return_attention_mask: Optional[bool] = None,
- return_overflowing_tokens: bool = False,
- return_special_tokens_mask: bool = False,
- return_offsets_mapping: bool = False,
- return_length: bool = False,
- verbose: bool = True,
- ) -> BatchEncoding:
- # This tokenizer converts input text pairs into Prefix input and subsequent input
- if isinstance(batch_text_or_text_pairs[0], tuple) or isinstance(tuple(batch_text_or_text_pairs[0]), list):
- # As a single text with an explicit un-prefix position
- batch_prefix_texts = []
- for pref, txt in batch_text_or_text_pairs:
- batch_prefix_texts.append(pref + self.sep_token + txt)
- batch_text_or_text_pairs = batch_prefix_texts
-
- return super()._batch_encode_plus(
- batch_text_or_text_pairs,
- add_special_tokens,
- padding_strategy,
- truncation_strategy,
- max_length,
- stride,
- is_split_into_words,
- pad_to_multiple_of,
- return_tensors,
- return_token_type_ids,
- return_attention_mask,
- return_overflowing_tokens,
- return_special_tokens_mask,
- return_offsets_mapping,
- return_length,
- verbose,
- )
-
-
-class SubWordJapaneseTokenizer(object):
- """
- This tokenizer is based on GPTNeoXJapaneseTokenizer and has the following modifications
- - Decoding byte0~byte255 tokens correctly
- - Added bagofword token handling
-
- https://github.com/tanreinama/Japanese-BPEEncoder_V2 This tokenizer class is under MIT Lisence according to the
- original repository.
-
- MIT License
-
- Copyright (c) 2020 tanreinama
-
- Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
- documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
- rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
- permit persons to whom the Software is furnished to do so, subject to the following conditions:
-
- The above copyright notice and this permission notice shall be included in all copies or substantial portions of
- the Software.
-
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
- THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
- TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- SOFTWARE.
- """
-
- # Copied from tokenization_gpt_neox_japanese.SubWordJapaneseTokenizer.__init__
- def __init__(self, vocab, ids_to_tokens, emoji):
- self.vocab = vocab # same as swe
- self.ids_to_tokens = ids_to_tokens # same as bpe
- self.emoji = emoji
- self.maxlen = np.max([len(w) for w in self.vocab.keys()])
- self.content_repatter1 = re.compile(r"(https?|ftp)(:\/\/[-_\.!~*\'()a-zA-Z0-9;\/?:\@&=\+$,%#]+)")
- self.content_repatter2 = re.compile(r"[A-Za-z0-9\._+]*@[\-_0-9A-Za-z]+(\.[A-Za-z]+)*")
- self.content_repatter3 = re.compile(r"[\(]{0,1}[0-9]{2,4}[\)\-\(]{0,1}[0-9]{2,4}[\)\-]{0,1}[0-9]{3,4}")
- self.content_repatter4 = re.compile(
- r"([12]\d{3}[/\-年])*(0?[1-9]|1[0-2])[/\-月]((0?[1-9]|[12][0-9]|3[01])日?)*(\d{1,2}|:|\d{1,2}時|\d{1,2}分|\(日\)|\(月\)|\(火\)|\(水\)|\(木\)|\(金\)|\(土\)|㈰|㈪|㈫|㈬|㈭|㈮|㈯)*"
- )
- self.content_repatter5 = re.compile(
- r"(明治|大正|昭和|平成|令和|㍾|㍽|㍼|㍻|\u32ff)\d{1,2}年(0?[1-9]|1[0-2])月(0?[1-9]|[12][0-9]|3[01])日(\d{1,2}|:|\d{1,2}時|\d{1,2}分|\(日\)|\(月\)|\(火\)|\(水\)|\(木\)|\(金\)|\(土\)|㈰|㈪|㈫|㈬|㈭|㈮|㈯)*"
- )
- self.content_repatter6 = re.compile(
- r"((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*億)*((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*万)*((0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*千)*(0|[1-9]\d*|[1-9]\d{0,2}(,\d{3})+)*(千円|万円|千万円|円|千ドル|万ドル|千万ドル|ドル|千ユーロ|万ユーロ|千万ユーロ|ユーロ)+(\(税込\)|\(税抜\)|\+tax)*"
- )
- keisen = "─━│┃┄┅┆┇┈┉┊┋┌┍┎┏┐┑┒┓└┕┖┗┘┙┚┛├┝┞┟┠┡┢┣┤┥┦┧┨┩┪┫┬┭┮┯┰┱┲┳┴┵┶┷┸┹┺┻┼┽┾┿╀╁╂╃╄╅╆╇╈╉╊╋╌╍╎╏═║╒╓╔╕╖╗╘╙╚╛╜╝╞╟╠╡╢╣╤╥╦╧╨╩╪╫╬╭╮╯╰╱╲╳╴╵╶╷╸╹╺╻╼╽╾╿"
- blocks = "▀▁▂▃▄▅▆▇█▉▊▋▌▍▎▏▐░▒▓▔▕▖▗▘▙▚▛▜▝▞▟"
- self.content_trans1 = str.maketrans({k: "" for k in keisen + blocks})
-
- # Copied from tokenization_gpt_neox_japanese.SubWordJapaneseTokenizer.__len__
- def __len__(self):
- return len(self.ids_to_tokens)
-
- # Copied from tokenization_gpt_neox_japanese.SubWordJapaneseTokenizer.clean_text
- def clean_text(self, content):
- content = self.content_repatter1.sub("", content)
- content = self.content_repatter2.sub("", content)
- content = self.content_repatter3.sub("", content)
- content = self.content_repatter4.sub("", content)
- content = self.content_repatter5.sub("", content)
- content = self.content_repatter6.sub("", content)
- content = content.translate(self.content_trans1)
- while "" in content:
- content = content.replace("", "")
- return content
-
- # Copied from tokenization_gpt_neox_japanese.SubWordJapaneseTokenizer.tokenize
- def tokenize(self, text, clean=False):
- text = text.replace(" ", "")
- text = text.replace(" ", "")
- text = text.replace("\r\n", "
")
- text = text.replace("\n", "
")
- text = text.replace("\r", "
")
- text = text.replace("\t", "")
- text = text.replace("—", "ー")
- text = text.replace("−", "ー")
- for k, v in self.emoji["emoji"].items():
- if k in text:
- text = text.replace(k, v)
- if clean:
- text = self.clean_text(text)
-
- def check_simbol(x):
- e = x.encode()
- if len(x) == 1 and len(e) == 2:
- c = (int(e[0]) << 8) + int(e[1])
- if (
- (c >= 0xC2A1 and c <= 0xC2BF)
- or (c >= 0xC780 and c <= 0xC783)
- or (c >= 0xCAB9 and c <= 0xCBBF)
- or (c >= 0xCC80 and c <= 0xCDA2)
- ):
- return True
- return False
-
- def checku2e(x):
- e = x.encode()
- if len(x) == 1 and len(e) == 3:
- c = (int(e[0]) << 16) + (int(e[1]) << 8) + int(e[2])
- if c >= 0xE28080 and c <= 0xE2B07F:
- return True
- return False
-
- pos = 0
- result = []
- while pos < len(text):
- end = min(len(text), pos + self.maxlen + 1) if text[pos] == "<" else pos + 3
- candidates = [] # (token_id, token, pos)
- for e in range(end, pos, -1):
- wd = text[pos:e]
- if wd in self.vocab:
- if wd[0] == "<" and len(wd) > 2:
- candidates = [(self.vocab[wd], wd, e)]
- break
- else:
- candidates.append((self.vocab[wd], wd, e))
- if len(candidates) > 0:
- # the smallest token_id is adopted
- _, wd, e = sorted(candidates, key=lambda x: x[0])[0]
- result.append(wd)
- pos = e
- else:
- end = pos + 1
- wd = text[pos:end]
- if check_simbol(wd):
- result.append("